diff --git a/Cargo.lock b/Cargo.lock index 8e617aeb..e5f9c0e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -181,15 +181,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "bitmaps" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "031043d04099746d8db04daf1fa424b2bc8bd69d92b25962dcde24da39ab64a2" -dependencies = [ - "typenum", -] - [[package]] name = "block-buffer" version = "0.10.4" @@ -512,7 +503,6 @@ dependencies = [ "futures", "futures-core", "futures-util", - "im", "opentelemetry", "opentelemetry_sdk", "parking_lot", @@ -966,21 +956,6 @@ dependencies = [ "icu_properties", ] -[[package]] -name = "im" -version = "15.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0acd33ff0285af998aaf9b57342af478078f53492322fafc47450e09397e0e9" -dependencies = [ - "bitmaps", - "rand_core 0.6.4", - "rand_xoshiro", - "serde", - "sized-chunks", - "typenum", - "version_check", -] - [[package]] name = "indexmap" version = "2.13.0" @@ -1473,15 +1448,6 @@ dependencies = [ "getrandom 0.3.4", ] -[[package]] -name = "rand_xoshiro" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa" -dependencies = [ - "rand_core 0.6.4", -] - [[package]] name = "redox_syscall" version = "0.5.18" @@ -1786,16 +1752,6 @@ dependencies = [ "rand_core 0.6.4", ] -[[package]] -name = "sized-chunks" -version = "0.6.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16d69225bde7a69b235da73377861095455d298f2b970996eec25ddbb42b3d1e" -dependencies = [ - "bitmaps", - "typenum", -] - [[package]] name = "slab" version = "0.4.12" diff --git a/Cargo.toml b/Cargo.toml index cbdc0256..ea83b23f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,7 +36,6 @@ serde = { workspace = true } serde_json = { workspace = true } derive_builder = { workspace = true } uuid = { workspace = true } -im = { workspace = true } pin-project = { workspace = true } parking_lot = { workspace = true } tokio = { workspace = true } @@ -83,7 +82,6 @@ sqlx = { version = "0.8", default-features = false, features = ["macros", "runti tokio = { version = "1.52", features = ["rt-multi-thread", "macros", "time"] } thiserror = "2.0" uuid = { version = "1.23", features = ["serde", "v7"] } -im = { version = "15.1", features = ["serde"] } pin-project = "1.1" tracing = { version = "0.1.41", default-features = false } tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/src/context/mod.rs b/src/context/mod.rs index 38f0f7b7..0b47d428 100644 --- a/src/context/mod.rs +++ b/src/context/mod.rs @@ -74,7 +74,7 @@ mod with_event_context; use serde::{Deserialize, Serialize}; -use std::{borrow::Cow, cell::RefCell, rc::Rc}; +use std::{borrow::Cow, cell::RefCell, rc::Rc, sync::Arc}; pub use tracing::*; pub use with_event_context::*; @@ -82,23 +82,41 @@ pub use with_event_context::*; /// Immutable context data that can be safely shared across thread boundaries. /// /// This struct holds key-value pairs of context information that gets attached -/// to events when they are persisted. It uses an immutable HashMap internally -/// for efficient cloning and thread-safe sharing of data snapshots. +/// to events when they are persisted. It uses a copy-on-write entry vector +/// internally: cloning is a single atomic refcount bump, and mutation only +/// clones the (tiny) entry vector when the data is currently shared. Context +/// maps hold a handful of entries in practice, so a linear-scan vector is +/// cheaper than a hashed or persistent map on every operation that matters +/// (clone, insert, lookup). /// /// `ContextData` is `Send` and can be passed between threads, unlike [`EventContext`] /// which is thread-local. This makes it suitable for transferring context across /// async boundaries via the [`WithEventContext`] trait. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(transparent)] -pub struct ContextData(im::HashMap, serde_json::Value>); +#[derive(Debug, Clone)] +pub struct ContextData(Arc, serde_json::Value)>>); impl ContextData { fn new() -> Self { - Self(im::HashMap::new()) + Self(Arc::new(Vec::new())) } fn insert(&mut self, key: &'static str, value: serde_json::Value) { - self.0 = self.0.update(Cow::Borrowed(key), value); + let entries = Arc::make_mut(&mut self.0); + if let Some((_, existing)) = entries.iter_mut().find(|(k, _)| *k == key) { + *existing = value; + } else { + entries.push((Cow::Borrowed(key), value)); + } + } + + /// Number of key-value pairs stored in this context. + pub fn len(&self) -> usize { + self.0.len() + } + + /// Returns `true` if the context holds no entries. + pub fn is_empty(&self) -> bool { + self.0.is_empty() } #[cfg(feature = "tracing-context")] @@ -115,16 +133,63 @@ impl ContextData { &self, key: &'static str, ) -> Result, serde_json::Error> { - let Some(val) = self.0.get(key) else { + let Some((_, val)) = self.0.iter().find(|(k, _)| *k == key) else { return Ok(None); }; serde_json::from_value(val.clone()).map(Some) } } +impl Serialize for ContextData { + fn serialize(&self, serializer: S) -> Result { + use serde::ser::SerializeMap; + let mut map = serializer.serialize_map(Some(self.0.len()))?; + for (key, value) in self.0.iter() { + map.serialize_entry(key, value)?; + } + map.end() + } +} + +impl<'de> Deserialize<'de> for ContextData { + fn deserialize>(deserializer: D) -> Result { + struct ContextDataVisitor; + + impl<'de> serde::de::Visitor<'de> for ContextDataVisitor { + type Value = ContextData; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("a map of context keys to JSON values") + } + + fn visit_map>( + self, + mut access: A, + ) -> Result { + let mut entries = Vec::with_capacity(access.size_hint().unwrap_or(0)); + while let Some((key, value)) = + access.next_entry::, serde_json::Value>()? + { + if let Some((_, existing)) = entries.iter_mut().find(|(k, _)| *k == key) { + *existing = value; + } else { + entries.push((key, value)); + } + } + Ok(ContextData(Arc::new(entries))) + } + } + + deserializer.deserialize_map(ContextDataVisitor) + } +} + struct StackEntry { id: Rc<()>, data: ContextData, + /// Set by [`EventContext::insert`]. Lets [`EventContext::data_if_dirty`] + /// skip the write-back clone when a poll left the context untouched. + dirty: bool, } thread_local! { @@ -216,6 +281,7 @@ impl EventContext { stack.push(StackEntry { id: id.clone(), data, + dirty: false, }); EventContext { id } @@ -248,6 +314,7 @@ impl EventContext { stack.push(StackEntry { id: id.clone(), data, + dirty: false, }); EventContext { id } @@ -318,6 +385,7 @@ impl EventContext { for entry in stack.iter_mut().rev() { if Rc::ptr_eq(&entry.id, &self.id) { entry.data.insert(key, json_value); + entry.dirty = true; return; } } @@ -356,6 +424,28 @@ impl EventContext { }) } + /// Returns a snapshot of the context data only if it was mutated since + /// the context was created (or since the last `data_if_dirty` call), + /// clearing the dirty flag. Returns `None` when untouched, letting + /// callers skip the clone entirely on the (overwhelmingly common) + /// read-only path. + pub(crate) fn data_if_dirty(&self) -> Option { + CONTEXT_STACK.with(|c| { + let mut stack = c.borrow_mut(); + for entry in stack.iter_mut().rev() { + if Rc::ptr_eq(&entry.id, &self.id) { + return if entry.dirty { + entry.dirty = false; + Some(entry.data.clone()) + } else { + None + }; + } + } + None + }) + } + #[allow(unused_mut)] pub(crate) fn data_for_storing() -> ContextData { let mut data = Self::current().data(); @@ -423,6 +513,69 @@ mod tests { ); } + #[test] + fn context_data_serializes_as_json_object_and_round_trips() { + let mut ctx = EventContext::current(); + ctx.insert("request_id", &"req-123").unwrap(); + ctx.insert("nested", &serde_json::json!({ "a": 1 })) + .unwrap(); + + let data = ctx.data(); + assert_eq!(data.len(), 2); + assert!(!data.is_empty()); + + let json = serde_json::to_value(&data).unwrap(); + assert!(json.is_object()); + assert_eq!( + json, + serde_json::json!({ "request_id": "req-123", "nested": { "a": 1 } }) + ); + + let round_tripped: ContextData = serde_json::from_value(json).unwrap(); + assert_eq!( + round_tripped.lookup::("request_id").unwrap(), + Some("req-123".to_string()) + ); + assert_eq!( + round_tripped.lookup::("nested").unwrap(), + Some(serde_json::json!({ "a": 1 })) + ); + assert_eq!(round_tripped.lookup::("missing").unwrap(), None); + } + + #[test] + fn context_data_insert_replaces_existing_key() { + let mut ctx = EventContext::current(); + ctx.insert("key", &"first").unwrap(); + ctx.insert("key", &"second").unwrap(); + let data = ctx.data(); + assert_eq!(data.len(), 1); + assert_eq!( + data.lookup::("key").unwrap(), + Some("second".to_string()) + ); + } + + #[test] + fn data_if_dirty_only_returns_data_after_mutation() { + let mut ctx = EventContext::current(); + ctx.insert("data", &"value").unwrap(); + + let seeded = EventContext::seed(ctx.data()); + assert!(seeded.data_if_dirty().is_none()); + + let mut inner = EventContext::current(); + inner.insert("inner", &"mutation").unwrap(); + + let dirty_data = seeded.data_if_dirty().expect("insert must mark dirty"); + assert_eq!( + dirty_data.lookup::("inner").unwrap(), + Some("mutation".to_string()) + ); + // Flag is cleared by the read + assert!(seeded.data_if_dirty().is_none()); + } + #[test] fn thread_isolation() { let mut ctx = EventContext::current(); diff --git a/src/context/with_event_context.rs b/src/context/with_event_context.rs index b1d0b4fc..81383b6c 100644 --- a/src/context/with_event_context.rs +++ b/src/context/with_event_context.rs @@ -84,7 +84,12 @@ impl Future for EventContextFuture { let this = self.project(); let ctx = EventContext::seed(this.context_data.clone()); let res = this.future.poll(cx); - *this.context_data = ctx.data(); + // Only write the context back when the poll actually mutated it — + // almost all polls are read-only, and the write-back costs a clone + // plus a stack walk. + if let Some(data) = ctx.data_if_dirty() { + *this.context_data = data; + } res } }