git.delta.rocks / jrsonnet / refs/commits / 1aab6a2e63b6

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};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		// fn dedent(s: String) -> String {74		// 	s.split('\n').filter(|s| !s.trim().is_empty()).map(|v| v.)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					// We should receive errors earlier than synchronization127					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					// Reader doesn't care about stdout, as this is cancelled.147					// Error still might be useful, to process leftover span closures?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		// Standard repl hello doesn't work with internal-json logger187		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					// Handle startup errors, but skip repl hello?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		self.stdin.write_all(cmd.as_ref()).await?;251		self.stdin.write_all(b"\n").await?;252		Ok(())253	}254	async fn read_until_delimiter(&mut self, err_handler: &mut dyn Handler) -> Result<String> {255		let mut out = String::new();256		while let Some(line) = self.out.next().await {257			let line = match line {258				OutputLine::Out(out) => out,259				OutputLine::Err(err) => {260					err_handler.handle_line(&err);261					continue;262				}263			};264			if line == self.full_delimiter {265				return Ok(out);266			}267			if !out.is_empty() {268				out.push('\n');269			}270			out.push_str(&line);271		}272		bail!("didn't reached delimiter");273	}274	async fn execute_expression_number(&mut self, expr: impl AsRef<[u8]>) -> Result<u64> {275		let num = self.number_wrapping.clone();276		let n = self.execute_expression_wrapping(expr, &num).await?;277		Ok(n.parse::<u64>()?)278	}279	async fn execute_expression_string(&mut self, expr: impl AsRef<[u8]>) -> Result<String> {280		let num = self.string_wrapping.clone();281		let n = self.execute_expression_wrapping(expr, &num).await?;282		let str: String = serde_json::from_str(&n)?;283		Ok(str)284	}285	async fn execute_expression_to_json<V: DeserializeOwned>(286		&mut self,287		expr: impl AsRef<[u8]>,288	) -> Result<V> {289		let mut fexpr = b"builtins.toJSON (".to_vec();290		fexpr.extend_from_slice(expr.as_ref());291		fexpr.push(b')');292		let v = self.execute_expression_string(fexpr).await?;293		Ok(serde_json::from_str(&v)?)294	}295	async fn execute_expression_wrapping(296		&mut self,297		expr: impl AsRef<[u8]>,298		wrapping: &(String, String),299	) -> Result<String> {300		let mut nix_handler = self.nix_handler.clone();301		let mut collected = ErrorCollector::new(&mut nix_handler);302		let res = self.execute_expression_raw(expr, &mut collected).await?;303		if res.is_empty() {304			collected.finish()?;305			bail!("expected expression, got nothing")306		} else {307			collected.flush()308		};309		let Some(res) = res.strip_prefix(&wrapping.0) else {310			bail!("invalid type")311		};312		let Some(res) = res.strip_suffix(&wrapping.1) else {313			bail!("invalid type")314		};315		Ok(res.to_owned())316	}317	async fn execute_expression_empty(&mut self, expr: impl AsRef<[u8]>) -> Result<()> {318		let mut nix_handler = self.nix_handler.clone();319		let mut collected = ErrorCollector::new(&mut nix_handler);320		let v = self.execute_expression_raw(expr, &mut collected).await?;321		collected.finish()?;322		ensure!(v.is_empty(), "unexpected expression result");323		Ok(())324	}325	async fn execute_expression_raw(326		&mut self,327		expr: impl AsRef<[u8]>,328		err_handler: &mut dyn Handler,329	) -> Result<String> {330		self.send_command(expr).await?;331		// It will be echoed332		self.send_command(REPL_DELIMITER).await?;333		self.read_until_delimiter(err_handler).await334	}335	async fn execute_assign(&mut self, expr: impl AsRef<str>) -> Result<u32> {336		let id = self.allocate_id();337		self.execute_expression_empty(format!("sess_field_{id} = {}", expr.as_ref()))338			.await?;339		Ok(id)340	}341342	/// Id should be immediately used343	fn allocate_id(&mut self) -> u32 {344		if let Some(free) = self.free_list.pop() {345			free346		} else {347			let v = self.next_id;348			self.next_id += 1;349			v350		}351	}352	// Nix has no way to deallocate variable, yet GC will correct everything not reachable.353	// async fn free_id(&mut self, id: u32) -> Result<()> {354	// 	self.execute_expression_empty(format!("sess_field_{id} = null"))355	// 		.await?;356	// 	self.free_list.push(id);357	// 	Ok(())358	// }359}360361#[derive(Clone)]362pub struct NixSession(Arc<tokio::sync::Mutex<PooledConnection<NixSessionPoolInner>>>);363364#[macro_export]365macro_rules! nix_path {366	(@o($o:ident) $var:ident $($tt:tt)*) => {{367		$o.push(Index::var(stringify!($var)));368		nix_path!(@o($o) $($tt)*);369	}};370	(@o($o:ident) . $var:ident $($tt:tt)*) => {{371		$o.push(Index::attr(stringify!($var)));372		nix_path!(@o($o) $($tt)*);373	}};374	(@o($o:ident) . $var:literal $($tt:tt)*) => {{375		$o.push(Index::attr($var));376		nix_path!(@o($o) $($tt)*);377	}};378	(@o($o:ident) . { $var:expr } $($tt:tt)*) => {{379		$o.push(Index::attr($var));380		nix_path!(@o($o) $($tt)*);381	}};382	(@o($o:ident) [ $var:literal ] $($tt:tt)*) => {{383		$o.push(Index::idx($var));384		nix_path!(@o($o) $($tt)*);385	}};386	(@o($o:ident) ($e:expr) $($tt:tt)*) => {387		$o.push(Index::apply($e));388		nix_path!(@o($o) $($tt)*);389	};390	(@o($o:ident)) => {};391	($($tt:tt)+) => {{392		use $crate::{nix_path, better_nix_eval::Index};393		let mut out = vec![];394		nix_path!(@o(out) $($tt)*);395		out396	}}397}398399#[derive(Clone)]400pub enum Index {401	Var(String),402	String(String),403	Apply(String),404	Idx(u32),405}406impl Index {407	pub fn var(v: impl AsRef<str>) -> Self {408		let v = v.as_ref();409		assert!(410			!(v.contains('.') | v.contains(' ')),411			"bad variable name: {v}"412		);413		Self::Var(v.to_owned())414	}415	pub fn attr(v: impl AsRef<str>) -> Self {416		Self::String(v.as_ref().to_owned())417	}418	pub fn idx(v: u32) -> Self {419		Self::Idx(v)420	}421	pub fn apply(v: impl Serialize) -> Self {422		let serialized = nixlike::serialize(v).expect("invalid value for apply");423		Self::Apply(serialized)424	}425}426impl Display for Index {427	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {428		match self {429			Index::Var(v) => {430				write!(f, "{v}")431			}432			Index::String(k) => {433				let v = nixlike::format_identifier(k.as_str());434				write!(f, ".{v}")435			}436			Index::Apply(o) => {437				let v = nixlike::serialize(o).map_err(|_| fmt::Error)?;438				write!(f, "<apply>({v})")439			}440			Index::Idx(i) => {441				write!(f, "[{i}]")442			}443		}444	}445}446impl fmt::Debug for Index {447	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {448		write!(f, "{self}")449	}450}451struct PathDisplay<'i>(&'i [Index]);452impl Display for PathDisplay<'_> {453	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {454		write!(f, "flake")?;455		for i in self.0 {456			write!(f, "{i}")?;457		}458		Ok(())459	}460}461pub struct Field {462	full_path: Vec<Index>,463	session: NixSession,464	value: Option<u32>,465}466impl Field {467	fn root(session: NixSession) -> Self {468		Self {469			full_path: vec![],470			session,471			value: None,472		}473	}474	pub async fn field(session: NixSession, field: &str) -> Result<Self> {475		Self::root(session)476			.select([Index::var(field)])477			.await478	}479	pub async fn get_json_deep<'a, V: DeserializeOwned>(480		&self,481		name: impl IntoIterator<Item = Index>,482	) -> Result<V> {483		let field = self.select(name).await?;484		field.as_json().await485	}486	pub async fn select<'a>(&self, name: impl IntoIterator<Item = Index>) -> Result<Self> {487		let mut name = name.into_iter();488489		let mut full_path = self.full_path.clone();490		let mut query = if let Some(id) = self.value {491			format!("sess_field_{id}")492		} else {493			let first = name.next();494			if let Some(Index::Var(i)) = first {495				full_path.push(Index::Var(i.clone()));496				i.clone()497			} else {498				panic!("first path item should be variable, got {first:?}")499			}500		};501		for v in name {502			full_path.push(v.clone());503			match v {504				Index::Var(_) => panic!("var item may only be first"),505				Index::String(s) => {506					let escaped = nixlike::serialize(s)?;507					query.push('.');508					query.push_str(escaped.trim());509				}510				Index::Apply(a) => {511					query.push(' ');512					query.push_str(&a);513				}514				Index::Idx(idx) => {515					query = format!("builtins.elemAt ({query}) {idx}");516				}517			}518		}519520		let vid = self521			.session522			.0523			.lock()524			.await525			.execute_assign(&query)526			.await527			.with_context(|| format!("full path: {}", PathDisplay(&full_path)))?;528		Ok(Self {529			full_path,530			session: self.session.clone(),531			value: Some(vid),532		})533	}534	pub async fn as_json<V: DeserializeOwned>(&self) -> Result<V> {535		let id = self.value.expect("can't serialize root field");536		self.session537			.0538			.lock()539			.await540			.execute_expression_to_json(&format!("sess_field_{id}"))541			.await542			.with_context(|| format!("full path: {}", PathDisplay(&self.full_path)))543	}544	pub async fn list_fields(&self) -> Result<Vec<String>> {545		let id = self.value.expect("can't list root fields");546		self.session547			.0548			.lock()549			.await550			.execute_expression_to_json(&format!("builtins.attrNames sess_field_{id}"))551			.await552			.with_context(|| format!("full path: {}", PathDisplay(&self.full_path)))553	}554	pub async fn build(&self) -> Result<HashMap<String, PathBuf>> {555		let id = self.value.expect("can't use build on not-value");556		let vid = self557			.session558			.0559			.lock()560			.await561			.execute_expression_raw(&format!(":b sess_field_{id}"), &mut NixHandler::default())562			.await?;563		ensure!(!vid.is_empty(), "build failed");564		let Some(vid) = vid.strip_prefix("This derivation produced the following outputs:\n")565		else {566			panic!("unexpected build output: {vid:?}");567		};568		let outputs = vid569			.split('\n')570			.filter(|v| !v.is_empty())571			.map(|v| v.split_once(" -> ").expect("unexpected build output"))572			.map(|(a, b)| (a.trim_start().to_owned(), PathBuf::from(b)))573			.collect();574		Ok(outputs)575	}576}577impl Drop for Field {578	fn drop(&mut self) {579		if let Some(id) = self.value {580			if let Ok(mut lock) = self.session.0.try_lock() {581				lock.free_list.push(id)582			}583			// Leaked584		}585	}586}587struct NixSessionPoolInner {588	flake: OsString,589	nix_args: Vec<OsString>,590}591592#[derive(Debug)]593pub struct NixPoolError(anyhow::Error);594impl From<anyhow::Error> for NixPoolError {595	fn from(value: anyhow::Error) -> Self {596		Self(value)597	}598}599impl std::error::Error for NixPoolError {}600impl std::fmt::Display for NixPoolError {601	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {602		self.0.fmt(f)603	}604}605impl r2d2::ManageConnection for NixSessionPoolInner {606	type Connection = NixSessionInner;607	type Error = NixPoolError;608	fn connect(&self) -> std::result::Result<Self::Connection, Self::Error> {609		let _v = TOKIO_RUNTIME610			.get()611			.expect("missed tokio runtime init!")612			.enter();613		Ok(futures::executor::block_on(NixSessionInner::new(614			self.flake.as_os_str(),615			self.nix_args.iter().map(OsString::as_os_str),616		))?)617	}618619	fn is_valid(&self, conn: &mut Self::Connection) -> std::result::Result<(), Self::Error> {620		let _v = TOKIO_RUNTIME621			.get()622			.expect("missed tokio runtime init!")623			.enter();624		let res = futures::executor::block_on(conn.execute_expression_number("2 + 2"))?;625		if res != 4 {626			return Err(anyhow!("sanity check failed").into());627		};628		Ok(())629	}630631	fn has_broken(&self, _conn: &mut Self::Connection) -> bool {632		false633	}634}635pub struct NixSessionPool(Pool<NixSessionPoolInner>);636impl NixSessionPool {637	pub async fn new(flake: OsString, nix_args: Vec<OsString>) -> Result<Self> {638		let inner = tokio::task::block_in_place(|| {639			r2d2::Builder::<NixSessionPoolInner>::new()640				.min_idle(Some(0))641				.build(NixSessionPoolInner { flake, nix_args })642		})?;643		Ok(Self(inner))644	}645	pub async fn get(&self) -> Result<NixSession> {646		let v = tokio::task::block_in_place(|| self.0.get())?;647		Ok(NixSession(Arc::new(tokio::sync::Mutex::new(v))))648	}649}650651pub static TOKIO_RUNTIME: OnceLock<tokio::runtime::Handle> = OnceLock::new();