1use std::collections::HashMap;2use std::ffi::{OsStr, OsString};3use std::fmt::{self, Display};4use std::path::PathBuf;5use std::process::Stdio;6use std::sync::{Arc, OnceLock};78use anyhow::{anyhow, bail, ensure, Context, Result};9use futures::StreamExt;10use itertools::Itertools;11use r2d2::{Pool, PooledConnection};12use serde::de::DeserializeOwned;13use serde::{Deserialize, Serialize};14use tokio::io::AsyncWriteExt;15use tokio::process::{ChildStderr, ChildStdin, ChildStdout, Command};16use tokio::select;17use tokio::sync::{mpsc, oneshot};18use tokio_util::codec::{FramedRead, LinesCodec};19use tracing::{debug, error, warn, Level};2021use crate::command::{ClonableHandler, Handler, NixHandler, NoopHandler};2223const REPL_DELIMITER: &str = "\"FLEET_MAGIC_REPL_DELIMITER\"";2425pub struct NixSessionInner {26 full_delimiter: String,27 nix_handler: ClonableHandler<NixHandler>,28 out: OutputHandler,29 stdin: ChildStdin,30 string_wrapping: (String, String),31 number_wrapping: (String, String),3233 next_id: u32,34 free_list: Vec<u32>,35}36const TRAIN_STRING: &str = "\"TRAIN_STRING\"";37const TRAIN_NUMBER: &str = "13141516";3839#[must_use]40struct ErrorCollector<'i, H> {41 collected: Vec<String>,42 inner: &'i mut H,43}44impl<'i, H> ErrorCollector<'i, H> {45 fn new(inner: &'i mut H) -> Self {46 Self {47 collected: vec![],48 inner,49 }50 }51}52impl<H> ErrorCollector<'_, H> {53 fn handle_line_inner(&mut self, msg: &str) -> bool {54 let Some(msg) = msg.strip_prefix("@nix ") else {55 return false;56 };57 #[derive(Deserialize)]58 struct ErrorAction {59 action: String,60 level: u32,61 msg: String,62 }63 let Ok(act) = serde_json::from_str::<ErrorAction>(msg) else {64 return false;65 };66 if act.action != "msg" || act.level != 0 {67 return false;68 }69 self.collected.push(act.msg);70 true71 }72 fn finish(self) -> Result<()> {73 74 75 76 if !self.collected.is_empty() {77 bail!(78 "{}",79 self.collected80 .iter()81 .map(|v| {82 if let Some(f) = v.strip_prefix("\u{1b}[31;1merror:\u{1b}[0m ") {83 let v = unindent::unindent(f.trim_start());84 v.trim().to_owned()85 } else {86 v.to_owned()87 }88 })89 .join("\n")90 );91 }92 Ok(())93 }94 fn flush(self) {95 for line in self.collected {96 warn!("{line}");97 }98 }99}100impl<H: Handler> Handler for ErrorCollector<'_, H> {101 fn handle_line(&mut self, e: &str) {102 if self.handle_line_inner(e) {103 return;104 }105 self.inner.handle_line(e)106 }107}108109enum OutputLine {110 Out(String),111 Err(String),112}113struct OutputHandler {114 rx: mpsc::Receiver<OutputLine>,115 _cancel_handle: oneshot::Receiver<()>,116}117impl OutputHandler {118 fn new(out: ChildStdout, err: ChildStderr) -> Self {119 let mut out = FramedRead::new(out, LinesCodec::new());120 let mut err = FramedRead::new(err, LinesCodec::new());121 let (tx, rx) = mpsc::channel(20);122 let (mut cancelled, _cancel_handle) = oneshot::channel();123 tokio::spawn(async move {124 loop {125 select! {126 127 biased;128 e = err.next() => {129 let Some(Ok(e)) = e else {130 if e.is_some() {131 error!("bad repl stderr: {e:?}");132 }133 continue;134 };135 let _ = tx.send(OutputLine::Err(e)).await;136 }137 o = out.next() => {138 let Some(Ok(o)) = o else {139 if o.is_some() {140 error!("bad repl stdout: {o:?}");141 }142 continue;143 };144 let _ = tx.send(OutputLine::Out(o)).await;145 }146 147 148 _ = cancelled.closed() => {149 break;150 }151 }152 }153 });154 Self { rx, _cancel_handle }155 }156 async fn next(&mut self) -> Option<OutputLine> {157 self.rx.recv().await158 }159}160161struct WarnHandler;162impl Handler for WarnHandler {163 fn handle_line(&mut self, e: &str) {164 warn!(target: "nix", "{e}")165 }166}167168impl NixSessionInner {169 async fn new(flake: &OsStr, extra_args: impl IntoIterator<Item = &OsStr>) -> Result<Self> {170 let mut cmd = Command::new("nix");171 cmd.arg("repl")172 .arg(flake)173 .arg("--log-format")174 .arg("internal-json");175 for arg in extra_args {176 cmd.arg(arg);177 }178 cmd.stdin(Stdio::piped());179 cmd.stdout(Stdio::piped());180 cmd.stderr(Stdio::piped());181 let cmd = cmd.spawn()?;182 let stdout = cmd.stdout.unwrap();183 let stderr = cmd.stderr.unwrap();184 let mut out = OutputHandler::new(stdout, stderr);185 let mut stdin = cmd.stdin.unwrap();186 187 stdin.write_all(REPL_DELIMITER.as_bytes()).await?;188 stdin.write_all(b"\n").await?;189 stdin.flush().await?;190 let nix_handler = NixHandler::default();191 let mut full_delimiter = None;192 let mut errors = vec![];193 while let Some(line) = out.next().await {194 let line = match line {195 OutputLine::Out(o) => o,196 OutputLine::Err(_e) => {197 198 errors.push(_e);199 continue;200 }201 };202 if line.contains(REPL_DELIMITER) {203 debug!("discovered repl delimiter with added colors: {line}");204 full_delimiter = Some(line.to_owned());205 break;206 }207 }208 let Some(full_delimiter) = full_delimiter else {209 for e in errors {210 error!("{e}");211 }212 bail!("failed to discover delimiter");213 };214 let mut res = Self {215 full_delimiter,216 nix_handler: ClonableHandler::new(nix_handler),217 out,218 stdin,219 string_wrapping: Default::default(),220 number_wrapping: Default::default(),221222 next_id: 0,223 free_list: vec![],224 };225 res.train().await?;226 Ok(res)227 }228 async fn train(&mut self) -> Result<()> {229 {230 let full_string = self231 .execute_expression_raw(TRAIN_STRING, &mut NoopHandler)232 .await?;233 let string_offset = full_string.find(TRAIN_STRING).expect("contained");234 let string_prefix = &full_string[..string_offset];235 let string_suffix = &full_string[string_offset + TRAIN_STRING.len()..];236 self.string_wrapping = (string_prefix.to_owned(), string_suffix.to_owned());237 }238 {239 let full_number = self240 .execute_expression_raw(TRAIN_NUMBER, &mut NoopHandler)241 .await?;242 let number_offset = full_number.find(TRAIN_NUMBER).expect("contained");243 let number_prefix = &full_number[..number_offset];244 let number_suffix = &full_number[number_offset + TRAIN_NUMBER.len()..];245 self.number_wrapping = (number_prefix.to_owned(), number_suffix.to_owned());246 }247 Ok(())248 }249 async fn send_command(&mut self, cmd: impl AsRef<[u8]>) -> Result<()> {250 if tracing::enabled!(Level::DEBUG) {251 let cmd_str = String::from_utf8_lossy(cmd.as_ref());252 tracing::debug!("{cmd_str}");253 };254 self.stdin.write_all(cmd.as_ref()).await?;255 self.stdin.write_all(b"\n").await?;256 Ok(())257 }258 async fn read_until_delimiter(&mut self, err_handler: &mut dyn Handler) -> Result<String> {259 let mut out = String::new();260 while let Some(line) = self.out.next().await {261 let line = match line {262 OutputLine::Out(out) => out,263 OutputLine::Err(err) => {264 err_handler.handle_line(&err);265 continue;266 }267 };268 if line == self.full_delimiter {269 return Ok(out);270 }271 if !out.is_empty() {272 out.push('\n');273 }274 out.push_str(&line);275 }276 bail!("didn't reached delimiter");277 }278 async fn execute_expression_number(&mut self, expr: impl AsRef<[u8]>) -> Result<u64> {279 let num = self.number_wrapping.clone();280 let n = self.execute_expression_wrapping(expr, &num).await?;281 Ok(n.parse::<u64>()?)282 }283 async fn execute_expression_string(&mut self, expr: impl AsRef<[u8]>) -> Result<String> {284 let num = self.string_wrapping.clone();285 let n = self.execute_expression_wrapping(expr, &num).await?;286 let str: String = serde_json::from_str(&n)?;287 Ok(str)288 }289 async fn execute_expression_to_json<V: DeserializeOwned>(290 &mut self,291 expr: impl AsRef<[u8]>,292 ) -> Result<V> {293 let mut fexpr = b"builtins.toJSON (".to_vec();294 fexpr.extend_from_slice(expr.as_ref());295 fexpr.push(b')');296 let v = self.execute_expression_string(fexpr).await?;297 Ok(serde_json::from_str(&v)?)298 }299 async fn execute_expression_wrapping(300 &mut self,301 expr: impl AsRef<[u8]>,302 wrapping: &(String, String),303 ) -> Result<String> {304 let mut nix_handler = self.nix_handler.clone();305 let mut collected = ErrorCollector::new(&mut nix_handler);306 let res = self.execute_expression_raw(expr, &mut collected).await?;307 if res.is_empty() {308 collected.finish()?;309 bail!("expected expression, got nothing")310 } else {311 collected.flush()312 };313 let Some(res) = res.strip_prefix(&wrapping.0) else {314 bail!("invalid type")315 };316 let Some(res) = res.strip_suffix(&wrapping.1) else {317 bail!("invalid type")318 };319 Ok(res.to_owned())320 }321 async fn execute_expression_empty(&mut self, expr: impl AsRef<[u8]>) -> Result<()> {322 let mut nix_handler = self.nix_handler.clone();323 let mut collected = ErrorCollector::new(&mut nix_handler);324 let v = self.execute_expression_raw(expr, &mut collected).await?;325 collected.finish()?;326 ensure!(v.is_empty(), "unexpected expression result");327 Ok(())328 }329 async fn execute_expression_raw(330 &mut self,331 expr: impl AsRef<[u8]>,332 err_handler: &mut dyn Handler,333 ) -> Result<String> {334 self.send_command(expr).await?;335 336 self.send_command(REPL_DELIMITER).await?;337 self.read_until_delimiter(err_handler).await338 }339 async fn execute_assign(&mut self, expr: impl AsRef<str>) -> Result<u32> {340 let id = self.allocate_id();341 self.execute_expression_empty(format!("sess_field_{id} = {}", expr.as_ref()))342 .await?;343 Ok(id)344 }345346 347 fn allocate_id(&mut self) -> u32 {348 if let Some(free) = self.free_list.pop() {349 free350 } else {351 let v = self.next_id;352 self.next_id += 1;353 v354 }355 }356 357 358 359 360 361 362 363}364365#[derive(Clone)]366pub struct NixSession(Arc<tokio::sync::Mutex<PooledConnection<NixSessionPoolInner>>>);367368#[macro_export]369macro_rules! nix_path {370 (@o($o:ident) $var:ident $($tt:tt)*) => {{371 $o.push(Index::var(stringify!($var)));372 nix_path!(@o($o) $($tt)*);373 }};374 (@o($o:ident) . $var:ident $($tt:tt)*) => {{375 $o.push(Index::attr(stringify!($var)));376 nix_path!(@o($o) $($tt)*);377 }};378 (@o($o:ident) . $var:literal $($tt:tt)*) => {{379 $o.push(Index::attr($var));380 nix_path!(@o($o) $($tt)*);381 }};382 (@o($o:ident) . { $var:expr } $($tt:tt)*) => {{383 $o.push(Index::attr($var));384 nix_path!(@o($o) $($tt)*);385 }};386 (@o($o:ident) [ $var:literal ] $($tt:tt)*) => {{387 $o.push(Index::idx($var));388 nix_path!(@o($o) $($tt)*);389 }};390 (@o($o:ident) ($e:expr) $($tt:tt)*) => {391 $o.push(Index::apply($e));392 nix_path!(@o($o) $($tt)*);393 };394 (@o($o:ident)) => {};395 ($($tt:tt)+) => {{396 use $crate::{nix_path, better_nix_eval::Index};397 let mut out = vec![];398 nix_path!(@o(out) $($tt)*);399 out400 }}401}402403#[derive(Clone)]404pub enum Index {405 Var(String),406 String(String),407 Apply(String),408 Idx(u32),409}410impl Index {411 pub fn var(v: impl AsRef<str>) -> Self {412 let v = v.as_ref();413 assert!(414 !(v.contains('.') | v.contains(' ')),415 "bad variable name: {v}"416 );417 Self::Var(v.to_owned())418 }419 pub fn attr(v: impl AsRef<str>) -> Self {420 Self::String(v.as_ref().to_owned())421 }422 pub fn idx(v: u32) -> Self {423 Self::Idx(v)424 }425 pub fn apply(v: impl Serialize) -> Self {426 let serialized = nixlike::serialize(v).expect("invalid value for apply");427 Self::Apply(serialized.trim_end().to_owned())428 }429}430impl Display for Index {431 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {432 match self {433 Index::Var(v) => {434 write!(f, "{v}")435 }436 Index::String(k) => {437 let v = nixlike::format_identifier(k.as_str());438 write!(f, ".{v}")439 }440 Index::Apply(o) => {441 write!(f, "<apply>({o})")442 }443 Index::Idx(i) => {444 write!(f, "[{i}]")445 }446 }447 }448}449impl fmt::Debug for Index {450 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {451 write!(f, "{self}")452 }453}454struct PathDisplay<'i>(&'i [Index]);455impl Display for PathDisplay<'_> {456 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {457 for i in self.0 {458 write!(f, "{i}")?;459 }460 Ok(())461 }462}463pub struct Field {464 full_path: Vec<Index>,465 session: NixSession,466 value: Option<u32>,467}468impl Field {469 fn root(session: NixSession) -> Self {470 Self {471 full_path: vec![],472 session,473 value: None,474 }475 }476 pub async fn field(session: NixSession, field: &str) -> Result<Self> {477 Self::root(session)478 .select([Index::var(field)])479 .await480 }481 pub async fn get_json_deep<'a, V: DeserializeOwned>(482 &self,483 name: impl IntoIterator<Item = Index>,484 ) -> Result<V> {485 let field = self.select(name).await?;486 field.as_json().await487 }488 pub async fn select<'a>(&self, name: impl IntoIterator<Item = Index>) -> Result<Self> {489 let mut name = name.into_iter();490491 let mut full_path = self.full_path.clone();492 let mut query = if let Some(id) = self.value {493 format!("sess_field_{id}")494 } else {495 let first = name.next();496 if let Some(Index::Var(i)) = first {497 full_path.push(Index::Var(i.clone()));498 i.clone()499 } else {500 panic!("first path item should be variable, got {first:?}")501 }502 };503 for v in name {504 full_path.push(v.clone());505 match v {506 Index::Var(_) => panic!("var item may only be first"),507 Index::String(s) => {508 let escaped = nixlike::serialize(s)?;509 query.push('.');510 query.push_str(escaped.trim());511 }512 Index::Apply(a) => {513 514 query = format!("({query} {a})");515 }516 Index::Idx(idx) => {517 query = format!("builtins.elemAt ({query}) {idx}");518 }519 }520 }521522 let vid = self523 .session524 .0525 .lock()526 .await527 .execute_assign(&query)528 .await529 .with_context(|| format!("full path: {}", PathDisplay(&full_path)))?;530 Ok(Self {531 full_path,532 session: self.session.clone(),533 value: Some(vid),534 })535 }536 pub async fn as_json<V: DeserializeOwned>(&self) -> Result<V> {537 let id = self.value.expect("can't serialize root field");538 self.session539 .0540 .lock()541 .await542 .execute_expression_to_json(&format!("sess_field_{id}"))543 .await544 .with_context(|| format!("full path: {}", PathDisplay(&self.full_path)))545 }546 pub async fn list_fields(&self) -> Result<Vec<String>> {547 let id = self.value.expect("can't list root fields");548 self.session549 .0550 .lock()551 .await552 .execute_expression_to_json(&format!("builtins.attrNames sess_field_{id}"))553 .await554 .with_context(|| format!("full path: {}", PathDisplay(&self.full_path)))555 }556 pub async fn build(&self) -> Result<HashMap<String, PathBuf>> {557 let id = self.value.expect("can't use build on not-value");558 let vid = self559 .session560 .0561 .lock()562 .await563 .execute_expression_raw(&format!(":b sess_field_{id}"), &mut NixHandler::default())564 .await?;565 ensure!(!vid.is_empty(), "build failed: {}", PathDisplay(&self.full_path));566 let Some(vid) = vid.strip_prefix("This derivation produced the following outputs:\n")567 else {568 panic!("unexpected build output: {vid:?}");569 };570 let outputs = vid571 .split('\n')572 .filter(|v| !v.is_empty())573 .map(|v| v.split_once(" -> ").expect("unexpected build output"))574 .map(|(a, b)| (a.trim_start().to_owned(), PathBuf::from(b)))575 .collect();576 Ok(outputs)577 }578}579impl Drop for Field {580 fn drop(&mut self) {581 if let Some(id) = self.value {582 if let Ok(mut lock) = self.session.0.try_lock() {583 lock.free_list.push(id)584 }585 586 }587 }588}589struct NixSessionPoolInner {590 flake: OsString,591 nix_args: Vec<OsString>,592}593594#[derive(Debug)]595pub struct NixPoolError(anyhow::Error);596impl From<anyhow::Error> for NixPoolError {597 fn from(value: anyhow::Error) -> Self {598 Self(value)599 }600}601impl std::error::Error for NixPoolError {}602impl std::fmt::Display for NixPoolError {603 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {604 self.0.fmt(f)605 }606}607impl r2d2::ManageConnection for NixSessionPoolInner {608 type Connection = NixSessionInner;609 type Error = NixPoolError;610 fn connect(&self) -> std::result::Result<Self::Connection, Self::Error> {611 let _v = TOKIO_RUNTIME612 .get()613 .expect("missed tokio runtime init!")614 .enter();615 Ok(futures::executor::block_on(NixSessionInner::new(616 self.flake.as_os_str(),617 self.nix_args.iter().map(OsString::as_os_str),618 ))?)619 }620621 fn is_valid(&self, conn: &mut Self::Connection) -> std::result::Result<(), Self::Error> {622 let _v = TOKIO_RUNTIME623 .get()624 .expect("missed tokio runtime init!")625 .enter();626 let res = futures::executor::block_on(conn.execute_expression_number("2 + 2"))?;627 if res != 4 {628 return Err(anyhow!("sanity check failed").into());629 };630 Ok(())631 }632633 fn has_broken(&self, _conn: &mut Self::Connection) -> bool {634 false635 }636}637pub struct NixSessionPool(Pool<NixSessionPoolInner>);638impl NixSessionPool {639 pub async fn new(flake: OsString, nix_args: Vec<OsString>) -> Result<Self> {640 let inner = tokio::task::block_in_place(|| {641 r2d2::Builder::<NixSessionPoolInner>::new()642 .min_idle(Some(0))643 .build(NixSessionPoolInner { flake, nix_args })644 })?;645 Ok(Self(inner))646 }647 pub async fn get(&self) -> Result<NixSession> {648 let v = tokio::task::block_in_place(|| self.0.get())?;649 Ok(NixSession(Arc::new(tokio::sync::Mutex::new(v))))650 }651}652653pub static TOKIO_RUNTIME: OnceLock<tokio::runtime::Handle> = OnceLock::new();