datastore: add generic key-value storage and persisted correlation-id by bfjelds · Pull Request #773 · microsoft/trident

bfjelds (bfjelds) · 2026-09-03T19:17:36Z

Add a generic keyvalue table to the SQLite datastore so arbitrary structured data (JSON-serialized) can be stored/retrieved by key, not just HostStatus. DataStore::get_value/set_value provide the generic API; keyvalue rows are carried over when a temporary datastore is persisted, same as HostStatus.

As a first consumer, add DataStore::correlation_id(), which generates and persists a UUID on first access and returns the same value on every subsequent call. The correlation ID is retrieved at trident CLI startup and attached to every trace/metric entry via TraceStream::set_correlation_id, so all tracing/telemetry for a given host installation can be correlated.

Future application: trident-acl-agent uses a file (state.json) to track state, replace state.json with this datastore storage.

Related PRs in stack:

azure-pipelines · 2026-09-03T19:17:58Z

Azure Pipelines: Successfully started running 1 pipeline(s). There may be pipelines that require an authorized user to comment /azp run to run.

Add a generic keyvalue table to the SQLite datastore so arbitrary structured data (JSON-serialized) can be stored/retrieved by key, not just HostStatus. DataStore::get_value/set_value provide the generic API; keyvalue rows are carried over when a temporary datastore is persisted, same as HostStatus. As a first consumer, add DataStore::correlation_id(), which generates and persists a UUID on first access and returns the same value on every subsequent call. The correlation ID is retrieved at trident CLI startup and attached to every trace/metric entry via TraceStream::set_correlation_id, so all tracing/telemetry for a given host installation can be correlated. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot

🟡 Changes recommended

There are correctness issues in datastore initialization/copy logic (schema typo and silent error swallowing during key/value copy) that can lead to incorrect behavior while still reporting success.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds a generic, JSON-backed key/value store to Trident’s SQLite datastore and introduces a persisted per-host correlation ID that gets attached to TraceStream telemetry, enabling cross-log/trace correlation for a given installation.

Changes:

Add a keyvalue table plus DataStore::get_value / DataStore::set_value for storing arbitrary JSON-serialized structured values by key.
Add DataStore::correlation_id() to generate/persist a UUID on first access and reuse it thereafter (including across temp→persisted datastore transition).
Plumb the correlation ID into tracing/metrics by adding it to TraceStream “additional_fields”, and initialize it at CLI startup.

File summaries

File	Description
crates/trident/src/main.rs	Loads/creates the datastore at startup and sets the correlation ID on TraceStream.
crates/trident/src/logging/tracestream.rs	Adds correlation ID plumbing so every trace/metric entry can include it in `additional_fields`.
crates/trident/src/datastore.rs	Adds `keyvalue` table creation, generic get/set APIs, correlation ID persistence, and carry-over on `persist()`.
crates/trident_api/src/error.rs	Extends structured error enums to cover key/value serialize/deserialize and key-specific datastore read/write failures.

Review details

Files reviewed: 4/4 changed files
Comments generated: 2
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

- Fix DEFALUT -> DEFAULT typo in hoststatus schema so timestamp auto-populates as intended. - copy_key_values now returns an error instead of warning and silently stopping when reading a keyvalue row fails, so persist() cannot report success after a partial/failed copy.

Copilot

🟡 Changes recommended

Existing datastores lack schema migration, and several telemetry paths do not receive the persisted correlation ID.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

crates/trident/src/datastore.rs:146

Issue: Existing datastore files are not upgraded with this table. Evidence: open_or_create calls open directly when the path exists, while only make_datastore executes this CREATE TABLE; consequently correlation_id() fails with “no such table: keyvalue” on every upgraded host. Suggestion: run a shared idempotent schema initializer for both newly created and existing connections, and cover opening a legacy hoststatus-only database.

        db.execute(
            "CREATE TABLE IF NOT EXISTS keyvalue (
                key TEXT PRIMARY KEY,
                contents TEXT NOT NULL
            )",
        )
        .structured(ServicingError::from(DatastoreError::InitializeDatastore))?;

Files reviewed: 4/4 changed files
Comments generated: 1
Review effort level: Balanced

