telemetry [5]: correlate staged servicing operations

crates/trident/Cargo.toml

@@ -35,7 +35,7 @@ oci-client = { workspace = true }
 procfs = { workspace = true }
 rayon = { workspace = true }
 regex = { workspace = true }
-reqwest = { workspace = true }
+reqwest = { workspace = true, features = ["stream"] }
 semver = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
@@ -64,7 +64,7 @@ hyper = { workspace = true }
 prost-types = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true }
-tokio-util = { workspace = true }
+tokio-util = { workspace = true, features = ["io-util"] }
 tonic = { workspace = true }
 tonic-middleware = { workspace = true }
 

crates/trident/src/datastore.rs

@@ -22,6 +22,14 @@ use crate::{logging::operation_context, TRIDENT_SEMVER_VERSION};
 /// can be correlated.
 const INSTALLATION_ID_KEY: &str = "installation-id";
 
+/// Key under which the current servicing ID is stored in the generic
+/// key-value table. Unlike `INSTALLATION_ID_KEY`, this is overwritten
+/// every time a new servicing operation begins staging (install, update,
+/// or manual rollback) via `DataStore::ensure_servicing_id` -- it identifies
+/// "the servicing operation in progress (or last completed)", not the
+/// host installation as a whole.
+const SERVICING_ID_KEY: &str = "servicing-id";
+
 pub struct DataStore {
     db: Option<sqlite::Connection>,
     host_status: HostStatus,
@@ -512,12 +520,11 @@ impl DataStore {
     /// `serde::Serialize`/`serde::de::DeserializeOwned` can be stored, not
     /// just `HostStatus`.
     ///
-    /// No first-party caller exists yet in this crate on its own -- this is
-    /// exercised only by tests for now. It's kept as part of the generic
-    /// key-value API (see `get_value`) for the first real consumer that
-    /// wants unconditional-overwrite semantics (as opposed to
-    /// `set_value_if_absent`'s insert-if-absent semantics).
-    #[allow(dead_code)]
+    /// `installation_id` needs insert-if-absent semantics instead (see
+    /// `set_value_if_absent`), but `servicing_id` -- a new servicing
+    /// operation genuinely does replace whatever ID (if any) came before
+    /// it -- is a first-party caller of this unconditional-overwrite
+    /// form.
     pub(crate) fn set_value<T: Serialize>(&self, key: &str, value: &T) -> Result<(), TridentError> {
         let contents = serde_json::to_string(value).structured(ServicingError::Datastore {
             inner: DatastoreError::SerializeValue {
@@ -688,6 +695,45 @@ impl DataStore {
         }
     }
 
+    /// Returns the currently persisted servicing ID, if any. `None` if no
+    /// servicing operation has ever staged (via `ensure_servicing_id`) on
+    /// this datastore.
+    pub fn servicing_id(&self) -> Result<Option<Uuid>, TridentError> {
+        self.get_value::<Uuid>(SERVICING_ID_KEY)
+    }
+
+    /// Generates a fresh servicing ID and persists it (unconditionally
+    /// overwriting any previous value), returning the new ID. Called once
+    /// at the start of staging for install, update, or manual rollback --
+    /// i.e. whenever `Operations::has_stage()` is true for that
+    /// invocation -- so a finalize-only call (or a later `commit`) must
+    /// use [`Self::servicing_id`] to read the value back instead of
+    /// calling this again.
+    ///
+    /// Like `ensure_installation_id`, the value used is this invocation's
+    /// own `operation_id` (see `logging::operation_context::current`),
+    /// not an unrelated freshly-generated UUID: whichever call actually
+    /// performs the staging is what should be correlatable, via this same
+    /// ID, with every later command (finalize, commit) that touches the
+    /// same servicing operation. Falls back to a fresh random UUID only if
+    /// called with no operation context active (should not happen for any
+    /// real caller).
+    ///
+    /// Unlike `ensure_installation_id`, this always generates a *new* ID
+    /// on every call -- there is no "first access wins"/get-or-create
+    /// semantics here, since a new servicing operation genuinely is a new
+    /// operation, not a value that should be stable for the datastore's
+    /// lifetime. (This also means, unlike installation_id, there is no
+    /// concurrent-first-access race to guard against: `set_value`'s plain
+    /// overwrite is sufficient.)
+    pub fn ensure_servicing_id(&mut self) -> Result<Uuid, TridentError> {
+        let id = operation_context::current()
+            .and_then(|(operation_id, _, _)| Uuid::parse_str(&operation_id).ok())
+            .unwrap_or_else(Uuid::new_v4);
+        self.set_value(SERVICING_ID_KEY, &id)?;
+        Ok(id)
+    }
+
     /// Parse a single HostStatus entry from a datastore query result.
     /// 1. Read each row as a string containing YAML-encoded Host Status.
     /// 2. Decode the YAML string into a serde_yaml Value.
@@ -1162,6 +1208,7 @@ mod tests {
                     let stale_id = staging.ensure_installation_id().unwrap();
                     assert_ne!(stale_id, installation_id);
                 }
+                let servicing_id = staging.ensure_servicing_id().unwrap();
                 staging.set_value("staged-image", &"image-hash").unwrap();
                 staging
                     .with_host_status(|status| {
@@ -1180,6 +1227,7 @@ mod tests {
                 let mut staging = super::DataStore::open_or_create(&staging_path).unwrap();
                 staging.inherit_installation_id(installation_id).unwrap();
                 assert_eq!(staging.installation_id().unwrap(), Some(installation_id));
+                assert_eq!(staging.servicing_id().unwrap(), Some(servicing_id));
                 assert_eq!(
                     staging.get_value::<String>("staged-image").unwrap(),
                     Some("image-hash".to_string())
@@ -1194,6 +1242,9 @@ mod tests {
                 for path in [&source_path, &staging_path, &target_path] {
                     let datastore = super::DataStore::open(path).unwrap();
                     assert_eq!(datastore.installation_id().unwrap(), Some(installation_id));
+                    if path != &source_path {
+                        assert_eq!(datastore.servicing_id().unwrap(), Some(servicing_id));
+                    }
                 }
             }
         }
@@ -1248,6 +1299,30 @@ mod tests {
 
         temp_dir.close().unwrap();
     }
+
+    #[test]
+    /// Regression test: `ensure_servicing_id` always generates a *fresh* ID
+    /// on every call, unlike `ensure_installation_id`'s get-or-create
+    /// semantics. This is what lets a retried clean install (e.g. a
+    /// previous attempt's staging failed, leaving
+    /// `ServicingState::NotProvisioned`) get its own distinct servicing ID
+    /// rather than reusing a stale one from the failed attempt.
+    fn test_ensure_servicing_id_generates_a_fresh_id_each_call() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let datastore_path = temp_dir.path().join("db.sqlite");
+
+        let mut datastore = super::DataStore::open_or_create(&datastore_path).unwrap();
+
+        let first = datastore.ensure_servicing_id().unwrap();
+        assert_eq!(datastore.servicing_id().unwrap(), Some(first));
+
+        let second = datastore.ensure_servicing_id().unwrap();
+        assert_ne!(
+            first, second,
+            "ensure_servicing_id should generate a fresh ID on every call, not reuse the previous one"
+        );
+        assert_eq!(datastore.servicing_id().unwrap(), Some(second));
+    }
 }
 
 #[cfg(feature = "functional-test")]

crates/trident/src/lib.rs

@@ -154,7 +154,8 @@ impl Trident {
     /// Construct without attaching the source datastore's installation ID.
     ///
     /// The CLI's multiboot path defers attachment until `install` has persisted
-    /// the shared identity in the selected staging datastore.
+    /// the shared identity in the selected staging datastore. Startup events
+    /// omit persisted identities regardless of which constructor is used.
     pub fn new_deferring_installation_id(
         config_source: Option<HostConfigurationSource>,
         datastore_path: &Path,
@@ -245,8 +246,33 @@ impl Trident {
             info!("Running Trident in a container");
         }
 
-        // Attach the persisted installation ID before startup metrics. Defer it
-        // when a multiboot install may select a different datastore.
+        // Attach this host's installation ID and its current servicing ID
+        // -- if either has already been stamped -- to the shared
+        // TraceStream before any startup metrics are emitted, so every
+        // trace/metric -- including this very "trident_start" event --
+        // carries them once available. Installation ID is normally
+        // created by `Trident::install` (at the start of staging), so
+        // every other caller of `Trident::new`
+        // (update/commit/rollback/rebuild-raid, and every daemon RPC
+        // handler) just attaches whatever was already persisted at
+        // install time. But not purely read-only: two other paths can
+        // also mint one here on first attach --
+        // `installation_id_or_migrate` performs a one-time migration
+        // write for a legacy/offline-provisioned datastore that predates
+        // this field (see `DataStore::installation_id_or_migrate`), and
+        // the CIH update-bootstrap path (below, in `update`) creates one
+        // via `ensure_and_attach_installation_id`. Neither ever creates a
+        // *datastore*, only (at most) writes into one that already
+        // exists. The servicing-ID half is fully read-only.
+        //
+        // Skips the installation-ID half specifically (but still attaches
+        // servicing ID) when `attach_installation_id` is false (see
+        // `new_deferring_installation_id`): attaching would stamp this
+        // invocation with the wrong (pre-swap) datastore's ID for a
+        // multiboot install that later swaps to a different datastore.
+        // Servicing ID is unaffected by that swap (a fresh temporary
+        // datastore simply has no servicing ID yet), so it's still worth
+        // attaching here even in that case.
         tracestream.attach_ids_if_present(datastore_path, !attach_installation_id);
 
         // Trace features enabled in the Host Configuration.
@@ -547,6 +573,14 @@ impl Trident {
                 warn!("Failed to attach installation ID: {e:?}");
             }
 
+            // Generate (if this call has `stage` allowed, from a source
+            // allowed to generate persistent IDs) or read back (otherwise)
+            // the servicing ID for this install, and attach it. See
+            // `TraceStream::refresh_servicing_id`'s doc comment for the
+            // full rationale (best-effort, fail-closed source gate,
+            // deliberate unconditional regeneration on `has_stage`).
+            tracestream.refresh_servicing_id(datastore, allowed_operations.has_stage());
+
             // Use a prefetched image if provided, otherwise load the image
             // specified in the Host Configuration.
             let image = match prefetched_image {
@@ -708,6 +742,18 @@ impl Trident {
                 }
             }
 
+            // Generate (if this call has `stage` allowed, from a source
+            // allowed to generate persistent IDs) or read back (otherwise)
+            // the servicing ID for this update, and attach it. Placed
+            // here -- after the datastore is guaranteed to exist (whether
+            // via the CIH bootstrap above or a pre-existing persistent
+            // one) -- so it covers both the CIH-bootstrap and normal
+            // update paths uniformly, unlike installation_id (which is
+            // attached earlier, in `Trident::new`, for the normal path).
+            // See `TraceStream::refresh_servicing_id`'s doc comment for
+            // the full rationale.
+            tracestream.refresh_servicing_id(datastore, allowed_operations.has_stage());
+
             // The storage section is optional for updates if COSI is in use.
             if host_config.image.is_some() && host_config.storage == Default::default() {
                 host_config.storage = datastore.host_status().spec.storage.clone();
@@ -861,6 +907,15 @@ impl Trident {
             return Ok((ExitKind::Done, ServicingType::NoActiveServicing));
         }
 
+        // `commit` never generates a servicing ID -- it only ever reads
+        // back whatever the preceding install/update/rollback stage
+        // persisted, so this invocation's own (fresh) `operation_id` can
+        // be correlated back to that servicing operation. Always passes
+        // `has_stage = false`, so this only ever reads (see
+        // `TraceStream::refresh_servicing_id`'s doc comment).
+        let tracestream = self.tracestream.clone();
+        tracestream.refresh_servicing_id(datastore, false);
+
         let servicing_type = match servicing_state {
             ServicingState::AbUpdateFinalized | ServicingState::AbUpdateHealthCheckFailed => {
                 ServicingType::AbUpdate
@@ -962,6 +1017,16 @@ impl Trident {
             return Ok((ExitKind::Done, ServicingType::NoActiveServicing));
         }
 
+        // Generate (if this call has `stage` allowed, from a source
+        // allowed to generate persistent IDs) or read back (otherwise) the
+        // servicing ID for this manual rollback, and attach it. Safe to do
+        // before `execute_and_record_error`: the datastore is already
+        // known-persistent per the state check above. See
+        // `TraceStream::refresh_servicing_id`'s doc comment for the full
+        // rationale.
+        let tracestream = self.tracestream.clone();
+        tracestream.refresh_servicing_id(datastore, allowed_operations.has_stage());
+
         let rollback_result = self.execute_and_record_error(datastore, |datastore| {
             manual_rollback::execute_rollback(
                 datastore,

crates/trident/src/logging/operation_context.rs

@@ -83,6 +83,27 @@ impl OperationSource {
     }
 }
 
+/// Whether the current invocation is allowed to generate/attach a
+/// persistent, datastore-backed ID (`installation_id`, `servicing_id`) --
+/// as opposed to merely reading one back for telemetry, which every
+/// source may do.
+///
+/// Only [`OperationSource::Cli`] and [`OperationSource::Daemon`] ever
+/// actually execute `Trident::install`/`update`/`rollback`/`commit`
+/// against a real datastore; [`OperationSource::GrpcClient`] only relays
+/// the request to a daemon, which is what performs the real work (tagged
+/// `Daemon`). There is therefore no real call site under `GrpcClient`
+/// today that would try to generate one of these IDs -- but that's an
+/// implicit property of today's call graph, not something the type system
+/// enforces. This gate makes the assumption explicit and centrally
+/// testable, rather than depending on every future call site
+/// independently getting it right (an equivalent implicit assumption is
+/// what caused `installation_id`'s multiboot mis-attach bug: see
+/// `Trident::new_deferring_installation_id`).
+pub fn should_generate_persistent_ids(source: OperationSource) -> bool {
+    matches!(source, OperationSource::Cli | OperationSource::Daemon)
+}
+
 thread_local! {
     static CURRENT_OPERATION: RefCell<Option<(String, CommandKind, OperationSource)>> =
         const { RefCell::new(None) };
@@ -260,6 +281,16 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_should_generate_persistent_ids() {
+        assert!(should_generate_persistent_ids(OperationSource::Cli));
+        assert!(should_generate_persistent_ids(OperationSource::Daemon));
+        assert!(
+            !should_generate_persistent_ids(OperationSource::GrpcClient),
+            "grpc-client only relays to a daemon; it must never generate installation_id/servicing_id itself"
+        );
+    }
+
     #[test]
     fn test_run_with_operation_clears_context_on_panic() {
         assert!(current().is_none());