feat: Store S3 metadata into segement metadata (#311)

* feat: Store S3 metadata into segement metadata

- Make segment metadata writing lazy so late-bound items will still be
  persisted into segment metadata
- update sgement metadata now merges entries
- copy s3 config options into segment metadata

* PR feedback

* Eagerly populate worker_ids when attaching a runtime

Pre-populate RuntimeContext::worker_ids in attach_runtime() right after
metrics_and_base is set, instead of lazily on each worker thread's first
poll/park event. We already know num_workers and base from RuntimeMetrics
at that point, so segment metadata is complete from the first flush cycle.

* Eagerly populate worker_ids when attaching a runtime

Pre-populate RuntimeContext::worker_ids in attach_runtime() right after
metrics_and_base is set, instead of lazily on each worker thread's first
poll/park event. We already know num_workers and base from RuntimeMetrics
at that point, so segment metadata is complete from the first flush cycle.

Strengthen the propagates_second_runtime_metadata test to assert exact
worker ID values (0,1 and 2,3) now that they're available immediately.

* Eagerly populate worker_ids when attaching a runtime

Pre-populate RuntimeContext::worker_ids in attach_runtime() right after
metrics_and_base is set, instead of lazily on each worker thread's first
poll/park event. We already know num_workers and base from RuntimeMetrics
at that point, so segment metadata is complete from the first flush cycle.

Strengthen tests to assert exact worker ID values now that they're
available immediately.

Author: rcoh

PROGRESS.md

@@ -0,0 +1,26 @@
+# Eager worker_ids population
+
+## What changed
+Pre-populate `RuntimeContext::worker_ids` in `attach_runtime()` immediately after `metrics_and_base` is set, instead of lazily on each worker thread's first poll/park event.
+
+## Why
+When a runtime is attached, we already know `num_workers` and `base` from `RuntimeMetrics`. Eagerly populating the map means `metadata_entry()` returns complete runtime→worker mappings from the very first flush cycle, rather than converging over time as workers come online.
+
+## Change
+One block added in `attach_runtime()` (recorder/mod.rs):
+```rust
+{
+    let mut ids = ctx.worker_ids.write().unwrap();
+    for i in 0..num_workers {
+        ids.insert(i as usize, base + i);
+    }
+}
+```
+
+`register_worker_if_needed` remains idempotent — guarded by a thread-local flag, and the insert is a no-op overwrite with the same value.
+
+## Verification
+- `cargo fmt --check` ✅
+- `cargo clippy --all-targets --all-features` ✅
+- `cargo nextest run` — 460/461 pass (1 pre-existing CPU-load-sensitive flaky test in dial9-perf-self-profile)
+- `cargo nextest run --stress-duration 20s` — 2 iterations, 460/460 pass

dial9-tokio-telemetry/Cargo.toml

@@ -49,7 +49,7 @@ tracing-layer = ["dep:tracing-subscriber"]
 worker-s3 = ["dep:aws-sdk-s3-transfer-manager", "dep:aws-sdk-s3", "dep:aws-config", "dep:time"]
 
 [dev-dependencies]
-dial9-tokio-telemetry = { path = ".", features = ["analysis", "tracing-layer"] }
+dial9-tokio-telemetry = { path = ".", features = ["analysis", "tracing-layer", "worker-s3"] }
 assert2 = { workspace = true }
 criterion = "0.5"
 clap = { version = "4", features = ["derive"] }

dial9-tokio-telemetry/src/background_task/s3.rs

@@ -125,6 +125,18 @@ impl S3Config {
         &self.bucket
     }
 
+    pub(crate) fn as_metadata(&self) -> impl Iterator<Item = (&str, &str)> {
+        [
+            ("bucket", self.bucket.as_str()),
+            ("service_name", self.service_name.as_str()),
+            ("instance_path", self.instance_path.as_str()),
+            ("boot_id", self.boot_id.as_str()),
+        ]
+        .into_iter()
+        .chain(self.prefix.as_ref().map(|p| ("prefix", p.as_str())))
+        .chain(self.region.as_ref().map(|r| ("region", r.as_str())))
+    }
+
     /// Optional region override for the S3 client.
     pub(crate) fn region(&self) -> Option<&str> {
         self.region.as_deref()

dial9-tokio-telemetry/src/telemetry/recorder/mod.rs

@@ -317,6 +317,16 @@ fn attach_runtime(
             });
         });
 
+    // Eagerly populate worker_ids so segment metadata is complete from the
+    // first flush cycle, rather than waiting for each worker thread to lazily
+    // register on its first poll/park event.
+    {
+        let mut ids = ctx.worker_ids.write().unwrap();
+        for i in 0..num_workers {
+            ids.insert(i as usize, base + i);
+        }
+    }
+
     shared.contexts.lock().unwrap().push(ctx);
 
     Ok(runtime)
@@ -1007,6 +1017,15 @@ impl TelemetryCore {
         let shared = Arc::new(SharedState::new(start_mono_ns));
         #[allow(unused_mut)]
         let mut event_writer = EventWriter::new(Box::new(writer));
+        #[cfg(feature = "worker-s3")]
+        if let Some(s3_config) = s3_config.as_ref() {
+            event_writer.update_segment_metadata(
+                s3_config
+                    .as_metadata()
+                    .map(|(k, v)| (k.to_string(), v.to_string()))
+                    .collect(),
+            )
+        }
 
         #[cfg(feature = "cpu-profiling")]
         {
@@ -1757,15 +1776,18 @@ mod tests {
             "expected at least one SegmentMetadata event in trace files"
         );
 
-        // At least one segment's metadata should contain both runtime mappings.
+        // At least one segment's metadata should contain both runtime mappings
+        // with the exact worker IDs (eagerly populated at attach time).
         let has_both = all_metadata.iter().any(|entries| {
-            let has_main = entries.iter().any(|(k, _)| k == "runtime.main");
-            let has_io = entries.iter().any(|(k, _)| k == "runtime.io");
+            let has_main = entries
+                .iter()
+                .any(|(k, v)| k == "runtime.main" && v == "0,1");
+            let has_io = entries.iter().any(|(k, v)| k == "runtime.io" && v == "2,3");
             has_main && has_io
         });
         assert!(
             has_both,
-            "expected segment metadata to contain both runtime.main and runtime.io, \
+            "expected segment metadata to contain runtime.main=0,1 and runtime.io=2,3, \
              got: {all_metadata:?}"
         );
     }