…on paths Trident::new emitted the "trident_start" metric before the CLI path (main.rs) had retrieved the persisted correlation ID and attached it to the shared TraceStream, so that very first startup event -- and any daemon RPC handler that never ran the CLI's correlation-ID block at all -- went out without it. Move the correlation ID retrieval into Trident::new itself, using the datastore_path it already receives, and set it on the TraceStream before "trident_start" is emitted. Every caller of Trident::new (the CLI path and each daemon gRPC service handler) now gets the same treatment for free, since they all supply datastore_path. main.rs's run_trident no longer needs its own post-hoc correlation-ID block or the pre-emptive tracestream clone that existed only to work around the ordering problem. Note: multiboot installs still open the persisted datastore, then swap in a fresh temporary one during install() (lib.rs) before the new installation is later persisted. Whether the same correlation ID should be carried forward across that swap (vs. each multiboot install getting its own) is a servicing-flow behavior question left as a follow-up rather than guessed at here.

Copilot

🟡 Changes recommended

Existing databases lack schema migration, one startup metric remains uncorrelated, and concurrent initialization can produce inconsistent IDs.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

crates/trident/src/datastore.rs:420

Issue: First-access ID creation is a non-atomic read-then-overwrite sequence. Evidence: Two processes can both observe no row, generate different UUIDs, and then set_value uses ON CONFLICT ... DO UPDATE; one caller returns an ID that the other immediately replaces, splitting telemetry for the same datastore. Suggestion: atomically insert only if absent (ideally in a transaction) and then read/return the stored winner.

        if let Some(id) = self.get_value::<Uuid>(CORRELATION_ID_KEY)? {
            return Ok(id);
        }

        let id = Uuid::new_v4();

Files reviewed: 5/5 changed files
Comments generated: 2
Review effort level: Balanced

Fix two issues flagged by Copilot review: - DataStore::open() (used for existing datastores) never ran the CREATE TABLE statements that make_datastore() runs on create, so an existing datastore created before the `keyvalue` table existed would fail correlation_id() with "no such table: keyvalue" on upgrade. Extract table creation into an idempotent ensure_schema() helper and call it from both open() and make_datastore(). Add a regression test covering a pre-existing datastore missing the keyvalue table. - Trident::new() called hc.feature_tracing() (which emits the host_config_feature_usage tracing event) before retrieving and attaching the correlation ID, so that first metric was missing the field despite every other trace/metric carrying it. Move correlation ID retrieval earlier, before feature_tracing() and any other tracing event.

Copilot

🔵 Needs a closer look

Concurrent first access can return inconsistent correlation IDs, and telemetry enrichment lacks direct regression coverage.

Review details

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

crates/trident/src/datastore.rs:435

Issue: First access is not atomic, so concurrent datastore connections can return different correlation IDs for the same installation. Evidence: Both callers can observe None, generate distinct UUIDs, and then this upsert path overwrites whichever UUID was inserted first; the first caller continues tracing with an ID that is no longer persisted. Suggestion: initialize with INSERT ... ON CONFLICT DO NOTHING, then read and return the row that actually won (ideally in a transaction), and add a concurrent-connection regression test.
crates/trident/src/logging/tracestream.rs:198
Issue: The core telemetry enrichment is not covered by the existing TraceStream tests. Evidence: The event/span tests only search for metric_name/value, so they still pass if the shared correlation ID is never copied into additional_fields. Suggestion: set a correlation ID before creating/emitting through the sender and assert the serialized entry contains that exact value under additional_fields.correlation_id.

Files reviewed: 5/5 changed files
Comments generated: 0 new
Review effort level: Balanced

correlation_id() previously read the keyvalue table, and if absent, generated a new UUID and wrote it unconditionally. Two connections racing on first access could each generate a different UUID and the second write would silently overwrite the first, leaving a caller who already captured the first UUID tracing with an ID no longer persisted. Fix: use an atomic INSERT ... ON CONFLICT DO NOTHING (set_value_if_absent) to claim the row, then re-read it, so all racing callers converge on whichever UUID actually got persisted. Also set a 5s SQLite busy timeout on every connection open, since the atomic insert path can hit SQLITE_BUSY under concurrent writes with the default 0ms timeout. Adds a concurrency regression test (two threads, separate connections, barrier-synchronized) asserting both observe the same correlation ID. logging/tracestream: add a regression test asserting set_correlation_id is actually copied into additional_fields on emitted trace/metric entries, closing a gap where existing tests only checked metric_name/ value and would pass even if the correlation ID were dropped.

