git.delta.rocks / jrsonnet / refs/commits / 064468b85673

1use std::{ffi::OsStr, num::ParseIntError, process::Stdio, sync::Arc};23use better_command::{ClonableHandler, Handler, NixHandler, NoopHandler};4use futures::StreamExt;5use itertools::Itertools as _;6use serde::{de::DeserializeOwned, Deserialize};7use thiserror::Error;8use tokio::{9	io::AsyncWriteExt,10	process::{ChildStderr, ChildStdin, ChildStdout, Command},11	select,12	sync::{mpsc, oneshot, Mutex},13};14use tokio_util::codec::{FramedRead, LinesCodec};15use tracing::{debug, error, warn, Level};1617#[derive(Error, Debug, Clone)]18pub enum Error {19	#[error("failed to create nix repl session: {0}")]20	SessionInit(&'static str),21	#[error("unexpected end of output, nix crashed?")]22	MissingDelimiter,2324	#[error("expression did'nt produce any output")]25	ExpectedOutput,26	#[error("expression produced output, which is unexpected")]27	UnexpectedOutput,2829	#[error("unexpected expression output type")]30	InvalidType,3132	#[error("failed to build attr {attribute}:\n{error}")]33	BuildFailed { attribute: String, error: String },3435	#[error("output: {0}")]36	Json(Arc<serde_json::Error>),37	// int outputs are too specific, and should not be used,38	// thus error is ok to be not informative.39	#[error("int output: {0}")]40	Int(ParseIntError),41	#[error("pool: {0}")]42	Pool(Arc<r2d2::Error>),43	#[error("io: {0}")]44	Io(Arc<std::io::Error>),4546	// TODO: Should be done by wrapper/in different type.47	#[error("at {0}: {1}")]48	InContext(String, Box<Self>),4950	#[error("error: {0}")]51	NixError(String),52}53impl From<r2d2::Error> for Error {54	fn from(value: r2d2::Error) -> Self {55		Self::Pool(Arc::new(value))56	}57}58impl From<std::io::Error> for Error {59	fn from(value: std::io::Error) -> Self {60		Self::Io(Arc::new(value))61	}62}63impl From<serde_json::Error> for Error {64	fn from(value: serde_json::Error) -> Self {65		Self::Json(Arc::new(value))66	}67}68impl Error {69	pub(crate) fn context(self, context: String) -> Self {70		Self::InContext(context, Box::new(self))71	}72}73pub type Result<T, E = Error> = std::result::Result<T, E>;7475enum OutputLine {76	Out(String),77	Err(String),78}79struct OutputHandler {80	rx: mpsc::Receiver<OutputLine>,81	_cancel_handle: oneshot::Receiver<()>,82}83impl OutputHandler {84	fn new(out: ChildStdout, err: ChildStderr) -> Self {85		let mut out = FramedRead::new(out, LinesCodec::new());86		let mut err = FramedRead::new(err, LinesCodec::new());87		let (tx, rx) = mpsc::channel(20);88		let (mut cancelled, _cancel_handle) = oneshot::channel();89		tokio::spawn(async move {90			loop {91				select! {92					// We should receive errors earlier than synchronization93					biased;94					e = err.next() => {95						let Some(Ok(e)) = e else {96							if e.is_some() {97								error!("bad repl stderr: {e:?}");98							}99							continue;100						};101						let _ = tx.send(OutputLine::Err(e)).await;102					}103					o = out.next() => {104						let Some(Ok(o)) = o else {105							if o.is_some() {106								error!("bad repl stdout: {o:?}");107							}108							continue;109						};110						let _ = tx.send(OutputLine::Out(o)).await;111					}112					// Reader doesn't care about stdout, as this is cancelled.113					// Error still might be useful, to process leftover span closures?114					_ = cancelled.closed() => {115						break;116					}117				}118			}119		});120		Self { rx, _cancel_handle }121	}122	async fn next(&mut self) -> Option<OutputLine> {123		self.rx.recv().await124	}125}126127#[must_use]128struct ErrorCollector<'i, H> {129	collected: Vec<String>,130	inner: &'i mut H,131}132impl<'i, H> ErrorCollector<'i, H> {133	fn new(inner: &'i mut H) -> Self {134		Self {135			collected: vec![],136			inner,137		}138	}139}140impl<H> ErrorCollector<'_, H> {141	fn handle_line_inner(&mut self, msg: &str) -> bool {142		let Some(msg) = msg.strip_prefix("@nix ") else {143			return false;144		};145		#[derive(Deserialize)]146		struct ErrorAction {147			action: String,148			level: u32,149			msg: String,150		}151		let Ok(act) = serde_json::from_str::<ErrorAction>(msg) else {152			return false;153		};154		if act.action != "msg" || act.level != 0 {155			return false;156		}157		self.collected.push(act.msg);158		true159	}160	fn finish(self) -> Result<()> {161		// fn dedent(s: String) -> String {162		// 	s.split('\n').filter(|s| !s.trim().is_empty()).map(|v| v.)163		// }164		if !self.collected.is_empty() {165			return Err(Error::NixError(166				self.collected167					.iter()168					.map(|v| {169						if let Some(f) = v.strip_prefix("\u{1b}[31;1merror:\u{1b}[0m ") {170							let v = unindent::unindent(f.trim_start());171							v.trim().to_owned()172						} else {173							v.to_owned()174						}175					})176					.join("\n")177					.to_string(),178			));179		}180		Ok(())181	}182	fn flush(self) {183		for line in self.collected {184			warn!("{line}");185		}186	}187}188impl<H: Handler> Handler for ErrorCollector<'_, H> {189	fn handle_line(&mut self, e: &str) {190		if self.handle_line_inner(e) {191			return;192		}193		self.inner.handle_line(e)194	}195}196197pub struct NixSessionInner {198	full_delimiter: String,199	nix_handler: ClonableHandler<NixHandler>,200	out: OutputHandler,201	stdin: ChildStdin,202	string_wrapping: (String, String),203	number_wrapping: (String, String),204205	executing_command: Arc<Mutex<()>>,206207	next_id: u32,208	pub(crate) free_list: Vec<u32>,209}210211/// Discover inter-message repl delimiter212const REPL_DELIMITER: &str = "\"FLEET_MAGIC_REPL_DELIMITER\"";213/// Discover formatting around strings214const TRAIN_STRING: &str = "\"TRAIN_STRING\"";215/// Discover formatting around numbers216const TRAIN_NUMBER: &str = "13141516";217// Other types of formatting are not discovered, because they are not used, JSON serialization is used instead218// Techically, number training is also not required, because numbers can be converted to string too...219// Eh, I'll remove it later.220221impl NixSessionInner {222	pub(crate) async fn new(223		flake: &OsStr,224		extra_args: impl IntoIterator<Item = &OsStr>,225	) -> Result<Self> {226		let mut cmd = Command::new("nix");227		cmd.arg("repl")228			.arg(flake)229			.arg("--log-format")230			.arg("internal-json");231		for arg in extra_args {232			cmd.arg(arg);233		}234		cmd.stdin(Stdio::piped());235		cmd.stdout(Stdio::piped());236		cmd.stderr(Stdio::piped());237		let cmd = cmd.spawn()?;238		let stdout = cmd.stdout.unwrap();239		let stderr = cmd.stderr.unwrap();240		let mut out = OutputHandler::new(stdout, stderr);241		let mut stdin = cmd.stdin.unwrap();242		// Standard repl hello doesn't work with internal-json logger243		stdin.write_all(REPL_DELIMITER.as_bytes()).await?;244		stdin.write_all(b"\n").await?;245		stdin.flush().await?;246		let nix_handler = NixHandler::default();247		let mut full_delimiter = None;248		let mut errors = vec![];249		while let Some(line) = out.next().await {250			let line = match line {251				OutputLine::Out(o) => o,252				OutputLine::Err(_e) => {253					// Handle startup errors, but skip repl hello?254					errors.push(_e);255					continue;256				}257			};258			if line.contains(REPL_DELIMITER) {259				debug!("discovered repl delimiter with added colors: {line}");260				full_delimiter = Some(line.to_owned());261				break;262			}263		}264		let Some(full_delimiter) = full_delimiter else {265			for e in errors {266				error!("{e}");267			}268			return Err(Error::SessionInit("failed to discover delimiter"));269		};270		let mut res = Self {271			full_delimiter,272			nix_handler: ClonableHandler::new(nix_handler),273			out,274			stdin,275			string_wrapping: Default::default(),276			number_wrapping: Default::default(),277278			executing_command: Arc::new(Mutex::new(())),279280			next_id: 0,281			free_list: vec![],282		};283		res.train().await?;284		Ok(res)285	}286	async fn train(&mut self) -> Result<()> {287		{288			let full_string = self289				.execute_expression_raw(TRAIN_STRING, &mut NoopHandler)290				.await?;291			let string_offset = full_string.find(TRAIN_STRING).expect("contained");292			let string_prefix = &full_string[..string_offset];293			let string_suffix = &full_string[string_offset + TRAIN_STRING.len()..];294			self.string_wrapping = (string_prefix.to_owned(), string_suffix.to_owned());295		}296		{297			let full_number = self298				.execute_expression_raw(TRAIN_NUMBER, &mut NoopHandler)299				.await?;300			let number_offset = full_number.find(TRAIN_NUMBER).expect("contained");301			let number_prefix = &full_number[..number_offset];302			let number_suffix = &full_number[number_offset + TRAIN_NUMBER.len()..];303			self.number_wrapping = (number_prefix.to_owned(), number_suffix.to_owned());304		}305		Ok(())306	}307	async fn send_command(&mut self, cmd: impl AsRef<[u8]>) -> Result<()> {308		if tracing::enabled!(Level::DEBUG) && cmd.as_ref() != REPL_DELIMITER.as_bytes() {309			let cmd_str = String::from_utf8_lossy(cmd.as_ref());310			tracing::debug!("{cmd_str}");311		};312		self.stdin.write_all(cmd.as_ref()).await?;313		self.stdin.write_all(b"\n").await?;314		Ok(())315	}316	async fn read_until_delimiter(&mut self, err_handler: &mut dyn Handler) -> Result<String> {317		let mut out = String::new();318		while let Some(line) = self.out.next().await {319			let line = match line {320				OutputLine::Out(out) => out,321				OutputLine::Err(err) => {322					err_handler.handle_line(&err);323					continue;324				}325			};326			if line == self.full_delimiter {327				return Ok(out);328			}329			if !out.is_empty() {330				out.push('\n');331			}332			out.push_str(&line);333		}334		Err(Error::MissingDelimiter)335	}336	pub(crate) async fn execute_expression_number(337		&mut self,338		expr: impl AsRef<[u8]>,339	) -> Result<u64> {340		let num = self.number_wrapping.clone();341		let n = self.execute_expression_wrapping(expr, &num).await?;342		n.parse::<u64>().map_err(Error::Int)343	}344	async fn execute_expression_string(&mut self, expr: impl AsRef<[u8]>) -> Result<String> {345		// builtins.toJSON escapes some thing in incorrect way, e.g escaped "$" in "\${" is being outputed as "\$",346		// while this escape should be removed as it is intended for nix itself, not for json output.347		//348		// This regex only allows \$ in the beginning of the string, it is easier to implement correctly.349		// TODO: Add peg parser for nix-produced JSON?..350		let regex = regex::Regex::new(r#"(?<prefix>[: {,\[]\\")\\\$"#).expect("fixup json");351352		let num = self.string_wrapping.clone();353		let n = self.execute_expression_wrapping(expr, &num).await?;354		let n = regex.replace_all(&n, "$prefix$$");355		let str: String = serde_json::from_str(&n)?;356		Ok(str)357	}358	pub(crate) async fn execute_expression_to_json<V: DeserializeOwned>(359		&mut self,360		expr: impl AsRef<[u8]>,361	) -> Result<V> {362		let mut fexpr = b"builtins.toJSON (".to_vec();363		fexpr.extend_from_slice(expr.as_ref());364		fexpr.push(b')');365366		Ok(serde_json::from_str(367			&self.execute_expression_string(fexpr).await?,368		)?)369	}370	async fn execute_expression_wrapping(371		&mut self,372		expr: impl AsRef<[u8]>,373		wrapping: &(String, String),374	) -> Result<String> {375		let mut nix_handler = self.nix_handler.clone();376		let mut collected = ErrorCollector::new(&mut nix_handler);377		let res = self.execute_expression_raw(expr, &mut collected).await?;378		if res.is_empty() {379			collected.finish()?;380			return Err(Error::ExpectedOutput);381		} else {382			collected.flush()383		};384		let Some(res) = res.strip_prefix(&wrapping.0) else {385			return Err(Error::InvalidType);386		};387		let Some(res) = res.strip_suffix(&wrapping.1) else {388			return Err(Error::InvalidType);389		};390		Ok(res.to_owned())391	}392	async fn execute_expression_empty(&mut self, expr: impl AsRef<[u8]>) -> Result<()> {393		let mut nix_handler = self.nix_handler.clone();394		let mut collected = ErrorCollector::new(&mut nix_handler);395		let v = self.execute_expression_raw(expr, &mut collected).await?;396		collected.finish()?;397		if !v.is_empty() {398			return Err(Error::UnexpectedOutput);399		}400		Ok(())401	}402	pub(crate) async fn execute_expression_raw(403		&mut self,404		expr: impl AsRef<[u8]>,405		err_handler: &mut dyn Handler,406	) -> Result<String> {407		// Prevent two commands from being executed in parallel, messing with each other.408		let _lock = self.executing_command.clone();409		let _guard = _lock.lock().await;410411		self.send_command(expr).await?;412		// It will be echoed413		self.send_command(REPL_DELIMITER).await?;414		self.read_until_delimiter(err_handler).await415	}416	pub(crate) async fn execute_assign(&mut self, expr: impl AsRef<str>) -> Result<u32> {417		let id = self.allocate_id();418		self.execute_expression_empty(format!("sess_field_{id} = {}", expr.as_ref()))419			.await?;420		Ok(id)421	}422423	/// Id should be immediately used424	fn allocate_id(&mut self) -> u32 {425		if let Some(free) = self.free_list.pop() {426			free427		} else {428			let v = self.next_id;429			self.next_id += 1;430			v431		}432	}433	// Nix has no way to deallocate variable, yet GC will correct everything not reachable.434	// async fn free_id(&mut self, id: u32) -> Result<()> {435	// 	self.execute_expression_empty(format!("sess_field_{id} = null"))436	// 		.await?;437	// 	self.free_list.push(id);438	// 	Ok(())439	// }440}