telemetry [14]: report CLI and reboot command failures

crates/trident/src/lib.rs

@@ -64,8 +64,8 @@ pub use crate::{
         logstream::Logstream,
         operation_context::{
             command_name, run_command, run_command_if, run_reboot_command,
-            run_with_captured_operation, run_with_operation, save_reboot_operation,
-            take_reboot_operation, OperationSource,
+            run_with_captured_operation, save_reboot_operation, take_reboot_operation,
+            OperationSource,
         },
         telemetry_uploader::TelemetryUploader,
         tracestream::TraceStream,

crates/trident/src/logging/appinsights.rs

@@ -480,7 +480,7 @@ mod tests {
                 )
                 .unwrap();
 
-                let mocks = operation_context::run_with_operation(&command, source, || {
+                operation_context::run_with_operation(&command, source, || {
                     let operation_id = operation_context::current().unwrap().0;
                     let servicing_id = if command.generates_servicing_id() {
                         operation_id.as_str()
@@ -525,12 +525,12 @@ mod tests {
                         mocks.push(mock);
                         sender.send_event(name, properties);
                     }
-                    mocks
+                    assert!(uploader.flush_with_deadline(Duration::from_secs(5)));
+                    for mock in mocks {
+                        mock.assert();
+                    }
                 });
                 uploader.shutdown_with_deadline(Duration::from_secs(5));
-                for mock in mocks {
-                    mock.assert();
-                }
             }
         }
     }

crates/trident/src/logging/telemetry_uploader.rs

@@ -223,7 +223,6 @@ impl TelemetryUploadHandle {
     /// treated as success; use [`Self::upload_with_validator`] if the
     /// destination's ingestion protocol can reject part of a request
     /// while still returning a 2xx status.
-    #[cfg(test)]
     pub fn upload(
         &self,
         url: &Url,

crates/trident/src/main.rs

@@ -10,10 +10,9 @@ use trident::{
     cli::{self, Cli, Commands, GetKind, TridentExitCodes},
     init::offline,
     manual_rollback::{self, utils::ManualRollbackRequestKind},
-    run_with_captured_operation, run_with_operation, save_reboot_operation, take_reboot_operation,
-    validation, AppInsightsSender, BackgroundLog, BackgroundUploader, DataStore, ExitKind,
-    LogForwarder, Logstream, OperationSource, TelemetryUploader, TraceStream, Trident,
-    TRIDENT_BACKGROUND_LOG_PATH,
+    run_command, run_reboot_command, save_reboot_operation, validation, AppInsightsSender,
+    BackgroundLog, BackgroundUploader, DataStore, ExitKind, LogForwarder, Logstream,
+    OperationSource, TelemetryUploader, TraceStream, Trident, TRIDENT_BACKGROUND_LOG_PATH,
 };
 use trident_api::{
     config::HostConfigurationSource,
@@ -49,7 +48,17 @@ fn run_trident(
         proxy_status("NO_PROXY"),
     );
 
-    // Catch exit fast commands
+    // Fast-exit commands: read-only/one-shot commands that never start a
+    // servicing run (validate, get, diagnose, offline-initialize, a manual
+    // rollback --check, start-network). These deliberately run outside
+    // run_command below -- no command_start/command_error telemetry is
+    // emitted for them. Their failures (a malformed --config, a datastore
+    // that can't be opened, a diagnostics bundle that can't be written,
+    // etc.) are operator-input or read errors, not servicing outcomes; a
+    // genuine underlying datastore/host problem still gets telemetry when
+    // the actual servicing operation (install/update/etc.) that triggered
+    // it runs. Handled here, before `command`/run_command are even set up,
+    // so none of that machinery needs to reason about them.
     match &args.command {
         Commands::Validate { config } => {
             return validation::validate_host_config_file(config).map(|()| ExitKind::Done);
@@ -123,96 +132,97 @@ fn run_trident(
         _ => (),
     }
 
-    let res = panic::catch_unwind(move || {
-        match &args.command {
-            Commands::Install { status, error, .. }
-            | Commands::Update { status, error, .. }
-            | Commands::Commit { status, error }
-            | Commands::RebuildRaid { status, error, .. }
-            | Commands::Rollback { status, error, .. } => {
-                // Determined before any preflight checks below, and used
-                // to wrap the *entire* servicing branch (preflight checks,
-                // Trident::new, and the actual command) in a single
-                // run_with_operation call -- not just the innermost
-                // install/update/commit/rollback/rebuild-raid call, as
-                // before. That previously left Trident::new (and
-                // everything it does, including firing "trident_start")
-                // outside any operation context: every event from CLI
-                // startup through to just before the actual command ran
-                // had no operation_id/command.
-                // `Commands::kind()` is total over every variant
-                // (including ones that can never actually reach this
-                // arm), so it renders the same wire name the old
-                // per-variant match did without needing an
-                // `unreachable!()` fallback to prove exhaustiveness.
-                let kind = args.command.kind();
-
-                // Determined up front, before any preflight checks below,
-                // so a missing/nonexistent --config is rejected immediately.
-                let config_path = match &args.command {
-                    Commands::Update { config, .. } | Commands::Install { config, .. } => {
-                        Some(config.clone())
-                    }
-                    Commands::RebuildRaid { config, .. } => config.clone(),
-                    _ => None,
-                };
-                if let Some(path) = &config_path {
-                    if !path.exists() {
-                        return run_with_operation(&kind, OperationSource::Cli, || {
-                            Err(TridentError::new(InvalidInputError::ReadInputFile {
-                                path: path.to_string_lossy().to_string(),
-                            }))
-                            .message("Config file does not exist")
-                        });
-                    }
-                }
+    // Only servicing commands reach here: Install, Update, Commit,
+    // RebuildRaid, and a non-check Rollback. These get command_start/
+    // command_error telemetry via run_command below; the fast-exit
+    // commands above already returned without any. `Commands::kind()` is
+    // total over every variant (including the fast-exit ones and
+    // Daemon/GrpcClient, none of which can actually reach here), so it
+    // renders the same wire name the old per-variant match did without
+    // needing `unreachable!()` arms to prove exhaustiveness.
+    let kind = args.command.kind();
+
+    // Determine this up front, before the pre-warm attach below: a
+    // multiboot install may swap to a brand-new temporary datastore
+    // inside `Trident::install` (see there), distinct from
+    // `datastore_path` here (the existing host's persistent datastore) --
+    // so skip attaching an installation ID from the pre-swap datastore
+    // until `install` has settled on which one it actually uses, rather
+    // than attaching the existing host's here and having it (and every
+    // event emitted before the swap decision) be wrong for the rest of
+    // the run. See `new_deferring_installation_id`'s doc comment for the
+    // full rationale.
+    let defer_installation_id = matches!(
+        args.command,
+        Commands::Install {
+            multiboot: true,
+            ..
+        }
+    );
 
-                // Determine this up front, before the pre-warm attach
-                // below: a multiboot install may swap to a brand-new
-                // temporary datastore inside `Trident::install` (see
-                // there), distinct from `datastore_path` here (the
-                // existing host's persistent datastore) -- so skip
-                // attaching an installation ID from the pre-swap datastore
-                // until `install` has settled on which one it actually
-                // uses, rather than attaching the existing host's here and
-                // having it (and every event emitted before the swap
-                // decision) be wrong for the rest of the run. See
-                // `new_deferring_installation_id`'s doc comment for the
-                // full rationale.
-                let defer_installation_id = matches!(
-                    args.command,
-                    Commands::Install {
-                        multiboot: true,
-                        ..
-                    }
-                );
+    // Pre-warm persisted IDs for errors that occur before Trident::new can
+    // attach them. Startup events omit both IDs even when they are cached.
+    // This does not create a datastore, but may migrate a legacy installation
+    // ID. Multiboot defers installation-ID attachment until its staging
+    // datastore inherits the source identity.
+    // Load once and reuse the same snapshot inside the run_command
+    // closure below, rather than reloading there: calling
+    // `AgentConfig::load()` a second time could observe a different
+    // `DatastorePath` (e.g. a CIH bootstrap swap between the two reads),
+    // leaving the ID attached to `tracestream` here attributed to a
+    // different datastore than the one the operation actually runs
+    // against.
+    let agent_config_result = AgentConfig::load();
+    if let Ok(agent_config) = &agent_config_result {
+        tracestream.attach_ids_if_present(agent_config.datastore_path(), defer_installation_id);
+    }
 
-                // Pre-warm persisted IDs for errors that occur before Trident::new can
-                // attach them. Startup events omit both IDs even when they are cached.
-                // This does not create a datastore, but may migrate a legacy installation
-                // ID. Multiboot defers installation-ID attachment until its staging
-                // datastore inherits the source identity.
-                // Load once and reuse the same snapshot for both the
-                // pre-warm attach here and the operation closure below:
-                // calling `AgentConfig::load()` a second time inside the
-                // closure could observe a different `DatastorePath` (e.g. a
-                // CIH bootstrap swap between the two reads), leaving the
-                // ID cached on `tracestream` here attributed to a
-                // different datastore than the one the operation actually
-                // runs against.
-                let agent_config_result = AgentConfig::load();
-                if let Ok(agent_config) = &agent_config_result {
-                    tracestream.attach_ids_if_present(
-                        agent_config.datastore_path(),
-                        defer_installation_id,
-                    );
-                }
+    // Determined up front so a missing/nonexistent --config is rejected
+    // immediately, before run_command below even fires command_start.
+    let config_path = match &args.command {
+        Commands::Update { config, .. } | Commands::Install { config, .. } => Some(config.clone()),
+        Commands::RebuildRaid { config, .. } => config.clone(),
+        _ => None,
+    };
+    if let Some(path) = &config_path {
+        if !path.exists() {
+            return run_command(&kind, OperationSource::Cli, || {
+                Err(TridentError::new(InvalidInputError::ReadInputFile {
+                    path: path.to_string_lossy().to_string(),
+                }))
+                .message("Config file does not exist")
+            });
+        }
+    }
 
-                run_with_operation(&kind, OperationSource::Cli, || {
+    // run_command itself now catches a panic from its closure (while the
+    // operation context is still active) and fires command_error before
+    // re-raising it, so a genuine panic gets the same telemetry as a
+    // normal Err. This outer catch_unwind remains as a safety net for a
+    // panic occurring outside run_command's closure (e.g. in run_command's
+    // own setup) and to keep converting an unwound panic into a non-zero
+    // exit code below.
+    let res = panic::catch_unwind(move || {
+        run_command(&kind, OperationSource::Cli, || {
+            match &args.command {
+                Commands::Install { status, error, .. }
+                | Commands::Update { status, error, .. }
+                | Commands::Commit { status, error }
+                | Commands::RebuildRaid { status, error, .. }
+                | Commands::Rollback { status, error, .. } => {
                     // config_path was already validated (existence-checked)
-                    // above. Reuse the same `AgentConfig` snapshot loaded
-                    // just above rather than reloading -- see the comment
-                    // there.
+                    // above.
+                    let config_path = match &args.command {
+                        Commands::Update { config, .. } | Commands::Install { config, .. } => {
+                            Some(config.clone())
+                        }
+                        Commands::RebuildRaid { config, .. } => config.clone(),
+                        _ => None,
+                    };
+
+                    // Reuse the same `AgentConfig` snapshot loaded above
+                    // (before this closure/panic::catch_unwind) rather
+                    // than reloading -- see the comment there.
                     let agent_config = agent_config_result?;
                     // For commands that cannot themselves stage a new
                     // install/update (see
@@ -316,7 +326,7 @@ fn run_trident(
 
                     // Capture this operation's identity while its context
                     // is still installed (this closure runs entirely
-                    // inside `run_with_operation`'s scope), so the reboot
+                    // inside `run_command`'s scope), so the reboot
                     // requested below by the caller can be tagged with the
                     // *original* install/update/etc.'s `operation_id`/
                     // `command` instead of a disconnected fresh one -- see
@@ -326,10 +336,10 @@ fn run_trident(
                     }
 
                     res.message(format!("Failed to execute '{}' command", args.command))
-                })
+                }
+                _ => unreachable!(),
             }
-            _ => unreachable!(),
-        }
+        })
     });
 
     match res {
@@ -740,16 +750,18 @@ fn main() -> ExitCode {
             Ok(ExitKind::NeedsReboot) => {
                 // Reuse the just-completed install/update/etc.'s own
                 // operation_id/command (captured via save_reboot_operation
-                // just before that command's own run_with_operation scope
-                // ended) rather than leaving `trident_system_reboot`
-                // untagged, or minting an unrelated fresh "reboot"
-                // identity: the reboot is a direct continuation of that
-                // same servicing operation, not an independent one, so
-                // telemetry should correlate it back to the same
-                // operation_id. Falls back to a plain, untagged call if
-                // nothing was captured (shouldn't happen on this path, but
-                // avoids losing the reboot attempt entirely if it does).
-                if let Err(e) = run_with_captured_operation(take_reboot_operation(), || {
+                // just before that command's own run_command scope ended)
+                // rather than leaving `trident_system_reboot` untagged, or
+                // minting an unrelated fresh "reboot" identity: the reboot
+                // is a direct continuation of that same servicing
+                // operation, not an independent one, so telemetry should
+                // correlate it back to the same operation_id.
+                // run_reboot_command also still fires command_error on a
+                // failed reboot -- falling back to a fresh, untagged
+                // command if nothing was captured -- matching every other
+                // command's error-reporting contract instead of silently
+                // dropping this one on the floor.
+                if let Err(e) = run_reboot_command(|| {
                     trident::request_reboot_with_wait_and_flush(|| telemetry_uploader.flush())
                 }) {
                     error!("Failed to reboot: {e:?}");
@@ -761,3 +773,47 @@ fn main() -> ExitCode {
         TridentExitCodes::Success.into()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pre_reboot_flush_keeps_uploader_open_for_errors() {
+        let mut server = mockito::Server::new();
+        let reboot = server
+            .mock("POST", "/reboot")
+            .with_status(200)
+            .expect(1)
+            .create();
+        let error = server
+            .mock("POST", "/error")
+            .with_status(200)
+            .expect(1)
+            .create();
+        let uploader = TelemetryUploader::new(4).unwrap();
+        let handle = uploader.get_handle().unwrap();
+        let guard = TelemetryUploaderGuard(Some(uploader));
+        let url = reqwest::Url::parse(&server.url()).unwrap();
+        handle
+            .upload(
+                &url.join("reboot").unwrap(),
+                b"reboot".to_vec(),
+                Duration::from_secs(2),
+                None,
+            )
+            .unwrap();
+        guard.flush();
+        reboot.assert();
+        handle
+            .upload(
+                &url.join("error").unwrap(),
+                b"failure".to_vec(),
+                Duration::from_secs(2),
+                None,
+            )
+            .unwrap();
+        drop(guard);
+        error.assert();
+    }
+}