bfjelds (bfjelds) · 2026-09-04T01:14:53Z

Addressed two issues Copilot flagged in the review summary ("suppressed comments" section, not posted as separate inline threads):

datastore.rs — concurrent first access to correlation_id() could return inconsistent IDs. Fixed by replacing the read-then-write pattern with an atomic INSERT ... ON CONFLICT DO NOTHING (new set_value_if_absent) followed by a re-read, so racing connections converge on whichever UUID actually got persisted instead of silently overwriting each other. Also added a 5s SQLite busy timeout on connection open, since the atomic-insert path can hit SQLITE_BUSY under concurrent writers with the default 0ms timeout. Added a two-thread, barrier-synchronized regression test.
logging/tracestream.rs — no regression test for correlation-ID enrichment. Added a test asserting set_correlation_id is actually copied into additional_fields on emitted trace/metric entries, since the existing tests only asserted on metric_name/value.

Fixed in 6a70c74.

Copilot

🟡 Changes recommended

Same-path persistence can fail with a SQLite lock while copying key-value rows.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Files reviewed: 5/5 changed files
Comments generated: 1
Review effort level: Balanced

copy_key_values kept the source SELECT statement active while writing to the destination connection. persist() supports a destination path equal to the currently-open (temporary) datastores own path -- the offline provisioning flow does this -- so source and destination can be two live connections to the same underlying file. SQLite locking is per-connection, so the destination write would wait on the sources still-active read lock, surfacing as "database is locked" (bounded only by the busy timeout, not resolved by it). Fix: fully read and finalize the source query before issuing any writes to the destination. Adds a regression test that persists a temporary datastore to its own path.

Copilot

🟢 Approval recommended

The implementation addresses schema upgrades, concurrent initialization, persistence, and telemetry propagation with regression coverage.

Review details

Files reviewed: 5/5 changed files
Comments generated: 0 new
Review effort level: Balanced

… tests Several tests marked #[functional_test] (VM-only) did not actually need a VM: they only touched a TempDir-based SQLite file, or wrote to /var/log/trident-metrics.jsonl only because that path was hardcoded into TraceSender::new, not because the test logic itself needed a real host path. Regular #[test]s are faster and can run directly on the host, so: - Add TraceStream::make_trace_sender_with_metrics_path so the local metrics file location is injectable. Production (main.rs) still goes through TRIDENT_METRICS_FILE_PATH via make_trace_sender(); tests now point it at a throwaway temp file. - Move test_tracestream_write_metric_event_to_file, test_tracestream_write_span_metric_to_file, and test_tracestream_correlation_id_written_to_additional_fields (and fix up test_tracestream/test_lock, which already silently touched the real path) from #[functional_test] to #[test]. Switch these from tracing::subscriber::set_global_default (process-wide, once-only -- conflicts across tests sharing a process) to set_default (thread- local, scoped to a guard), since they can now run alongside other tests in the same process. test_populate_additional_fields and test_populate_platform_info stay functional tests: they assert against the real hosts hardware/platform info, which cannot be faked. - Move test_persist_to_same_path_does_not_deadlock (added this session) from #[functional_test] to #[test]: it only used a TempDir SQLite file and never needed VM isolation.

bfjelds (bfjelds) force-pushed the user/bfjelds/datastore-generic-storage branch from a2d59c7 to 720309e Compare September 3, 2026 19:57

bfjelds (bfjelds) changed the title ~~datastore: add generic key-value storage and persisted machine ID~~ datastore: add generic key-value storage and persisted correlation-id Sep 3, 2026

bfjelds (bfjelds) requested a lite review from Copilot September 3, 2026 23:40

Copilot started reviewing on behalf of bfjelds (bfjelds) September 3, 2026 23:40 View session

Copilot AI reviewed Sep 3, 2026

View reviewed changes

Comment thread crates/trident/src/datastore.rs

Comment thread crates/trident/src/datastore.rs Outdated

bfjelds (bfjelds) requested a balanced review from Copilot September 4, 2026 00:01

Copilot started reviewing on behalf of bfjelds (bfjelds) September 4, 2026 00:02 View session