diff --git a/crates/trident/src/datastore.rs b/crates/trident/src/datastore.rs index 9136dba1f..712a6b9a0 100644 --- a/crates/trident/src/datastore.rs +++ b/crates/trident/src/datastore.rs @@ -1,7 +1,9 @@ use std::{fs, path::Path}; use log::{debug, warn}; +use serde::{de::DeserializeOwned, Serialize}; use sqlite::State; +use uuid::Uuid; use trident_api::{ error::{ @@ -12,6 +14,13 @@ use trident_api::{ use crate::TRIDENT_SEMVER_VERSION; +/// Key under which the datastore's unique database ID is stored in the +/// generic key-value table. This ID is generated once (on first access) 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 DATASTORE_ID_KEY: &str = "datastore-id"; + pub struct DataStore { db: Option, host_status: HostStatus, @@ -37,11 +46,38 @@ impl DataStore { pub(crate) fn open(path: &Path) -> Result { debug!("Loading datastore from {}", path.display()); - let db = sqlite::open(path).structured(ServicingError::Datastore { + let mut db = sqlite::open(path).structured(ServicingError::Datastore { inner: DatastoreError::LoadDatastore { path: path.to_string_lossy().into(), }, })?; + // Multiple connections to the same datastore file (e.g. concurrent + // daemon RPC handlers) can briefly contend for the write lock; wait + // for it rather than failing immediately with "database is locked". + Self::set_busy_timeout(&mut db)?; + // Require the `hoststatus` table to already be present before + // applying any schema migration. This distinguishes a genuinely + // pre-existing datastore (which may only be missing a table added + // in a later Trident version, e.g. `keyvalue`) from a zero-byte or + // otherwise corrupt/truncated file, which must fail loudly here + // rather than silently succeeding as an empty, freshly-provisioned + // datastore. + if !Self::table_exists(&db, "hoststatus")? { + return Err(TridentError::new(ServicingError::Datastore { + inner: DatastoreError::LoadDatastore { + path: path.to_string_lossy().into(), + }, + })) + .message( + "Existing datastore file is missing its 'hoststatus' table; \ + the file may be corrupt, truncated, or not a Trident datastore", + ); + } + // Existing datastores may predate a table added in a later Trident + // version (e.g. `keyvalue`). Idempotently ensure the full schema is + // present so upgraded hosts don't fail with "no such table" the + // first time a new table is accessed. + Self::ensure_schema(&db)?; let host_status_yaml: Option = db .prepare("SELECT contents FROM hoststatus ORDER BY id DESC LIMIT 1") .structured(ServicingError::Datastore { @@ -118,17 +154,62 @@ impl DataStore { DatastoreError::CreateDatastoreDirectory, ))?; - let db = + let mut db = sqlite::open(path).structured(ServicingError::from(DatastoreError::OpenDatastore))?; + Self::set_busy_timeout(&mut db)?; + Self::ensure_schema(&db)?; + Ok(db) + } + + /// Wait (rather than immediately failing with "database is locked") for + /// up to five seconds when another connection holds the write lock on + /// this datastore file. Needed because more than one connection to the + /// same datastore can legitimately exist at once (e.g. concurrent + /// daemon RPC handlers each calling `Trident::new`), and SQLite's + /// default busy timeout is zero. + fn set_busy_timeout(db: &mut sqlite::Connection) -> Result<(), TridentError> { + db.set_busy_timeout(5000) + .structured(ServicingError::from(DatastoreError::OpenDatastore)) + } + + /// Returns whether a table with the given name currently exists in the + /// datastore's schema (queried via `sqlite_master`). + fn table_exists(db: &sqlite::Connection, table: &str) -> Result { + let mut statement = db + .prepare("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?") + .structured(ServicingError::from(DatastoreError::InitializeDatastore))?; + statement + .bind((1, table)) + .structured(ServicingError::from(DatastoreError::InitializeDatastore))?; + Ok(matches!( + statement + .next() + .structured(ServicingError::from(DatastoreError::InitializeDatastore))?, + State::Row + )) + } + + /// Idempotently create any tables that don't already exist. Safe to call + /// on both newly-created and pre-existing datastores, so that a + /// datastore created by an older Trident version picks up tables added + /// by a newer version the next time it is opened. + fn ensure_schema(db: &sqlite::Connection) -> Result<(), TridentError> { db.execute( "CREATE TABLE IF NOT EXISTS hoststatus ( id INTEGER PRIMARY KEY AUTOINCREMENT, - timestamp DATETIME DEFALUT CURRENT_TIMESTAMP, + timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, contents TEXT NOT NULL )", ) .structured(ServicingError::from(DatastoreError::InitializeDatastore))?; - Ok(db) + db.execute( + "CREATE TABLE IF NOT EXISTS keyvalue ( + key TEXT PRIMARY KEY, + contents TEXT NOT NULL + )", + ) + .structured(ServicingError::from(DatastoreError::InitializeDatastore))?; + Ok(()) } pub(crate) fn persist(&mut self, path: &Path) -> Result<(), TridentError> { @@ -139,6 +220,14 @@ impl DataStore { TridentVersion::SemVer(TRIDENT_SEMVER_VERSION.clone()); Self::write_host_status(&persistent_db, self.host_status())?; + // Carry over any generic key-value entries (e.g. the database ID) + // recorded in the temporary datastore into the persistent one, so + // they survive the transition from temporary to persistent + // storage. + if let Some(temporary_db) = self.db.as_ref() { + Self::copy_key_values(temporary_db, &persistent_db)?; + } + self.db = Some(persistent_db); self.temporary = false; } @@ -146,6 +235,82 @@ impl DataStore { Ok(()) } + /// Copy all rows of the generic key-value table from `source` into + /// `destination`, overwriting any conflicting keys already present in + /// `destination`. + /// + /// `source` and `destination` may be two live connections to the *same* + /// underlying SQLite file (e.g. an offline `persist` whose destination + /// path is the currently-open datastore's own path). SQLite's locking is + /// per-connection, so a `SELECT` left active on `source` holds a read + /// lock that a write from `destination` on the same file would have to + /// wait on -- and since both connections are driven from this single + /// thread, that wait can never be satisfied ("database is locked"). + /// To avoid this, fully read and finalize the source query *before* + /// issuing any writes to `destination`. + fn copy_key_values( + source: &sqlite::Connection, + destination: &sqlite::Connection, + ) -> Result<(), TridentError> { + let mut rows: Vec<(String, String)> = Vec::new(); + { + let mut query_statement = source + .prepare("SELECT key, contents FROM keyvalue") + .structured(ServicingError::from(DatastoreError::ReadDatastore))?; + + loop { + match query_statement.next() { + Ok(State::Done) => break, + Err(e) => { + return Err(e) + .structured(ServicingError::from(DatastoreError::ReadDatastore)) + .message("Failed to get next keyvalue row while copying datastore"); + } + Ok(State::Row) => {} // continue below + } + + let key = query_statement + .read::(0) + .structured(ServicingError::from(DatastoreError::ReadDatastore))?; + let contents = query_statement + .read::(1) + .structured(ServicingError::from(DatastoreError::ReadDatastore))?; + + rows.push((key, contents)); + } + // `query_statement` is dropped here, finalizing the source + // SELECT and releasing its read lock before any writes below. + } + + for (key, contents) in rows { + let mut insert_statement = destination + .prepare( + "INSERT INTO keyvalue (key, contents) VALUES (?, ?) \ + ON CONFLICT(key) DO UPDATE SET contents = excluded.contents", + ) + .structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { key: key.clone() }, + })?; + insert_statement + .bind((1, &*key)) + .structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { key: key.clone() }, + })?; + insert_statement + .bind((2, &*contents)) + .structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { key: key.clone() }, + })?; + insert_statement + .next() + .structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { key }, + })?; + } + + Ok(()) + } + fn write_host_status( db: &sqlite::Connection, host_status: &HostStatus, @@ -219,6 +384,183 @@ impl DataStore { self.db = None; } + /// Retrieve a structured value stored under `key` in the datastore's + /// generic key-value table, if present. + /// + /// Values are serialized as JSON, so any type implementing + /// `serde::Serialize`/`serde::de::DeserializeOwned` can be stored, not + /// just `HostStatus`. + pub(crate) fn get_value( + &self, + key: &str, + ) -> Result, TridentError> { + let db = self + .db + .as_ref() + .structured(ServicingError::from(DatastoreError::OpenDatastore))?; + + let mut statement = db + .prepare("SELECT contents FROM keyvalue WHERE key = ?") + .structured(ServicingError::Datastore { + inner: DatastoreError::ReadKeyValue { + key: key.to_string(), + }, + })?; + statement + .bind((1, key)) + .structured(ServicingError::Datastore { + inner: DatastoreError::ReadKeyValue { + key: key.to_string(), + }, + })?; + + match statement.next().structured(ServicingError::Datastore { + inner: DatastoreError::ReadKeyValue { + key: key.to_string(), + }, + })? { + State::Row => { + let contents = + statement + .read::(0) + .structured(ServicingError::Datastore { + inner: DatastoreError::ReadKeyValue { + key: key.to_string(), + }, + })?; + let value = serde_json::from_str(&contents).structured( + InternalError::DeserializeValue { + key: key.to_string(), + }, + )?; + Ok(Some(value)) + } + State::Done => Ok(None), + } + } + + /// Store a structured value under `key` in the datastore's generic + /// key-value table, overwriting any previous value stored under the + /// same key. + /// + /// Values are serialized as JSON, so any type implementing + /// `serde::Serialize`/`serde::de::DeserializeOwned` can be stored, not + /// just `HostStatus`. + /// + /// `datastore_id` is currently the only first-party caller of the + /// generic key-value store, and it needs insert-if-absent semantics + /// (see `set_value_if_absent`) rather than an unconditional overwrite, + /// so this unconditional-overwrite variant is presently exercised only + /// by tests. It's kept as part of the generic key-value API (see + /// `get_value`) for future callers that do want overwrite semantics. + #[allow(dead_code)] + pub(crate) fn set_value(&self, key: &str, value: &T) -> Result<(), TridentError> { + let contents = serde_json::to_string(value).structured(InternalError::SerializeValue { + key: key.to_string(), + })?; + self.write_key_value_row( + key, + &contents, + "INSERT INTO keyvalue (key, contents) VALUES (?, ?) \ + ON CONFLICT(key) DO UPDATE SET contents = excluded.contents", + ) + } + + /// 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 (see [`Self::datastore_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, + value: &T, + ) -> Result<(), TridentError> { + let contents = serde_json::to_string(value).structured(InternalError::SerializeValue { + key: key.to_string(), + })?; + self.write_key_value_row( + key, + &contents, + "INSERT INTO keyvalue (key, contents) VALUES (?, ?) \ + ON CONFLICT(key) DO NOTHING", + ) + } + + /// Execute a parameterized `INSERT ... ON CONFLICT ...` against the + /// `keyvalue` table, binding `key` and `contents` as the two `?` + /// placeholders in `sql`. + fn write_key_value_row( + &self, + key: &str, + contents: &str, + sql: &str, + ) -> Result<(), TridentError> { + let db = self + .db + .as_ref() + .structured(ServicingError::from(DatastoreError::WriteToClosedDatastore))?; + + let mut statement = db.prepare(sql).structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { + key: key.to_string(), + }, + })?; + statement + .bind((1, key)) + .structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { + key: key.to_string(), + }, + })?; + statement + .bind((2, contents)) + .structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { + key: key.to_string(), + }, + })?; + statement.next().structured(ServicingError::Datastore { + inner: DatastoreError::WriteKeyValue { + key: key.to_string(), + }, + })?; + + Ok(()) + } + + /// Retrieve this datastore's unique database ID, generating and + /// persisting a new one on first access. + /// + /// This ID is stable for the lifetime of the datastore (surviving the + /// temporary-to-persistent transition performed by `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. + /// + /// 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 datastore_id(&mut self) -> Result { + if let Some(id) = self.get_value::(DATASTORE_ID_KEY)? { + return Ok(id); + } + + self.set_value_if_absent(DATASTORE_ID_KEY, &Uuid::new_v4())?; + + self.get_value::(DATASTORE_ID_KEY)? + .structured(InternalError::Internal( + "Datastore ID missing immediately after being inserted", + )) + } + /// 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. @@ -276,6 +618,31 @@ mod tests { temp_dir.close().unwrap(); } + #[test] + /// Regression test: `persist` supports a destination path equal to the + /// currently-open (temporary) datastore's own path -- the offline + /// provisioning flow does this. `copy_key_values` fully reads and + /// finalizes the source `SELECT` before writing to the destination + /// connection, so a self-persist (both connections pointing at the + /// same file) does not block the destination write on the source's + /// read lock. + fn test_persist_to_same_path_does_not_deadlock() { + 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 datastore_id = datastore.datastore_id().unwrap(); + + // Persist to the exact same path the datastore is currently open at. + datastore.persist(&datastore_path).unwrap(); + + assert_eq!( + datastore.datastore_id().unwrap(), + datastore_id, + "Datastore ID should survive a self-persist" + ); + } + #[test] fn test_parse_host_status() { let ds = super::DataStore { @@ -306,6 +673,147 @@ mod tests { .parse_host_status(Ok(serde_yaml::to_string(&valid_host_status).unwrap())) .is_some()); } + + #[test] + fn test_generic_key_value_store() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().join("db.sqlite"); + let db = super::DataStore::make_datastore(&path).unwrap(); + let datastore = super::DataStore { + db: Some(db), + host_status: Default::default(), + temporary: false, + }; + + // No value stored yet. + assert_eq!(datastore.get_value::("some-key").unwrap(), None); + + // Store and retrieve a value. + datastore + .set_value("some-key", &"some-value".to_string()) + .unwrap(); + assert_eq!( + datastore.get_value::("some-key").unwrap(), + Some("some-value".to_string()) + ); + + // Overwrite the value. + datastore + .set_value("some-key", &"other-value".to_string()) + .unwrap(); + assert_eq!( + datastore.get_value::("some-key").unwrap(), + Some("other-value".to_string()) + ); + + temp_dir.close().unwrap(); + } + + #[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, `datastore_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"); + + // Simulate a datastore created before the `keyvalue` table existed: + // create only the `hoststatus` table. + { + let db = sqlite::open(&path).unwrap(); + db.execute( + "CREATE TABLE IF NOT EXISTS hoststatus ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, + contents TEXT NOT NULL + )", + ) + .unwrap(); + } + + let mut datastore = super::DataStore::open(&path).unwrap(); + // Should not fail with "no such table: keyvalue". + datastore.datastore_id().unwrap(); + + temp_dir.close().unwrap(); + } + + #[test] + /// Regression test: a zero-byte or otherwise schema-less file at the + /// datastore path must be rejected by `open()` rather than silently + /// treated as a valid, freshly-provisioned datastore. + fn test_open_rejects_file_missing_hoststatus_table() { + let temp_dir = tempfile::tempdir().unwrap(); + let path = temp_dir.path().join("db.sqlite"); + + // Create a valid, but empty (no tables), SQLite file at the path -- + // simulating a truncated/corrupt existing datastore. + sqlite::open(&path).unwrap(); + + let result = super::DataStore::open(&path); + assert!( + result.is_err(), + "open() must reject a datastore file with no 'hoststatus' table" + ); + + temp_dir.close().unwrap(); + } + + #[test] + fn test_datastore_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 daemon RPC handlers + // concurrently calling `Trident::new` (and therefore + // `datastore_id`) against the same datastore path. + super::DataStore::make_datastore(&path).unwrap(); + + let barrier = std::sync::Arc::new(std::sync::Barrier::new(2)); + let handles: Vec<_> = (0..2) + .map(|_| { + let path = path.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + let mut datastore = super::DataStore::open(&path).unwrap(); + // Synchronize so both threads attempt "first access" + // (no database ID persisted yet) as close together + // as possible. + barrier.wait(); + datastore.datastore_id().unwrap() + }) + }) + .collect(); + + let ids: Vec = handles.into_iter().map(|h| h.join().unwrap()).collect(); + assert_eq!( + ids[0], ids[1], + "concurrent first access returned inconsistent database IDs" + ); + + temp_dir.close().unwrap(); + } + + #[test] + fn test_datastore_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.datastore_id().unwrap(); + // Calling datastore_id again should return the same ID, not generate a + // new one. + assert_eq!(datastore.datastore_id().unwrap(), id); + + temp_dir.close().unwrap(); + } } #[cfg(feature = "functional-test")] @@ -381,4 +889,24 @@ mod functional_test { ServicingState::Provisioned ); } + + #[functional_test] + fn test_datastore_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"); + + // Generate a database ID in the temporary datastore, then persist it. + let datastore_id = { + let mut datastore = DataStore::open_or_create(&datastore_temp_path).unwrap(); + let datastore_id = datastore.datastore_id().unwrap(); + datastore.persist(&datastore_path).unwrap(); + datastore_id + }; + + // Re-open the persisted datastore and verify the same database ID is + // returned, rather than a new one being generated. + let mut datastore = DataStore::open(&datastore_path).unwrap(); + assert_eq!(datastore.datastore_id().unwrap(), datastore_id); + } } diff --git a/crates/trident/src/lib.rs b/crates/trident/src/lib.rs index 8982b3f58..bb3f62fb4 100644 --- a/crates/trident/src/lib.rs +++ b/crates/trident/src/lib.rs @@ -207,6 +207,23 @@ impl Trident { info!("Running Trident in a container"); } + // Retrieve (or create, on first run) this host's unique + // database ID from the datastore actually used for servicing, + // and attach it to the shared TraceStream before any startup + // metrics are emitted, so every trace/metric -- including this + // very "trident_start" event -- carries it. This runs for every + // caller of `Trident::new` (both the CLI path and each daemon + // RPC handler), since they all supply `datastore_path`. + match DataStore::open_or_create(datastore_path).and_then(|mut ds| ds.datastore_id()) { + Ok(datastore_id) => { + info!("Datastore ID: {datastore_id}"); + tracestream.set_datastore_id(datastore_id.to_string()); + } + Err(e) => { + warn!("Failed to get or create database ID: {e:?}"); + } + } + // Trace features enabled in the Host Configuration. if let Some(hc) = &host_config { hc.feature_tracing(); diff --git a/crates/trident/src/logging/tracestream.rs b/crates/trident/src/logging/tracestream.rs index 132696833..d6f87fdf0 100644 --- a/crates/trident/src/logging/tracestream.rs +++ b/crates/trident/src/logging/tracestream.rs @@ -84,6 +84,7 @@ pub struct TraceStream { // TODO: Consider changing this to a LockOnce when rustc is updated to // >=1.70 target: Arc>>, + datastore_id: Arc>>, disabled: bool, } @@ -125,14 +126,46 @@ impl TraceStream { Ok(()) } + /// Set the database 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 database ID has been retrieved (see + /// `DataStore::datastore_id`). + pub fn set_datastore_id(&self, datastore_id: String) { + match self.datastore_id.write() { + Ok(mut val) => { + val.replace(datastore_id); + } + Err(_) => warn!("Failed to lock tracestream to set database ID"), + } + } + /// Create a Boxed TraceSender pub fn make_trace_sender(&self) -> Box { - Box::new(TraceSender::new(self.target.clone())) + self.make_trace_sender_with_metrics_path(TRIDENT_METRICS_FILE_PATH) + } + + /// Like `make_trace_sender`, but writes the local metrics file to + /// `metrics_file_path` instead of the real host path + /// (`TRIDENT_METRICS_FILE_PATH`). This lets tests exercise the full + /// metrics-writing pipeline against a throwaway temp file instead of a + /// real, shared host path, so they can be plain `#[test]`s instead of + /// needing a VM. + pub(crate) fn make_trace_sender_with_metrics_path( + &self, + metrics_file_path: &str, + ) -> Box { + Box::new(TraceSender::new( + self.target.clone(), + self.datastore_id.clone(), + metrics_file_path, + )) } } pub struct TraceSender { server: Arc>>, + datastore_id: Arc>>, client: reqwest::blocking::Client, metrics_file: Option, } @@ -144,11 +177,16 @@ struct ExecutionTime(Instant); /// the tracing-subscriber crate to handle the events and send them to the /// server. impl TraceSender { - fn new(server: Arc>>) -> Self { + fn new( + server: Arc>>, + datastore_id: Arc>>, + metrics_file_path: &str, + ) -> Self { Self { server, + datastore_id, client: reqwest::blocking::Client::new(), - metrics_file: match files::create_file(TRIDENT_METRICS_FILE_PATH) { + metrics_file: match files::create_file(metrics_file_path) { Ok(f) => Some(f), Err(err) => { eprintln!( @@ -164,6 +202,20 @@ impl TraceSender { self.server.read().map(|s| s.clone()).unwrap_or_default() } + /// Build the `additional_fields` map for a trace entry: the static + /// `ADDITIONAL_FIELDS`, plus the database ID (if one has been set via + /// `TraceStream::set_datastore_id`), so entries can be correlated back to a + /// specific host installation. + fn additional_fields(&self) -> BTreeMap { + let mut fields = ADDITIONAL_FIELDS.clone(); + if let Ok(datastore_id) = self.datastore_id.read() { + if let Some(datastore_id) = datastore_id.as_ref() { + fields.insert("datastore_id".to_string(), json!(datastore_id)); + } + } + fields + } + fn write_metric_to_file(&self, metric: String) { if let Some(mut file) = self.metrics_file.as_ref() { if let Err(e) = file.write_all(format!("{metric}\n").as_bytes()) { @@ -229,7 +281,7 @@ where timestamp: Utc::now(), metric_name, value: json!(value), - additional_fields: ADDITIONAL_FIELDS.clone(), + additional_fields: self.additional_fields(), platform_info: PLATFORM_INFO.clone(), }; @@ -309,7 +361,7 @@ where timestamp: Utc::now(), metric_name: span.name().to_string(), value: json!(visitor.fields), - additional_fields: ADDITIONAL_FIELDS.clone(), + additional_fields: self.additional_fields(), platform_info: PLATFORM_INFO.clone(), }; @@ -417,12 +469,20 @@ fn populate_platform_info() -> BTreeMap { mod tests { use super::*; - use std::{fs::File, io::Write}; + use std::{ + fs::File, + io::{BufRead, BufReader, Write}, + }; + + use tracing_subscriber::{filter, layer::SubscriberExt}; #[test] fn test_tracestream() { + let temp_dir = tempfile::tempdir().unwrap(); + let metrics_path = temp_dir.path().join("metrics.jsonl"); let tracestream = TraceStream::default(); - let trace_sender = tracestream.make_trace_sender(); + let trace_sender = + tracestream.make_trace_sender_with_metrics_path(metrics_path.to_str().unwrap()); assert!( trace_sender.get_server().is_none(), "tracestream should not have a server" @@ -441,8 +501,11 @@ mod tests { #[test] fn test_lock() { + let temp_dir = tempfile::tempdir().unwrap(); + let metrics_path = temp_dir.path().join("metrics.jsonl"); let mut tracestream = TraceStream::default(); - let trace_sender = tracestream.make_trace_sender(); + let trace_sender = + tracestream.make_trace_sender_with_metrics_path(metrics_path.to_str().unwrap()); assert!( trace_sender.get_server().is_none(), @@ -476,31 +539,24 @@ mod tests { let uuid = read_product_uuid(filepath.to_str().unwrap().to_string()); assert_eq!(uuid, "test_uuid"); } -} - -#[cfg(feature = "functional-test")] -#[cfg_attr(not(test), allow(unused_imports, dead_code))] -mod functional_test { - use super::*; - use std::io::{BufRead, BufReader}; - - use tracing_subscriber::{filter, layer::SubscriberExt}; - - use pytest_gen::functional_test; - - #[functional_test] + #[test] fn test_tracestream_write_metric_event_to_file() { + let temp_dir = tempfile::tempdir().unwrap(); + let metrics_path = temp_dir.path().join("metrics.jsonl"); let tracestream = TraceStream::default(); let trace_sender = tracestream - .make_trace_sender() + .make_trace_sender_with_metrics_path(metrics_path.to_str().unwrap()) .with_filter(filter::LevelFilter::INFO); - tracing::subscriber::set_global_default( + // Use a thread-local scoped default subscriber (rather than + // `set_global_default`) since this is a plain `#[test]` that may run + // concurrently with other tests in the same process -- the global + // default can only be set once per process, but the scoped default + // is per-thread and automatically restored when `_guard` drops. + let _guard = tracing::subscriber::set_default( tracing_subscriber::Registry::default().with(trace_sender), - ) - .context("Failed to set global default subscriber") - .unwrap(); + ); tracing::info!(metric_name = "test_metric", value = true); @@ -508,7 +564,7 @@ mod functional_test { std::thread::sleep(std::time::Duration::from_millis(100)); // Check if the specific metric exists in the file. - let file = File::open(TRIDENT_METRICS_FILE_PATH).unwrap(); + let file = File::open(&metrics_path).unwrap(); let reader = BufReader::new(file); let lines: Vec = reader.lines().map(|l| l.unwrap()).collect(); @@ -522,52 +578,61 @@ mod functional_test { ); } - #[functional_test] - fn test_populate_additional_fields() { - let additional_fields = populate_additional_fields(); - assert_eq!( - additional_fields.get("trident_version").unwrap(), - &json!(TRIDENT_VERSION) - ); - } + #[test] + /// Regression test: `TraceStream::set_datastore_id` must actually + /// reach the serialized trace entry's `additional_fields.datastore_id` + /// -- the metric/span tests above only assert on `metric_name`/`value` + /// and would still pass even if the database ID were never copied + /// into `additional_fields`. + fn test_tracestream_datastore_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_datastore_id("test-datastore-id".to_string()); + let trace_sender = tracestream + .make_trace_sender_with_metrics_path(metrics_path.to_str().unwrap()) + .with_filter(filter::LevelFilter::INFO); - #[functional_test] - fn test_populate_platform_info() { - let mut expected_platform_info = BTreeMap::new(); - expected_platform_info.insert( - "asset_id".to_string(), - json!(read_product_uuid(PRODUCT_UUID_FILE.into())), - ); - expected_platform_info.insert("os_release".to_string(), json!(get_os_release())); - expected_platform_info.insert("total_cpu".to_string(), json!(4)); - expected_platform_info.insert("total_memory_gib".to_string(), json!(6)); - expected_platform_info.insert( - "kernel_version".to_string(), - json!(uname::kernel_release().unwrap().trim()), + // 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), ); - // Call the function to get the actual result. - let platform_info = populate_platform_info(); + tracing::info!(metric_name = "test_metric_with_datastore_id", value = true); - // Assert that the actual result matches the expected result. - assert_eq!( - platform_info, expected_platform_info, - "Platform info does not match the expected result" + // 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_datastore_id""#) + && line.contains(r#""datastore_id":"test-datastore-id""#) + }); + + assert!( + metric_found, + "Expected metric with datastore_id field not found in the local metrics file" ); } - #[functional_test] + #[test] fn test_tracestream_write_span_metric_to_file() { + let temp_dir = tempfile::tempdir().unwrap(); + let metrics_path = temp_dir.path().join("metrics.jsonl"); let tracestream = TraceStream::default(); let trace_sender = tracestream - .make_trace_sender() + .make_trace_sender_with_metrics_path(metrics_path.to_str().unwrap()) .with_filter(filter::LevelFilter::INFO); - tracing::subscriber::set_global_default( + // 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), - ) - .context("Failed to set global default subscriber") - .unwrap(); + ); // Call test function that will create a span simulate_function_span(); @@ -576,7 +641,7 @@ mod functional_test { std::thread::sleep(std::time::Duration::from_millis(100)); // Check if the specific metric exists in the file. - let file = File::open(TRIDENT_METRICS_FILE_PATH).unwrap(); + let file = File::open(&metrics_path).unwrap(); let reader = BufReader::new(file); let lines: Vec = reader.lines().map(|l| l.unwrap()).collect(); @@ -594,3 +659,51 @@ mod functional_test { #[tracing::instrument(name = "test_span", skip_all)] fn simulate_function_span() {} } + +#[cfg(feature = "functional-test")] +#[cfg_attr(not(test), allow(unused_imports, dead_code))] +mod functional_test { + use super::*; + + use pytest_gen::functional_test; + + // These two remain functional tests (VM-only) because they assert + // against the actual host's hardware/platform info (CPU count, memory, + // product UUID, os-release, kernel version) -- unlike the metrics-file + // tests above, there's no way to inject a fake value here, so the + // result is inherently host-dependent. + + #[functional_test] + fn test_populate_additional_fields() { + let additional_fields = populate_additional_fields(); + assert_eq!( + additional_fields.get("trident_version").unwrap(), + &json!(TRIDENT_VERSION) + ); + } + + #[functional_test] + fn test_populate_platform_info() { + let mut expected_platform_info = BTreeMap::new(); + expected_platform_info.insert( + "asset_id".to_string(), + json!(read_product_uuid(PRODUCT_UUID_FILE.into())), + ); + expected_platform_info.insert("os_release".to_string(), json!(get_os_release())); + expected_platform_info.insert("total_cpu".to_string(), json!(4)); + expected_platform_info.insert("total_memory_gib".to_string(), json!(6)); + expected_platform_info.insert( + "kernel_version".to_string(), + json!(uname::kernel_release().unwrap().trim()), + ); + + // Call the function to get the actual result. + let platform_info = populate_platform_info(); + + // Assert that the actual result matches the expected result. + assert_eq!( + platform_info, expected_platform_info, + "Platform info does not match the expected result" + ); + } +} diff --git a/crates/trident/src/main.rs b/crates/trident/src/main.rs index 9d3f2b7bf..c1d996459 100644 --- a/crates/trident/src/main.rs +++ b/crates/trident/src/main.rs @@ -165,6 +165,10 @@ fn run_trident( ) .message("Failed to initialize Trident")?; + // `Trident::new` has already retrieved (or created) this + // host's persisted database ID and attached it to the + // shared TraceStream, so every trace/metric emitted from + // here on -- including "trident_start" -- carries it. let mut datastore = DataStore::open_or_create(agent_config.datastore_path()) .message("Failed to open datastore")?; diff --git a/crates/trident/src/subsystems/management.rs b/crates/trident/src/subsystems/management.rs index dbe408d2f..ab66e65aa 100644 --- a/crates/trident/src/subsystems/management.rs +++ b/crates/trident/src/subsystems/management.rs @@ -92,43 +92,67 @@ fn configure_agent_config( datastore_path: &Path, is_root_verity: bool, ) -> Result<(), TridentError> { - // Ensure that Trident agent config exists with correct datastore path - if Path::new(agent_config_path).exists() { - // If the agent config exists, check that the datastore matches the expected path. - if let Ok(contents) = std::fs::read_to_string(agent_config_path) { - let mut datastore_path_configured = TRIDENT_DATASTORE_PATH_DEFAULT; - for line in contents.lines() { - if let Some(path) = line.strip_prefix("DatastorePath=") { - datastore_path_configured = path.trim(); - break; - } - } - // If the datastore path in the agent config does not match the expected path, - // return an error. - if datastore_path != Path::new(datastore_path_configured) { - return Err(TridentError::new( - InvalidInputError::ImageBadAgentConfiguration, - )) - .message(format!( - "Datastore path in agent config ({}) does not match expected path ({})", - datastore_path_configured, - datastore_path.display() - )); - } + // Ensure that Trident agent config exists with correct datastore path. + // A missing file and an existing-but-unreadable file are both treated + // as "no contents to preserve" -- except an unreadable file is left + // untouched entirely, matching the prior behavior of silently skipping + // the whole check when the file exists but can't be read. + let contents = if Path::new(agent_config_path).exists() { + match std::fs::read_to_string(agent_config_path) { + Ok(contents) => contents, + Err(_) => return Ok(()), } - } else if datastore_path != Path::new(TRIDENT_DATASTORE_PATH_DEFAULT) { - // Only attempt to create the agent config if the datastore path is not the default. + } else { + String::new() + }; + + let mut datastore_path_line_present = false; + let mut datastore_path_configured = TRIDENT_DATASTORE_PATH_DEFAULT; + for line in contents.lines() { + if let Some(path) = line.strip_prefix("DatastorePath=") { + datastore_path_line_present = true; + datastore_path_configured = path.trim(); + break; + } + } + if datastore_path_line_present { + // An explicit DatastorePath= line is present: it must match + // the expected path exactly. + if datastore_path != Path::new(datastore_path_configured) { + return Err(TridentError::new( + InvalidInputError::ImageBadAgentConfiguration, + )) + .message(format!( + "Datastore path in agent config ({}) does not match expected path ({})", + datastore_path_configured, + datastore_path.display() + )); + } + } else if datastore_path != Path::new(TRIDENT_DATASTORE_PATH_DEFAULT) { + // No DatastorePath= line -- either the file doesn't exist yet, or + // it exists but carries other settings (e.g. Telemetry=) that must + // be preserved. Missing DatastorePath only implies the default + // path, so if a non-default path is expected, merge a + // DatastorePath= line into whatever contents already exist (empty, + // if the file didn't exist) rather than treating it as a mismatch. if is_root_verity { - // For root-verity, do not attempt to create the agent config. + // For root-verity, do not attempt to create or modify the agent config. return Err(TridentError::new( InvalidInputError::ImageBadAgentConfiguration, )) - .message("Agent configuration file does not exist and root filesystem is verity"); + .message( + "Agent configuration file does not set a non-default datastore path \ + and root filesystem is verity", + ); } - let datastore_configuration = format!("DatastorePath={}", datastore_path.display()); - fs::write(agent_config_path, datastore_configuration).structured( + let mut updated_contents = contents; + if !updated_contents.is_empty() && !updated_contents.ends_with('\n') { + updated_contents.push('\n'); + } + updated_contents.push_str(&format!("DatastorePath={}\n", datastore_path.display())); + fs::write(agent_config_path, updated_contents).structured( ServicingError::CreateConfigurationFile { path: agent_config_path.into(), }, @@ -270,5 +294,60 @@ mod tests { ) .unwrap_err(); } + + { + // Regression test: agent config exists but only carries other + // settings (e.g. Telemetry=OptIn), with no DatastorePath= line. + // A non-default datastore path must be merged in, preserving the + // existing settings, rather than treated as a mismatch. + let agent_config_folder = tempfile::tempdir().unwrap(); + let agent_config_path = agent_config_folder.path().join("trident.conf"); + fs::write( + &agent_config_path, + "Telemetry=OptIn +", + ) + .unwrap(); + + configure_agent_config( + &agent_config_path.to_string_lossy(), + Path::new(nonstandard_datastore_path), + false, + ) + .unwrap(); + + let contents = std::fs::read_to_string(&agent_config_path).unwrap(); + assert!(contents.contains("Telemetry=OptIn")); + assert!(contents.contains(&format!("DatastorePath={nonstandard_datastore_path}"))); + } + + { + // Same as above, but root-verity: must not attempt to modify + // the agent config, and must error like the "file does not + // exist" root-verity case. + let agent_config_folder = tempfile::tempdir().unwrap(); + let agent_config_path = agent_config_folder.path().join("trident.conf"); + fs::write( + &agent_config_path, + "Telemetry=OptIn +", + ) + .unwrap(); + + configure_agent_config( + &agent_config_path.to_string_lossy(), + Path::new(nonstandard_datastore_path), + true, + ) + .unwrap_err(); + + // The file must be left untouched. + let contents = std::fs::read_to_string(&agent_config_path).unwrap(); + assert_eq!( + contents, + "Telemetry=OptIn +" + ); + } } } diff --git a/crates/trident_api/src/error.rs b/crates/trident_api/src/error.rs index 064be997d..1a8e4b9f3 100644 --- a/crates/trident_api/src/error.rs +++ b/crates/trident_api/src/error.rs @@ -124,6 +124,12 @@ pub enum InternalError { #[error("Failed to serialize Host Status")] SerializeHostStatus, + #[error("Failed to serialize value for datastore key '{key}'")] + SerializeValue { key: String }, + + #[error("Failed to deserialize value for datastore key '{key}'")] + DeserializeValue { key: String }, + #[error("Failed to set up extension images on the target OS")] SetUpExtensionImages, @@ -742,11 +748,17 @@ pub enum DatastoreError { #[error("Failed to read from datastore")] ReadDatastore, + #[error("Failed to read key '{key}' from datastore")] + ReadKeyValue { key: String }, + #[error("Failed to write to datastore as it is closed")] WriteToClosedDatastore, #[error("Failed to write to datastore")] WriteToDatastore, + + #[error("Failed to write key '{key}' to datastore")] + WriteKeyValue { key: String }, } impl ServicingError {