feat: proactively detect when stdout goes away on Unix / fixes for streaming command output (#53)

Author: cablehead

Cargo.lock

@@ -58,6 +58,9 @@ webpki-roots = "0.26.6"
 base64 = "0.22.1"
 tracing-subscriber = "0.3.19"
 
+[target.'cfg(unix)'.dependencies]
+nix = { version = "0.29", default-features = false, features = ["poll"] }
+
 [dev-dependencies]
 assert_cmd = "2.0.14"
 duct = "0.13.7"

Cargo.toml

@@ -2,17 +2,27 @@ use scru128::Scru128Id;
 use std::collections::HashMap;
 use tracing::instrument;
 
+use serde::{Deserialize, Serialize};
+
 use crate::error::Error;
 use crate::nu;
 use crate::nu::commands;
 use crate::nu::util::value_to_json;
-use crate::store::{FollowOption, Frame, ReadOptions, Store};
+use crate::store::{FollowOption, Frame, ReadOptions, Store, TTL};
+
+// TODO: DRY with handlers
+#[derive(Clone, Debug, Serialize, Deserialize, Default)]
+pub struct ReturnOptions {
+    pub suffix: Option<String>,
+    pub ttl: Option<TTL>,
+}
 
 #[derive(Clone)]
 struct Command {
     id: Scru128Id,
     engine: nu::Engine,
     definition: String,
+    return_options: Option<ReturnOptions>,
 }
 
 async fn handle_define(
@@ -59,7 +69,6 @@ pub async fn serve(
         if frame.topic == "xs.threshold" {
             break;
         }
-
         if let Some(name) = frame.topic.strip_suffix(".define") {
             handle_define(&frame, name, &base_engine, &store, &mut commands).await;
         }
@@ -70,8 +79,23 @@ pub async fn serve(
         if let Some(name) = frame.topic.strip_suffix(".define") {
             handle_define(&frame, name, &base_engine, &store, &mut commands).await;
         } else if let Some(name) = frame.topic.strip_suffix(".call") {
-            if let Some(command) = commands.get(name) {
-                execute_command(command.clone(), frame, &store).await?;
+            let name = name.to_owned();
+            if let Some(command) = commands.get(&name) {
+                let store = store.clone();
+                let frame = frame.clone();
+                let command = command.clone();
+                tokio::spawn(async move {
+                    if let Err(e) = execute_command(command, &frame, &store).await {
+                        tracing::error!("Failed to execute command '{}': {:?}", name, e);
+                        let _ = store.append(
+                            Frame::builder(format!("{}.error", name), frame.context_id)
+                                .meta(serde_json::json!({
+                                    "error": e.to_string(),
+                                }))
+                                .build(),
+                        );
+                    }
+                });
             }
         }
     }
@@ -86,12 +110,12 @@ async fn register_command(
 ) -> Result<Command, Error> {
     // Get definition from CAS
     let hash = frame.hash.as_ref().ok_or("Missing hash field")?;
-    let definition = store.cas_read(hash).await?;
-    let definition = String::from_utf8(definition)?;
+    let definition_bytes = store.cas_read(hash).await?;
+    let definition = String::from_utf8(definition_bytes)?;
 
     let mut engine = base_engine.clone();
 
-    // Add addtional commands, scoped to this command's context
+    // Add additional commands, scoped to this command's context
     engine.add_commands(vec![
         Box::new(commands::cat_command::CatCommand::new(
             store.clone(),
@@ -103,10 +127,14 @@ async fn register_command(
         )),
     ])?;
 
+    // Parse the command configuration to extract return_options (ignore the process closure here)
+    let (_closure, return_options) = parse_command_definition(&mut engine, &definition)?;
+
     Ok(Command {
         id: frame.id,
         engine,
         definition,
+        return_options,
     })
 }
 
@@ -120,8 +148,9 @@ async fn register_command(
         )
     )
 )]
-async fn execute_command(command: Command, frame: Frame, store: &Store) -> Result<(), Error> {
+async fn execute_command(command: Command, frame: &Frame, store: &Store) -> Result<(), Error> {
     let store = store.clone();
+    let frame = frame.clone();
 
     tokio::task::spawn_blocking(move || {
         let base_meta = serde_json::json!({
@@ -139,19 +168,34 @@ async fn execute_command(command: Command, frame: Frame, store: &Store) -> Resul
             ),
         )])?;
 
-        let closure = parse_command_definition(&mut engine, &command.definition)?;
+        let (closure, _) = parse_command_definition(&mut engine, &command.definition)?;
 
         // Run command and process pipeline
         match run_command(&engine, closure, &frame) {
             Ok(pipeline_data) => {
+                let recv_suffix = command
+                    .return_options
+                    .as_ref()
+                    .and_then(|opts| opts.suffix.as_deref())
+                    .unwrap_or(".recv");
+                let ttl = command
+                    .return_options
+                    .as_ref()
+                    .and_then(|opts| opts.ttl.clone());
+
                 // Process each value as a .recv event
                 for value in pipeline_data {
                     let hash = store.cas_insert_sync(value_to_json(&value).to_string())?;
                     let _ = store.append(
                         Frame::builder(
-                            format!("{}.recv", frame.topic.strip_suffix(".call").unwrap()),
+                            format!(
+                                "{}{}",
+                                frame.topic.strip_suffix(".call").unwrap(),
+                                recv_suffix
+                            ),
                             frame.context_id,
                         )
+                        .maybe_ttl(ttl.clone())
                         .hash(hash)
                         .meta(serde_json::json!({
                             "command_id": command.id.to_string(),
@@ -224,7 +268,7 @@ fn run_command(
 fn parse_command_definition(
     engine: &mut nu::Engine,
     script: &str,
-) -> Result<nu_protocol::engine::Closure, Error> {
+) -> Result<(nu_protocol::engine::Closure, Option<ReturnOptions>), Error> {
     let mut working_set = nu_protocol::engine::StateWorkingSet::new(&engine.state);
     let block = nu_parser::parse(&mut working_set, None, script.as_bytes(), false);
 
@@ -240,12 +284,36 @@ fn parse_command_definition(
 
     let config = result.into_value(nu_protocol::Span::unknown())?;
 
+    // Get the process closure (required)
     let process = config
         .get_data_by_key("process")
         .ok_or("No 'process' field found in command configuration")?
         .into_closure()?;
 
+    // Optionally parse return_options (using the same approach as in handlers)
+    let return_options = if let Some(return_config) = config.get_data_by_key("return_options") {
+        let record = return_config
+            .as_record()
+            .map_err(|_| "return must be a record")?;
+
+        let suffix = record
+            .get("suffix")
+            .map(|v| v.as_str().map_err(|_| "suffix must be a string"))
+            .transpose()?
+            .map(String::from);
+
+        let ttl = record
+            .get("ttl")
+            .map(|v| serde_json::from_str(&value_to_json(v).to_string()))
+            .transpose()
+            .map_err(|e| format!("invalid TTL: {}", e))?;
+
+        Some(ReturnOptions { suffix, ttl })
+    } else {
+        None
+    };
+
     engine.state.merge_env(&mut stack)?;
 
-    Ok(process)
+    Ok((process, return_options))
 }

src/commands/serve.rs

@@ -247,7 +247,6 @@ async fn cat(args: CommandCat) -> Result<(), Box<dyn std::error::Error + Send +
         .as_deref()
         .and_then(|context| scru128::Scru128Id::from_str(context).ok())
         .or_else(|| (!args.all).then_some(ZERO_CONTEXT));
-
     let last_id = if let Some(last_id) = &args.last_id {
         match scru128::Scru128Id::from_str(last_id) {
             Ok(id) => Some(id),
@@ -256,7 +255,6 @@ async fn cat(args: CommandCat) -> Result<(), Box<dyn std::error::Error + Send +
     } else {
         None
     };
-
     // Build options in one chain
     let options = ReadOptions::builder()
         .tail(args.tail)
@@ -271,19 +269,77 @@ async fn cat(args: CommandCat) -> Result<(), Box<dyn std::error::Error + Send +
         .maybe_limit(args.limit.map(|l| l as usize))
         .maybe_context_id(context_id)
         .build();
-
     let mut receiver = xs::client::cat(&args.addr, options, args.sse).await?;
     let mut stdout = tokio::io::stdout();
 
-    match async {
-        while let Some(bytes) = receiver.recv().await {
-            stdout.write_all(&bytes).await?;
-            stdout.flush().await?;
+    #[cfg(unix)]
+    let result = {
+        use nix::unistd::dup;
+        use std::io::Write;
+        use std::os::unix::io::{AsRawFd, FromRawFd};
+        use tokio::io::unix::AsyncFd;
+
+        let stdout_fd = std::io::stdout().as_raw_fd();
+        // Create a duplicate of the file descriptor so we can check it separately
+        let dup_fd = dup(stdout_fd)?;
+        let stdout_file = unsafe { std::fs::File::from_raw_fd(dup_fd) };
+        let async_fd = AsyncFd::new(stdout_file)?;
+
+        async {
+            loop {
+                tokio::select! {
+                    maybe_bytes = receiver.recv() => {
+                        match maybe_bytes {
+                            Some(bytes) => {
+                                if let Err(e) = stdout.write_all(&bytes).await {
+                                    if e.kind() == std::io::ErrorKind::BrokenPipe {
+                                        break;
+                                    }
+                                    return Err(e);
+                                }
+                                stdout.flush().await?;
+                            }
+                            None => break,
+                        }
+                    },
+                    Ok(mut guard) = async_fd.writable() => {
+                        // Try a zero-byte write to check if stdout is closed
+                        match guard.try_io(|inner| {
+                            // We're just doing a test write to check for EPIPE
+                            match inner.get_ref().write(&[]) {
+                                Ok(_) => Ok(false), // Stdout is still open
+                                Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => Ok(true), // Stdout closed
+                                Err(e) => Err(e), // Other error
+                            }
+                        }) {
+                            Ok(Ok(true)) => break, // Stdout is closed
+                            Ok(Ok(false)) => { /* Stdout still open */ },
+                            Ok(Err(e)) => return Err(e), // Error occurred
+                            Err(_) => { /* Would block, try again later */ }
+                        }
+
+                        guard.clear_ready();
+                    }
+                }
+            }
+            Ok::<_, std::io::Error>(())
         }
-        Ok::<_, std::io::Error>(())
-    }
-    .await
-    {
+        .await
+    };
+
+    #[cfg(not(unix))]
+    let result = {
+        async {
+            while let Some(bytes) = receiver.recv().await {
+                stdout.write_all(&bytes).await?;
+                stdout.flush().await?;
+            }
+            Ok::<_, std::io::Error>(())
+        }
+        .await
+    };
+
+    match result {
         Ok(_) => Ok(()),
         Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
         Err(e) => Err(e.into()),

src/main.rs

@@ -17,7 +17,7 @@ def conditional-pipe [
   condition: bool
   action: closure
 ] {
-  if $condition { do $action } else { $in }
+  if $condition { do $action } else { }
 }
 
 export def xs-addr [] {
@@ -94,7 +94,7 @@ export def .cat [
     limit: $limit
     context: (if not $all { (xs-context $context (metadata $context).span) })
     all: $all
-  } | conditional-pipe (not ($detail or $all)) { reject context_id ttl }
+  } | conditional-pipe (not ($detail or $all)) { each { reject context_id ttl } }
 }
 
 def read_hash [hash?: any] {