Exit successfully when user closes stdout (#1360)

Author: ia0

rust/cli/CHANGELOG.md

@@ -1,9 +1,14 @@
 # Changelog
 
-## 1.0.3-dev
+## 1.1.0-dev
+
+### Minor
+
+- Exit successfully when the standard output is closed by the user
 
 ### Patch
 
+- Join all threads before shutdown to avoid segmentation faults in ONNX Runtime
 - Update dependencies
 - Fix new clippy lints
 

rust/cli/Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "magika-cli"
-version = "1.0.3-dev"
+version = "1.1.0-dev"
 authors = ["Magika Developers <magika-dev@google.com>"]
 license = "Apache-2.0"
 edition = "2021"

rust/cli/Cargo.toml

@@ -15,7 +15,7 @@
 use std::borrow::Cow;
 use std::collections::HashMap;
 use std::fmt::Write;
-use std::io::ErrorKind;
+use std::io::{ErrorKind, Write as _};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
@@ -30,6 +30,7 @@ use ort::session::builder::GraphOptimizationLevel;
 use serde::Serialize;
 use tokio::fs::File;
 use tokio::io::AsyncReadExt;
+use tokio::task::JoinSet;
 
 /// Determines file content types using AI.
 #[derive(Parser)]
@@ -158,6 +159,13 @@ struct Experimental {
 
 #[tokio::main]
 async fn main() -> Result<()> {
+    let mut tasks = JoinSet::new();
+    let result = start(&mut tasks).await;
+    while tasks.join_next().await.is_some() {}
+    result
+}
+
+async fn start(tasks: &mut JoinSet<()>) -> Result<()> {
     let mut flags = Flags::parse();
     ensure!(0 < flags.experimental.batch_size, "--batch-size cannot be zero");
     // If --num-tasks is set, we don't do any guessing.
@@ -183,33 +191,49 @@ async fn main() -> Result<()> {
     if flags.colors.disable {
         colored::control::set_override(false);
     }
-    let (result_sender, mut result_receiver) =
+    let (result_sender, result_receiver) =
         tokio::sync::mpsc::channel::<Result<Response>>(num_tasks * flags.experimental.batch_size);
     let (batch_sender, batch_receiver) = async_channel::bounded::<Batch>(num_tasks);
-    tokio::spawn({
+    tasks.spawn({
         let flags = flags.clone();
         let result_sender = result_sender.clone();
         async move {
             if let Err(e) = extract_features(&flags, &batch_sender, &result_sender).await {
-                result_sender.send(Err(e)).await.unwrap();
+                let _ = result_sender.send(Err(e)).await;
             }
         }
     });
     for _ in 0..num_tasks {
         let mut magika = build_session(&flags)?;
-        tokio::spawn({
+        tasks.spawn({
             let batch_receiver = batch_receiver.clone();
             let result_sender = result_sender.clone();
             async move {
                 if let Err(e) = infer_batch(&mut magika, &batch_receiver, &result_sender).await {
-                    result_sender.send(Err(e)).await.unwrap();
+                    let _ = result_sender.send(Err(e)).await;
                 }
             }
         });
     }
     drop(result_sender);
+    match print(&flags, result_receiver).await {
+        Err(e)
+            if e.root_cause()
+                .downcast_ref::<std::io::Error>()
+                .is_some_and(|x| x.kind() == std::io::ErrorKind::BrokenPipe) =>
+        {
+            Ok(())
+        }
+        x => x,
+    }
+}
+
+async fn print(
+    flags: &Flags, mut result_receiver: tokio::sync::mpsc::Receiver<Result<Response>>,
+) -> Result<()> {
+    let mut stdout = std::io::stdout().lock();
     if flags.format.json {
-        print!("[");
+        write!(stdout, "[")?;
     }
     let mut reorder = Reorder::default();
     let mut errors = false;
@@ -219,22 +243,22 @@ async fn main() -> Result<()> {
             errors |= response.result.is_err();
             if flags.format.json {
                 if reorder.next != 1 {
-                    print!(",");
+                    write!(stdout, ",")?;
                 }
                 for line in serde_json::to_string_pretty(&response.json()?)?.lines() {
-                    print!("\n  {line}");
+                    write!(stdout, "\n  {line}")?;
                 }
             } else {
-                println!("{}", response.format(&flags)?);
+                writeln!(stdout, "{}", response.format(flags)?)?;
             }
         }
     }
     debug_assert!(reorder.is_empty());
     if flags.format.json {
         if reorder.next != 0 {
-            println!();
+            writeln!(stdout)?;
         }
-        println!("]");
+        writeln!(stdout, "]")?;
     }
     if errors {
         std::process::exit(1);

rust/cli/src/main.rs

@@ -1,4 +1,4 @@
-#!/bin/sh
+#!/bin/bash
 # Copyright 2024 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -54,3 +54,7 @@ test_error '/etc/shadow' "\
 test_error 'non_existent src/main.rs' "\
 non_existent: No such file or directory (os error 2) (error)
 src/main.rs: Rust source (code)"
+
+info "Test exit code with broken pipe"
+magika -r ../../tests_data | head -n1 >/dev/null
+[ "${PIPESTATUS[0]}" -eq 0 ] || error "non-zero exit code with broken pipe"