encapsulate flush into exit()

Author: SteveL-MSFT

dsc/src/main.rs

@@ -7,10 +7,10 @@ use clap_complete::generate;
 use dsc_lib::progress::ProgressFormat;
 use mcp::start_mcp_server;
 use rust_i18n::{i18n, t};
-use std::{io, process::exit};
+use std::io;
 use sysinfo::{Process, RefreshKind, System, get_current_pid, ProcessRefreshKind};
 use tracing::{error, info, warn, debug};
-use util::flush_and_shutdown_tracing;
+use util::exit;
 
 use crate::util::{EXIT_INVALID_INPUT, get_input};
 
@@ -69,7 +69,6 @@ fn main() {
                         Ok(merged) => Some(merged),
                         Err(err) => {
                             error!("{}: {err}", t!("main.failedMergingParameters"));
-                            flush_and_shutdown_tracing();
                             exit(EXIT_INVALID_INPUT);
                         }
                     }
@@ -90,10 +89,8 @@ 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 } => {
@@ -105,15 +102,13 @@ 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);
 }
 
@@ -125,18 +120,15 @@ 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);
 }
 
@@ -201,7 +193,6 @@ 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

@@ -52,8 +52,6 @@ use std::collections::HashMap;
 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,
@@ -692,22 +690,17 @@ pub fn merge_parameters(file_params: &str, inline_params: &str) -> Result<String
     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);
-        }
+/// Exit the process with the given code after flushing and shutting down tracing.
+pub fn exit(code: i32) -> ! {
+    // 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));
-    });
+    // Small delay to ensure async writes complete
+    std::thread::sleep(std::time::Duration::from_millis(50));
+    std::process::exit(code);
 }