diff --git a/crates/trident/src/datastore.rs b/crates/trident/src/datastore.rs index df811bed1a..fdb2b3be4f 100644 --- a/crates/trident/src/datastore.rs +++ b/crates/trident/src/datastore.rs @@ -3,15 +3,24 @@ use std::{fs, path::Path}; use log::{debug, warn}; use serde::{de::DeserializeOwned, Serialize}; use sqlite::State; +use uuid::Uuid; use trident_api::{ error::{ DatastoreError, InternalError, ReportError, ServicingError, TridentError, TridentResultExt, }, - status::{decode_host_status, HostStatus, TridentVersion}, + status::{decode_host_status, HostStatus, ServicingState, TridentVersion}, }; -use crate::TRIDENT_SEMVER_VERSION; +use crate::{logging::operation_context, TRIDENT_SEMVER_VERSION}; + +/// Key under which the shared installation ID is stored in the generic +/// key-value table. This ID is generated once, at the start of +/// `Trident::install` (see `DataStore::ensure_installation_id`), and +/// persisted for the lifetime of the datastore. It is intended to be added +/// to tracing/telemetry so that all activity for a given host installation +/// can be correlated. +const INSTALLATION_ID_KEY: &str = "installation-id"; pub struct DataStore { db: Option, @@ -206,8 +215,8 @@ impl DataStore { pub(crate) fn persist(&mut self, path: &Path) -> Result<(), TridentError> { if self.temporary { - // Read every generic key-value entry (e.g. an installation ID) - // out of the temporary datastore before starting any + // Read every generic key-value entry (e.g. the installation + // ID) out of the temporary datastore before starting any // destination writes below: this both avoids a same-file // read/write lock deadlock when `path` is the datastore's own // current path (see `write_key_values`'s doc comment) and lets @@ -453,11 +462,6 @@ impl DataStore { /// Values are serialized as JSON, so any type implementing /// `serde::Serialize`/`serde::de::DeserializeOwned` can be stored, not /// just `HostStatus`. - /// - /// No first-party caller exists yet in this crate on its own (see - /// `set_value`/`set_value_if_absent` for why); kept as part of the - /// generic key-value API for the first real consumer to build on. - #[allow(dead_code)] pub(crate) fn get_value( &self, key: &str, @@ -531,14 +535,9 @@ impl DataStore { /// Like [`Self::set_value`], but only inserts a row if `key` does not /// already have one; an existing row is left untouched. Used where two /// datastore connections could race to perform "first access" - /// initialization of a key: whichever connection's insert commits first - /// wins, and the other's insert becomes a no-op instead of overwriting - /// the winner's value. - /// - /// No first-party caller exists yet in this crate on its own -- this is - /// exercised only by tests for now. Kept as part of the generic - /// key-value API for the first real get-or-create-style consumer. - #[allow(dead_code)] + /// initialization of a key (see [`Self::ensure_installation_id`]): + /// whichever connection's insert commits first wins, and the other's + /// insert becomes a no-op instead of overwriting the winner's value. pub(crate) fn set_value_if_absent( &self, key: &str, @@ -599,6 +598,96 @@ impl DataStore { Ok(()) } + /// Returns this datastore's installation ID, if one has already been + /// persisted. Read-only: never generates or persists one -- callers + /// that need get-or-create semantics must call + /// [`Self::ensure_installation_id`] instead. + pub fn installation_id(&self) -> Result, TridentError> { + self.get_value::(INSTALLATION_ID_KEY) + } + + /// Retrieve this datastore's installation ID, generating and + /// persisting a new one on first access. + /// + /// Multiboot datastores inherit the source installation's ID instead. + /// The ID survives the temporary-to-persistent transition in `persist` and is + /// intended to be attached to tracing/telemetry so that activity for a + /// given host installation can be correlated across logs and traces. + /// + /// The value used on first access is this invocation's own + /// `operation_id` (see `logging::operation_context::current`), not an + /// unrelated freshly-generated UUID: whichever command actually + /// creates the datastore (normally `Trident::install`) has, by the + /// time it gets here, already been tagging its own telemetry with + /// `operation_id` since `command_start`. Reusing it here means that + /// invocation's telemetry never needs a second, disconnected + /// installation ID -- and every later command on this host just reads + /// back this same value. Falls back to a fresh random UUID only if + /// called with no operation context active (should not happen for any + /// real caller). + /// + /// First access is not a simple read-then-write: `Trident::new` may be + /// invoked concurrently (e.g. by multiple daemon RPC handlers), each + /// opening its own connection to the same datastore file. A naive + /// "read, and if absent generate + write" would let two connections + /// both observe no row, generate different UUIDs, and each overwrite + /// the other -- leaving one caller tracing with an ID that was never + /// actually persisted. Instead, unconditionally attempt to insert a + /// freshly generated ID with `ON CONFLICT DO NOTHING` (a no-op if + /// another connection already inserted one first), then read back + /// whichever ID actually won that race. + pub fn ensure_installation_id(&mut self) -> Result { + if let Some(id) = self.get_value::(INSTALLATION_ID_KEY)? { + return Ok(id); + } + + let new_id = operation_context::current() + .and_then(|(operation_id, _, _)| Uuid::parse_str(&operation_id).ok()) + .unwrap_or_else(Uuid::new_v4); + self.set_value_if_absent(INSTALLATION_ID_KEY, &new_id)?; + + self.get_value::(INSTALLATION_ID_KEY)? + .structured(InternalError::Internal( + "Installation ID missing immediately after being inserted", + )) + } + + /// Persist the source installation's ID in a multiboot staging datastore. + /// + /// The source datastore is authoritative, even when the reused staging + /// datastore contains an older ID. Do not reset its servicing state or + /// servicing ID when adopting the shared installation identity. + pub fn inherit_installation_id(&mut self, installation_id: Uuid) -> Result<(), TridentError> { + self.set_value(INSTALLATION_ID_KEY, &installation_id) + } + + /// Returns this datastore's installation ID, creating one as a + /// one-time migration if it is missing. Single source of truth for + /// every caller that attaches an installation ID to the shared + /// `TraceStream` on datastore open, so the same migration logic isn't + /// hand-copied (and drifting) across `Trident::new`, the CLI's + /// installation-ID pre-warm, and the daemon's per-request backstop. + /// + /// A datastore can be genuinely provisioned (`servicing_state != + /// NotProvisioned`) without ever having gone through + /// `Trident::install` -- offline initialization and the CIH update + /// bootstrap both create/adopt a datastore directly. Those hosts + /// would otherwise never get an installation ID, since nothing else + /// ever calls `ensure_installation_id`. Create one as a one-time + /// migration in that specific case, while leaving a genuinely + /// unprovisioned (temporary or not-yet-installed) datastore alone -- + /// callers should not observe an installation ID appear before the + /// host actually completes an install. + pub fn installation_id_or_migrate(&mut self) -> Result, TridentError> { + match self.installation_id()? { + Some(id) => Ok(Some(id)), + None if self.host_status().servicing_state != ServicingState::NotProvisioned => { + self.ensure_installation_id().map(Some) + } + None => Ok(None), + } + } + /// 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. @@ -704,17 +793,15 @@ mod tests { let datastore_path = temp_dir.path().join("db.sqlite"); let mut datastore = super::DataStore::open_or_create(&datastore_path).unwrap(); - datastore - .set_value("test-key", &"test-value".to_string()) - .unwrap(); + let installation_id = datastore.ensure_installation_id().unwrap(); // Persist to the exact same path the datastore is currently open at. datastore.persist(&datastore_path).unwrap(); assert_eq!( - datastore.get_value::("test-key").unwrap(), - Some("test-value".to_string()), - "keyvalue row should survive a self-persist" + datastore.ensure_installation_id().unwrap(), + installation_id, + "Installation ID should survive a self-persist" ); } @@ -964,8 +1051,8 @@ mod tests { #[test] /// Regression test: a datastore created by an older Trident version that /// predates the `keyvalue` table (only `hoststatus` exists) must still - /// be usable after `open()` -- in particular, `get_value`/`set_value` - /// must not fail with "no such table: keyvalue". + /// be usable after `open()` -- in particular, `ensure_installation_id()` must + /// not fail with "no such table: keyvalue". fn test_open_upgrades_pre_existing_datastore_schema() { let temp_dir = tempfile::tempdir().unwrap(); let path = temp_dir.path().join("db.sqlite"); @@ -984,11 +1071,9 @@ mod tests { .unwrap(); } - let datastore = super::DataStore::open(&path).unwrap(); + let mut datastore = super::DataStore::open(&path).unwrap(); // Should not fail with "no such table: keyvalue". - datastore - .set_value("test-key", &"test-value".to_string()) - .unwrap(); + datastore.ensure_installation_id().unwrap(); temp_dir.close().unwrap(); } @@ -1014,33 +1099,15 @@ mod tests { temp_dir.close().unwrap(); } - /// Test helper mirroring the "get-or-create" pattern a real consumer - /// (e.g. a future `installation_id`-style accessor) would build on top - /// of `get_value`/`set_value_if_absent`: read the value back if - /// present, otherwise attempt to insert a freshly generated one with - /// insert-if-absent semantics, then read back whichever value actually - /// won. - fn get_or_create_test_value(datastore: &mut super::DataStore) -> String { - const KEY: &str = "test-key"; - if let Some(v) = datastore.get_value::(KEY).unwrap() { - return v; - } - datastore - .set_value_if_absent(KEY, &uuid::Uuid::new_v4().to_string()) - .unwrap(); - datastore.get_value::(KEY).unwrap().unwrap() - } - #[test] - fn test_set_value_if_absent_concurrent_first_access_is_consistent() { + fn test_installation_id_concurrent_first_access_is_consistent() { let temp_dir = tempfile::tempdir().unwrap(); let path = temp_dir.path().join("db.sqlite"); // Create the datastore (and its schema) up front, then open two - // separate connections to it, simulating two callers concurrently - // performing "first access" get-or-create initialization (e.g. two - // daemon RPC handlers both calling `Trident::new`) against the same - // datastore path. + // separate connections to it, simulating two daemon RPC handlers + // concurrently calling `Trident::new` (and therefore + // `installation_id`) against the same datastore path. super::DataStore::make_datastore(&path).unwrap(); let barrier = std::sync::Arc::new(std::sync::Barrier::new(2)); @@ -1051,25 +1118,111 @@ mod tests { std::thread::spawn(move || { let mut datastore = super::DataStore::open(&path).unwrap(); // Synchronize so both threads attempt "first access" - // (no value persisted yet) as close together as - // possible. + // (no installation ID persisted yet) as close together + // as possible. barrier.wait(); - get_or_create_test_value(&mut datastore) + datastore.ensure_installation_id().unwrap() }) }) .collect(); - let values: Vec = handles.into_iter().map(|h| h.join().unwrap()).collect(); + let ids: Vec = handles.into_iter().map(|h| h.join().unwrap()).collect(); assert_eq!( - values[0], values[1], - "concurrent first access returned inconsistent values" + ids[0], ids[1], + "concurrent first access returned inconsistent installation IDs" ); temp_dir.close().unwrap(); } #[test] - fn test_set_value_if_absent_is_stable() { + fn test_multiboot_installation_id_survives_reopen_and_persist() { + for existing_source_id in [false, true] { + for existing_staging_id in [false, true] { + let directory = tempfile::tempdir().unwrap(); + let source_path = directory.path().join("source.sqlite"); + let staging_path = directory.path().join("staging.sqlite"); + let target_path = directory.path().join("target.sqlite"); + let mut source = super::DataStore::open_or_create(&source_path).unwrap(); + if existing_source_id { + source.ensure_installation_id().unwrap(); + } + source + .with_host_status(|status| { + status.servicing_state = super::ServicingState::Provisioned; + }) + .unwrap(); + source.persist(&source_path).unwrap(); + + let installation_id = source.ensure_installation_id().unwrap(); + source.close(); + + let mut staging = super::DataStore::open_or_create(&staging_path).unwrap(); + if existing_staging_id { + let stale_id = staging.ensure_installation_id().unwrap(); + assert_ne!(stale_id, installation_id); + } + staging.set_value("staged-image", &"image-hash").unwrap(); + staging + .with_host_status(|status| { + status.servicing_state = super::ServicingState::CleanInstallStaged; + }) + .unwrap(); + let host_status = serde_yaml::to_string(staging.host_status()).unwrap(); + + staging.inherit_installation_id(installation_id).unwrap(); + assert_eq!( + serde_yaml::to_string(staging.host_status()).unwrap(), + host_status + ); + drop(staging); + + 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.get_value::("staged-image").unwrap(), + Some("image-hash".to_string()) + ); + assert_eq!( + staging.host_status().servicing_state, + super::ServicingState::CleanInstallStaged + ); + staging.persist(&target_path).unwrap(); + drop(staging); + + 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)); + } + } + } + } + + #[test] + fn test_installation_id_is_stable() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().join("db.sqlite"); + let db = super::DataStore::make_datastore(&path).unwrap(); + let mut datastore = super::DataStore { + db: Some(db), + host_status: Default::default(), + temporary: false, + }; + + let id = datastore.ensure_installation_id().unwrap(); + // Calling installation_id again should return the same ID, not generate a + // new one. + assert_eq!(datastore.ensure_installation_id().unwrap(), id); + + temp_dir.close().unwrap(); + } + + #[test] + /// The read-only getter must not create an installation ID -- only + /// `ensure_installation_id` (called specifically at the start of + /// `Trident::install`) does that. + fn test_installation_id_read_only_getter_does_not_create() { let temp_dir = tempfile::tempdir().unwrap(); let path = temp_dir.path().join("db.sqlite"); let db = super::DataStore::make_datastore(&path).unwrap(); @@ -1079,10 +1232,19 @@ mod tests { temporary: false, }; - let value = get_or_create_test_value(&mut datastore); - // Calling get_or_create again should return the same value, not - // generate a new one. - assert_eq!(get_or_create_test_value(&mut datastore), value); + assert_eq!( + datastore.installation_id().unwrap(), + None, + "no installation ID should exist before ensure_installation_id is called" + ); + + let created = datastore.ensure_installation_id().unwrap(); + + assert_eq!( + datastore.installation_id().unwrap(), + Some(created), + "the read-only getter should now see the created ID" + ); temp_dir.close().unwrap(); } @@ -1163,28 +1325,22 @@ mod functional_test { } #[functional_test] - fn test_generic_key_value_survives_persist() { + fn test_installation_id_survives_persist() { let temp_dir = TempDir::new().unwrap(); let datastore_temp_path = temp_dir.path().join("db-tmp.sqlite"); let datastore_path = temp_dir.path().join("db.sqlite"); - // Store a generic keyvalue entry in the temporary datastore, then - // persist it. - { + // Generate a installation ID in the temporary datastore, then persist it. + let installation_id = { let mut datastore = DataStore::open_or_create(&datastore_temp_path).unwrap(); - datastore - .set_value("test-key", &"test-value".to_string()) - .unwrap(); + let installation_id = datastore.ensure_installation_id().unwrap(); datastore.persist(&datastore_path).unwrap(); + installation_id }; - // Re-open the persisted datastore and verify the same value is - // returned, rather than being lost in the temporary-to-persistent - // transition. - let datastore = DataStore::open(&datastore_path).unwrap(); - assert_eq!( - datastore.get_value::("test-key").unwrap(), - Some("test-value".to_string()) - ); + // Re-open the persisted datastore and verify the same installation ID is + // returned, rather than a new one being generated. + let mut datastore = DataStore::open(&datastore_path).unwrap(); + assert_eq!(datastore.ensure_installation_id().unwrap(), installation_id); } } diff --git a/crates/trident/src/lib.rs b/crates/trident/src/lib.rs index b678286a41..de03e0c88b 100644 --- a/crates/trident/src/lib.rs +++ b/crates/trident/src/lib.rs @@ -135,6 +135,10 @@ pub struct Trident { host_config: Option, orchestrator: Option, is_stream_image: bool, + /// Kept so `Trident::install` can attach a newly-created installation + /// ID to it -- see `DataStore::ensure_installation_id`, only ever + /// called from `install`/staging, unlike this constructor. + tracestream: TraceStream, } impl Trident { @@ -143,6 +147,29 @@ impl Trident { datastore_path: &Path, logstream: Logstream, tracestream: TraceStream, + ) -> Result { + Self::new_impl(config_source, datastore_path, logstream, tracestream, true) + } + + /// 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. + pub fn new_deferring_installation_id( + config_source: Option, + datastore_path: &Path, + logstream: Logstream, + tracestream: TraceStream, + ) -> Result { + Self::new_impl(config_source, datastore_path, logstream, tracestream, false) + } + + fn new_impl( + config_source: Option, + datastore_path: &Path, + logstream: Logstream, + tracestream: TraceStream, + attach_installation_id: bool, ) -> Result { let host_config = config_source .map(|source| Self::load_host_config(&source)) @@ -218,6 +245,10 @@ 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. + tracestream.attach_ids_if_present(datastore_path, !attach_installation_id); + // Trace features enabled in the Host Configuration. if let Some(hc) = &host_config { hc.feature_tracing(); @@ -244,6 +275,7 @@ impl Trident { host_config, orchestrator, is_stream_image: false, + tracestream, }) } @@ -453,6 +485,7 @@ impl Trident { ))?; let is_stream_image = self.is_stream_image; + let tracestream = self.tracestream.clone(); self.execute_and_record_error(datastore, |datastore| { host_config @@ -474,6 +507,7 @@ impl Trident { // provisioned host. If the host is not provisioned, we can proceed // with a clean install. If the host is provisioned, we need to // check if a multiboot install was requested. + let mut inherited_installation_id = None; if datastore.is_persistent() { if !multiboot { // If the host IS provisioned and multiboot is NOT requested @@ -491,12 +525,28 @@ impl Trident { multiboot install" ); + inherited_installation_id = Some(datastore.ensure_installation_id()); datastore.close(); *datastore = DataStore::open_or_create(Path::new(TEMPORARY_DATASTORE_PATH)) .message("Failed to create temporary datastore for multiboot install")?; } } + // Persist inheritance independently of the telemetry cache. A reused + // multiboot datastore may still contain another installation's ID. + let installation_id_result = match inherited_installation_id { + Some(installation_id) => installation_id.and_then(|id| { + datastore.inherit_installation_id(id)?; + tracestream.set_installation_id(id.to_string()); + Ok(()) + }), + None => tracestream.ensure_and_attach_installation_id(datastore), + }; + if let Err(e) = installation_id_result { + // Do not mint a different ID when multiboot inheritance fails. + warn!("Failed to attach installation ID: {e:?}"); + } + // Use a prefetched image if provided, otherwise load the image // specified in the Host Configuration. let image = match prefetched_image { @@ -593,6 +643,8 @@ impl Trident { "update called without Host Configuration set", ))?; + let tracestream = self.tracestream.clone(); + self.execute_and_record_error(datastore, |datastore| { // Ensure that the datastore exists. if !datastore.is_persistent() { @@ -635,6 +687,20 @@ impl Trident { status.is_management_os = false; }) .message("Failed to initialize datastore")?; + + // This host has just adopted a datastore via the CIH + // bootstrap path, bypassing `Trident::install` entirely + // -- nothing else will ever create an installation ID for + // it otherwise. Best-effort: a failure here must not + // block the update itself, since telemetry attribution + // is not load-bearing for servicing outcomes. Known + // remaining gap: this invocation's own + // command_start/trident_start (fired even earlier, in + // the CLI/daemon dispatch and Trident::new respectively) + // still won't carry it -- both fire before this point. + if let Err(e) = tracestream.ensure_and_attach_installation_id(datastore) { + warn!("Failed to create installation ID during CIH bootstrap: {e:?}"); + } } else { // For non-ACL images, if the datastore is not persistent, return error return Err(TridentError::new(InvalidInputError::HostNotProvisioned)) diff --git a/crates/trident/src/logging/tracestream.rs b/crates/trident/src/logging/tracestream.rs index a108818dca..edafb531de 100644 --- a/crates/trident/src/logging/tracestream.rs +++ b/crates/trident/src/logging/tracestream.rs @@ -3,6 +3,7 @@ use std::{ collections::BTreeMap, fs::File, io::Write, + path::Path, sync::{Arc, RwLock}, time::Instant, }; @@ -19,13 +20,20 @@ use tracing::{ }; use tracing_subscriber::{layer::Layer, registry::LookupSpan}; -use crate::{logging::operation_context, TRIDENT_METRICS_FILE_PATH, TRIDENT_VERSION}; +#[cfg(test)] +use crate::command_kind::CommandKind; +use trident_api::error::TridentError; + use osutils::{ dmi, files, osrelease::{OsRelease, OS_RELEASE_PATH}, uname, }; +use crate::{ + datastore::DataStore, logging::operation_context, TRIDENT_METRICS_FILE_PATH, TRIDENT_VERSION, +}; + lazy_static::lazy_static! { static ref ADDITIONAL_FIELDS: BTreeMap = populate_additional_fields(); pub static ref PLATFORM_INFO: BTreeMap = populate_platform_info(); @@ -82,6 +90,7 @@ pub struct TraceStream { // TODO: Consider changing this to a LockOnce when rustc is updated to // >=1.70 target: Arc>>, + installation_id: Arc>>, disabled: bool, } @@ -123,6 +132,92 @@ impl TraceStream { Ok(()) } + /// Set the installation ID to attach to every trace entry sent from this + /// point forward, as an additional field, so that all traces/metrics for + /// a given host installation can be correlated. Expected to be called + /// once the datastore's persisted installation ID has been retrieved + /// (see [`Self::attach_ids_if_present`] and + /// [`Self::ensure_and_attach_installation_id`]). + pub fn set_installation_id(&self, installation_id: String) { + match self.installation_id.write() { + Ok(mut val) => { + val.replace(installation_id); + } + Err(_) => warn!("Failed to lock tracestream to set installation ID"), + } + } + + /// Returns a clone of the shared installation-ID handle -- the same + /// underlying `Arc>` written by `set_installation_id` -- so + /// other telemetry sinks (namely `AppInsightsSender`) can read the + /// current value at send-time without needing their own copy of the + /// logic that sets it. + pub fn installation_id_handle(&self) -> Arc>> { + self.installation_id.clone() + } + + fn installation_id_cached(&self) -> bool { + self.installation_id + .read() + .map(|v| v.is_some()) + .unwrap_or(false) + } + + /// Attach an existing installation ID, migrating provisioned legacy + /// datastores when needed. Never creates a datastore. Multiboot + /// callers may defer attachment until their datastore is selected. + pub fn attach_ids_if_present(&self, datastore_path: &Path, skip_installation_id: bool) { + if !datastore_path.exists() { + return; + } + let mut datastore = match DataStore::open(datastore_path) { + Ok(ds) => ds, + Err(e) => { + warn!("Failed to open datastore to attach IDs: {e:?}"); + return; + } + }; + + if !skip_installation_id && !self.installation_id_cached() { + match datastore.installation_id_or_migrate() { + Ok(Some(installation_id)) => { + info!("Installation ID: {installation_id}"); + self.set_installation_id(installation_id.to_string()); + } + Ok(None) => { + debug!("No installation ID persisted yet (host not yet installed)"); + } + Err(e) => { + warn!("Failed to read installation ID: {e:?}"); + } + } + } + } + + /// Ensures `datastore` has an installation ID -- creating one if this + /// is the first access, or reading back the existing one otherwise -- + /// and attaches it. Unlike + /// [`Self::attach_ids_if_present`], this is only for the + /// one caller that already knows a command genuinely allowed to + /// initialize a brand-new datastore (per + /// [`crate::datastore::DataStore::may_initialize_datastore_for_command`]) + /// is proceeding, and already holds (or just created) the datastore + /// handle for it -- so this attaches the new install/update's own ID + /// instead of leaving the trace stream untagged until some later + /// read-only attach happens to run. + pub fn ensure_and_attach_installation_id( + &self, + datastore: &mut DataStore, + ) -> Result<(), TridentError> { + if self.installation_id_cached() { + return Ok(()); + } + let installation_id = datastore.ensure_installation_id()?; + info!("Installation ID: {installation_id}"); + self.set_installation_id(installation_id.to_string()); + Ok(()) + } + /// Create a Boxed TraceSender pub fn make_trace_sender(&self) -> Box { self.make_trace_sender_with_metrics_path(TRIDENT_METRICS_FILE_PATH) @@ -135,12 +230,17 @@ impl TraceStream { /// real, shared host path, so they can be plain `#[test]`s instead of /// needing a VM. fn make_trace_sender_with_metrics_path(&self, metrics_file_path: &str) -> Box { - Box::new(TraceSender::new(self.target.clone(), metrics_file_path)) + Box::new(TraceSender::new( + self.target.clone(), + self.installation_id.clone(), + metrics_file_path, + )) } } pub struct TraceSender { server: Arc>>, + installation_id: Arc>>, client: reqwest::blocking::Client, metrics_file: Option, } @@ -152,9 +252,14 @@ struct ExecutionTime(Instant); /// the tracing-subscriber crate to handle the events and send them to the /// server. impl TraceSender { - fn new(server: Arc>>, metrics_file_path: &str) -> Self { + fn new( + server: Arc>>, + installation_id: Arc>>, + metrics_file_path: &str, + ) -> Self { Self { server, + installation_id, client: reqwest::blocking::Client::new(), metrics_file: match files::create_file(metrics_file_path) { Ok(f) => Some(f), @@ -174,6 +279,11 @@ impl TraceSender { fn additional_fields(&self) -> BTreeMap { let mut fields = ADDITIONAL_FIELDS.clone(); + if let Ok(installation_id) = self.installation_id.read() { + if let Some(installation_id) = installation_id.as_ref() { + fields.insert(fields::INSTALLATION_ID.to_string(), json!(installation_id)); + } + } merge_operation_context(&mut fields); fields } @@ -391,7 +501,7 @@ fn product_uuid_or_unknown(uuid: Result) -> String { /// /// Shared by both local telemetry sinks (`TraceSender::additional_fields` /// below and `AppInsightsSender::send_event`) so the -/// operation_id/command/source/installation_id-fallback rule has one +/// operation_id/command/source enrichment rule has one /// implementation instead of being hand-duplicated between them -- a /// prior version of this function existed independently in each sink, /// which risked the two silently diverging if the rule ever changed in @@ -474,6 +584,24 @@ mod tests { use tracing_subscriber::{filter, layer::SubscriberExt}; + #[test] + fn test_operation_context_does_not_invent_an_installation_id() { + operation_context::run_with_operation( + &CommandKind::update(), + operation_context::OperationSource::Cli, + || { + let mut fields = BTreeMap::new(); + merge_operation_context(&mut fields); + assert!(fields.contains_key("operation_id")); + assert!(!fields.contains_key("installation_id")); + + fields.insert("installation_id".to_string(), json!("persisted-id")); + merge_operation_context(&mut fields); + assert_eq!(fields["installation_id"], json!("persisted-id")); + }, + ); + } + #[test] fn test_tracestream() { let temp_dir = tempfile::tempdir().unwrap(); @@ -580,6 +708,50 @@ mod tests { ); } + #[test] + /// Regression test: `TraceStream::set_installation_id` must actually + /// reach the serialized trace entry's `additional_fields.installation_id` + /// -- the metric/span tests above only assert on `metric_name`/`value` + /// and would still pass even if the installation ID were never copied + /// into `additional_fields`. + fn test_tracestream_installation_id_written_to_additional_fields() { + let temp_dir = tempfile::tempdir().unwrap(); + let metrics_path = temp_dir.path().join("metrics.jsonl"); + let tracestream = TraceStream::default(); + tracestream.set_installation_id("test-installation-id".to_string()); + let trace_sender = tracestream + .make_trace_sender_with_metrics_path(metrics_path.to_str().unwrap()) + .with_filter(filter::LevelFilter::INFO); + + // See test_tracestream_write_metric_event_to_file for why a scoped + // (not global) default subscriber is used here. + let _guard = tracing::subscriber::set_default( + tracing_subscriber::Registry::default().with(trace_sender), + ); + + tracing::info!( + metric_name = "test_metric_with_installation_id", + value = true + ); + + // Ensure the trace system has time to write the file. + std::thread::sleep(std::time::Duration::from_millis(100)); + + let file = File::open(&metrics_path).unwrap(); + let reader = BufReader::new(file); + let lines: Vec = reader.lines().map(|l| l.unwrap()).collect(); + + let metric_found = lines.iter().any(|line| { + line.contains(r#""metric_name":"test_metric_with_installation_id""#) + && line.contains(r#""installation_id":"test-installation-id""#) + }); + + assert!( + metric_found, + "Expected metric with installation_id field not found in the local metrics file" + ); + } + #[test] fn test_tracestream_write_span_metric_to_file() { let temp_dir = tempfile::tempdir().unwrap(); @@ -619,6 +791,84 @@ mod tests { // Helper function to test span metrics #[tracing::instrument(name = "test_span", skip_all)] fn simulate_function_span() {} + + #[test] + /// Same cross-invocation join gap as + /// `test_servicing_id_stage_then_read_back_across_separate_tracestreams`, + /// but for `installation_id`: every existing installation_id test + /// either only exercises the datastore layer (`datastore.rs`) or sets + /// the value directly on a single `TraceStream` + /// (`test_tracestream_installation_id_written_to_additional_fields`) -- + /// none of them prove that a *separate* `TraceStream` calling + /// `attach_ids_if_present` actually recovers an installation ID + /// created by a different one. + fn test_installation_id_stage_then_read_back_across_separate_tracestreams() { + let temp_dir = tempfile::tempdir().unwrap(); + let db_path = temp_dir.path().join("db.sqlite"); + + // Stage: a fresh TraceStream + datastore, as if this were + // `Trident::install` creating the installation ID for the first + // time. + let mut staging_datastore = DataStore::open_or_create(&db_path).unwrap(); + let staging_tracestream = TraceStream::default(); + operation_context::run_with_operation( + &CommandKind::install(), + operation_context::OperationSource::Cli, + || { + staging_tracestream + .ensure_and_attach_installation_id(&mut staging_datastore) + .unwrap(); + }, + ); + let staged_id = staging_tracestream + .installation_id_handle() + .read() + .unwrap() + .clone() + .expect("staging call should have created an installation ID"); + + // Read back: a completely separate TraceStream, reading only from + // disk (via the datastore path, matching the real + // attach_ids_if_present call site signature), as if this were a + // later command or daemon request in the same process + // reattaching to an already-provisioned host. + let readback_tracestream = TraceStream::default(); + readback_tracestream.attach_ids_if_present(&db_path, false); + + assert_eq!( + readback_tracestream + .installation_id_handle() + .read() + .unwrap() + .clone(), + Some(staged_id), + "a separate TraceStream reading the same datastore should recover the installation ID" + ); + } + + #[test] + /// Negative counterpart: a `TraceStream` that only ever reads (never + /// creates) against a datastore that has never had an installation ID + /// created, and is not otherwise provisioned, must leave the handle + /// `None` rather than fabricating one. + fn test_installation_id_absent_when_never_created() { + let temp_dir = tempfile::tempdir().unwrap(); + let db_path = temp_dir.path().join("db.sqlite"); + + // Create the datastore file itself (so attach_ids_if_present's + // `datastore_path.exists()` check passes), but never create an + // installation ID on it. + DataStore::open_or_create(&db_path).unwrap(); + + let tracestream = TraceStream::default(); + tracestream.attach_ids_if_present(&db_path, false); + + assert_eq!( + tracestream.installation_id_handle().read().unwrap().clone(), + None, + "installation ID must stay absent when none has ever been created on this datastore" + ); + } } #[cfg(feature = "functional-test")] diff --git a/crates/trident/src/main.rs b/crates/trident/src/main.rs index 0eb730dbe4..41ccace3f9 100644 --- a/crates/trident/src/main.rs +++ b/crates/trident/src/main.rs @@ -166,7 +166,60 @@ fn run_trident( } } + // 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, + .. + } + ); + + // Attach this host's installation ID and current servicing + // ID to the shared TraceStream before run_with_operation + // below fires command_start: Trident::new (further down, + // inside the closure) is the usual place installation ID + // gets attached, but that's too late for command_start, + // which run_with_operation fires immediately, before the + // closure even runs. Does not create a *datastore* (see + // `TraceStream::attach_ids_if_present`) -- silently does + // nothing if the datastore doesn't exist yet, which is + // expected for a host's first-ever install. But not purely + // read-only: on an existing datastore that predates + // `installation_id`, this can perform a one-time migration + // write to mint one (see + // `DataStore::installation_id_or_migrate`). Skips the + // installation-ID half specifically (but still attaches + // servicing ID, unaffected by a multiboot swap) when + // `defer_installation_id` is set, for the same reason + // `Trident::new_deferring_installation_id` is used below. + // + // 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, + ); + } run_with_operation(&kind, OperationSource::Cli, || { // config_path was already validated (existence-checked) @@ -185,14 +238,29 @@ fn run_trident( .message("Datastore file does not exist"); } - let mut trident = Trident::new( - config_path.map(HostConfigurationSource::File), - agent_config.datastore_path(), - logstream.clone(), - tracestream.clone(), - ) + let mut trident = if defer_installation_id { + Trident::new_deferring_installation_id( + config_path.map(HostConfigurationSource::File), + agent_config.datastore_path(), + logstream.clone(), + tracestream.clone(), + ) + } else { + Trident::new( + config_path.map(HostConfigurationSource::File), + agent_config.datastore_path(), + logstream.clone(), + tracestream.clone(), + ) + } .message("Failed to initialize Trident")?; + // `Trident::new` (or `Trident::new_deferring_installation_id` + // for a multiboot install) has already attached this + // host's persisted installation ID -- if any -- to the + // shared TraceStream, so every trace/metric emitted + // from here on -- including "trident_start" -- carries + // it once available. let mut datastore = DataStore::open_or_create(agent_config.datastore_path()) .message("Failed to open datastore")?; diff --git a/crates/trident/src/server/mod.rs b/crates/trident/src/server/mod.rs index 4a94e8dadc..f8a0f723ba 100644 --- a/crates/trident/src/server/mod.rs +++ b/crates/trident/src/server/mod.rs @@ -107,6 +107,27 @@ pub fn server_main( } }; + // Attach this host's installation ID and current servicing ID to the + // shared TraceStream before accepting any RPCs. `Trident::new()` + // attaches installation ID the same way on every request, but the + // very first servicing request this daemon process ever handles would + // otherwise have its command_start (fired by run_with_operation before + // that request's own Trident::new() call runs) go out untagged, so + // attach it up front instead. Every later request is unaffected either + // way, since the shared TraceStream keeps whatever was set here (or by + // the first request) for the rest of the daemon's lifetime. Does not + // create a datastore (see `TraceStream::attach_ids_if_present`), but + // is not purely read-only: starting the daemon can mutate a legacy/ + // offline-provisioned datastore that predates `installation_id`, + // performing a one-time migration write to mint one (see + // `DataStore::installation_id_or_migrate`) before accepting RPCs. + // Never skips the installation-ID half here (unlike the CLI's + // pre-warm): the daemon never receives a multiboot install request -- + // every `install.rs` service handler call site hardcodes + // `multiboot: false` -- so there's no swap-datastore scenario to guard + // against on this path. + tracestream.attach_ids_if_present(agent_config.datastore_path(), false); + let shutdown_signals = match ShutdownSignals::setup_signal_handlers() { Ok(signals) => signals, Err(e) => { diff --git a/crates/trident/src/server/tridentserver/mod.rs b/crates/trident/src/server/tridentserver/mod.rs index 5c8e651e7a..60b705e602 100644 --- a/crates/trident/src/server/tridentserver/mod.rs +++ b/crates/trident/src/server/tridentserver/mod.rs @@ -4,7 +4,7 @@ //! services. It defines the `TridentServer` struct which implements the gRPC //! service interface generated by Tonic from the protobuf definitions. -use std::{panic, sync::Arc, time::SystemTime}; +use std::{panic, path::Path, sync::Arc, time::SystemTime}; use log::{error, info, warn}; use prost_types::Timestamp; @@ -178,6 +178,31 @@ impl TridentServer { }) } + /// Re-attaches a persisted installation ID and current servicing ID + /// to `tracestream`, if either is now available but wasn't at + /// daemon-startup time (`server_main`'s one-time attach runs before + /// any request has had a chance to create a datastore, so a request + /// that arrives before the very first install/update -- and whose + /// own handler goes on to create that datastore -- would otherwise + /// still be missing them). Takes an explicit `tracestream`/ + /// `datastore_path` (rather than `&self`) so it can be called from + /// inside the `'static` closure `servicing_request` hands to + /// `spawn_servicing_task` -- see the call site there for why this + /// runs inside `spawn_blocking` instead of synchronously on the + /// gRPC handler's async Tokio worker thread. Does not create a + /// datastore: silently does nothing if the datastore doesn't exist + /// yet. But on an existing datastore, this may still *persist* a + /// missing installation ID, via + /// `DataStore::installation_id_or_migrate`'s legacy-ID migration (see + /// `TraceStream::attach_ids_if_present`); the servicing-ID half is + /// fully read-only. Never skips the installation-ID half (unlike the + /// CLI's pre-warm) -- the daemon never receives a multiboot install + /// request, so there's no swap-datastore scenario to guard against + /// here. + fn refresh_ids(tracestream: &TraceStream, datastore_path: &Path) { + tracestream.attach_ids_if_present(datastore_path, false); + } + /// Handles a servicing request by acquiring the necessary locks, /// setting up log forwarding, and spawning the provided servicing task. /// @@ -232,6 +257,29 @@ impl TridentServer { return Err(Status::failed_precondition("Host is not provisioned")); } + // Re-check for a persisted installation ID and current servicing + // ID before this request fires its own command_start (below, via + // run_with_operation). server_main's daemon-startup attach only + // ever runs once, at startup -- so a request that arrives before + // any datastore exists (e.g. this daemon's very first install) + // would otherwise never see one, even after that request's own + // handler goes on to create the datastore. + // + // Deferred into the `f` closure below (called right before + // run_with_operation, preserving "before command_start") rather + // than run synchronously here: `attach_ids_if_present` opens the + // datastore, and datastore connections use a 5-second busy + // timeout, so lock contention could otherwise stall this + // request's shared Tokio worker thread for that long before the + // real servicing work even reaches `spawn_blocking`. The closure + // itself runs entirely inside `spawn_servicing_task`'s + // `spawn_blocking` (see there), which has its own dedicated OS + // thread, so the same I/O there can't block anything else. + // Captures owned clones (not `self`) since the closure must be + // `'static`. + let tracestream_for_ids = self.tracestream.clone(); + let datastore_path_for_ids = self.agent_config.datastore_path().to_path_buf(); + // Tag every metric/tracing event `f` fires (on whatever thread it // ultimately runs on -- see `spawn_servicing_task`, which runs it // via `tokio::task::spawn_blocking`, giving it a dedicated OS @@ -248,6 +296,7 @@ impl TridentServer { // operation's identity instead of leaving it untagged. let kind_for_closure = kind.clone(); let f = move || { + Self::refresh_ids(&tracestream_for_ids, &datastore_path_for_ids); operation_context::run_with_operation( &kind_for_closure, operation_context::OperationSource::Daemon, diff --git a/crates/trident/src/telemetry.rs b/crates/trident/src/telemetry.rs index ee339173c7..0954482674 100644 --- a/crates/trident/src/telemetry.rs +++ b/crates/trident/src/telemetry.rs @@ -17,6 +17,7 @@ pub(crate) mod fields { pub const ASSET_ID: &str = "asset_id"; pub const COMMAND: &str = "command"; pub const EXECUTION_TIME: &str = "execution_time"; + pub const INSTALLATION_ID: &str = "installation_id"; pub const KERNEL_VERSION: &str = "kernel_version"; pub const METRIC_NAME: &str = "metric_name"; pub const OPERATION_ID: &str = "operation_id";