Flush traces before dsc exits

Author: SteveL-MSFT

dsc/src/main.rs

@@ -10,6 +10,7 @@ use rust_i18n::{i18n, t};
 use std::{io, process::exit};
 use sysinfo::{Process, RefreshKind, System, get_current_pid, ProcessRefreshKind};
 use tracing::{error, info, warn, debug};
+use util::flush_and_shutdown_tracing;
 
 use crate::util::{EXIT_INVALID_INPUT, get_input};
 
@@ -68,6 +69,7 @@ fn main() {
                         Ok(merged) => Some(merged),
                         Err(err) => {
                             error!("{}: {err}", t!("main.failedMergingParameters"));
+                            flush_and_shutdown_tracing();
                             exit(EXIT_INVALID_INPUT);
                         }
                     }
@@ -88,8 +90,10 @@ fn main() {
         SubCommand::Mcp => {
             if let Err(err) = start_mcp_server() {
                 error!("{}", t!("main.failedToStartMcpServer", error = err));
+                flush_and_shutdown_tracing();
                 exit(util::EXIT_MCP_FAILED);
             }
+            flush_and_shutdown_tracing();
             exit(util::EXIT_SUCCESS);
         }
         SubCommand::Resource { subcommand } => {
@@ -101,13 +105,15 @@ fn main() {
                 Ok(json) => json,
                 Err(err) => {
                     error!("JSON: {err}");
+                    flush_and_shutdown_tracing();
                     exit(util::EXIT_JSON_ERROR);
                 }
             };
             util::write_object(&json, output_format.as_ref(), false);
         },
     }
 
+    flush_and_shutdown_tracing();
     exit(util::EXIT_SUCCESS);
 }
 
@@ -119,15 +125,18 @@ fn ctrlc_handler() {
     info!("{}: {}", t!("main.foundProcesses"), sys.processes().len());
     let Ok(current_pid) = get_current_pid() else {
         error!("{}", t!("main.failedToGetPid"));
+        flush_and_shutdown_tracing();
         exit(util::EXIT_CTRL_C);
     };
     info!("{}: {}", t!("main.currentPid"), current_pid);
     let Some(current_process) = sys.process(current_pid) else {
         error!("{}", t!("main.failedToGetProcess"));
+        flush_and_shutdown_tracing();
         exit(util::EXIT_CTRL_C);
     };
 
     terminate_subprocesses(&sys, current_process);
+    flush_and_shutdown_tracing();
     exit(util::EXIT_CTRL_C);
 }
 
@@ -192,6 +201,7 @@ fn check_store() {
         eprintln!("{}", t!("main.storeMessage"));
         // wait for keypress
         let _ = io::stdin().read(&mut [0u8]).unwrap();
+        flush_and_shutdown_tracing();
         exit(util::EXIT_INVALID_ARGS);
     }
 }

dsc/src/util.rs

@@ -53,6 +53,7 @@ use std::env;
 use std::io::{IsTerminal, Read, stdout, Write};
 use std::path::Path;
 use std::process::exit;
+use std::sync::Once;
 use syntect::{
     easy::HighlightLines,
     highlighting::ThemeSet,
@@ -690,3 +691,23 @@ pub fn merge_parameters(file_params: &str, inline_params: &str) -> Result<String
     let merged = Value::Object(file_map);
     Ok(serde_json::to_string(&merged)?)
 }
+
+static FLUSH_ONCE: Once = Once::new();
+
+/// Flush and shutdown tracing to ensure all traces are written before exit.
+/// This function ensures that any pending trace writes are completed and
+/// background writer threads have time to finish their work.
+pub fn flush_and_shutdown_tracing() {
+    FLUSH_ONCE.call_once(|| {
+        // Force any pending writes to complete
+        if let Err(e) = std::io::stderr().flush() {
+            eprintln!("Failed to flush stderr: {}", e);
+        }
+        if let Err(e) = std::io::stdout().flush() {
+            eprintln!("Failed to flush stdout: {}", e);
+        }
+
+        // Small delay to ensure async writes complete
+        std::thread::sleep(std::time::Duration::from_millis(50));
+    });
+}