diff --git a/Cargo.lock b/Cargo.lock index 8e617aeb..9ce47768 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -26,6 +35,18 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.103" @@ -220,6 +241,12 @@ dependencies = [ "serde", ] +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.56" @@ -248,6 +275,58 @@ dependencies = [ "windows-link", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91e0c145792ef73a6ad36d27c75ac09f1832222a3c209689d90f534685ee5b7" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -302,6 +381,61 @@ version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "tokio", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.12" @@ -317,6 +451,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -507,6 +647,7 @@ dependencies = [ "async-trait", "base64", "chrono", + "criterion", "derive_builder", "es-entity-macros", "futures", @@ -747,6 +888,17 @@ dependencies = [ "wasip3", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -993,6 +1145,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.17" @@ -1198,6 +1359,12 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "opentelemetry" version = "0.32.0" @@ -1230,6 +1397,16 @@ dependencies = [ "tokio-stream", ] +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "parking" version = "2.2.1" @@ -1337,6 +1514,34 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "pluralizer" version = "0.5.0" @@ -1482,6 +1687,26 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1635,6 +1860,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schemars" version = "1.2.1" @@ -2155,6 +2389,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.10.0" @@ -2418,6 +2662,16 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" @@ -2527,6 +2781,16 @@ dependencies = [ "semver", ] +[[package]] +name = "web-sys" +version = "0.3.90" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "705eceb4ce901230f8625bd1d665128056ccbe4b7408faa625eec1ba80f59a97" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "web-time" version = "1.1.0" @@ -2565,6 +2829,37 @@ dependencies = [ "wasite", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-core" version = "0.62.2" diff --git a/Cargo.toml b/Cargo.toml index cbdc0256..fc130650 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,6 +57,11 @@ anyhow = { workspace = true } async-trait = { workspace = true } futures = { workspace = true } tracing-subscriber = { workspace = true } +criterion = { version = "0.8", features = ["html_reports", "async_tokio"] } + +[[bench]] +name = "context" +harness = false [workspace] resolver = "2" diff --git a/benches/context.rs b/benches/context.rs new file mode 100644 index 00000000..48552481 --- /dev/null +++ b/benches/context.rs @@ -0,0 +1,232 @@ +//! Micro-benchmarks for the event-context machinery. +//! +//! Structured after cala's `cala-perf` criterion setup. These benches only use +//! public API that exists both before and after PR #163 (`im::HashMap` vs +//! copy-on-write `Arc` inside `ContextData`), so the identical bench +//! commit can be cherry-picked onto either implementation for an A/B run. +//! +//! # Running +//! +//! ```text +//! cargo bench --bench context # core paths +//! cargo bench --bench context --features tracing-context # + per-persist tracing insert +//! ``` +//! +//! # A/B protocol against the `im`-based implementation +//! +//! ```text +//! git switch perf/event-context-benches-im # parent of PR #163 + this bench commit +//! cargo bench --bench context -- --save-baseline im +//! git switch perf/event-context-benches # PR #163 + this bench commit +//! cargo bench --bench context -- --baseline im +//! ``` +//! +//! # What each group isolates +//! +//! 1. `per_poll_overhead` — the PR's headline claim. A future that returns +//! `Pending` N times is driven to completion by a raw no-op-waker loop, +//! bare vs wrapped in `with_event_context`. (wrapped - bare) / (N + 1) +//! is the per-poll cost of the wrapper (seed + inner poll + write-back + +//! drop). +//! 2. `tokio_yield` — same comparison on a real tokio runtime, so the wrapper +//! cost can be read as a fraction of realistic scheduler overhead. +//! 3. `context_data_clone` — the clone that `EventContextFuture::poll` pays +//! on every poll (`im` map clone vs `Arc` refcount bump). +//! 4. `context_lifecycle` — `seed`+`drop` (TLS push/pop + `Rc` alloc, a cost +//! the PR does NOT remove) and `seed`+`insert`+`drop` (the copy-on-write / +//! HAMT path-copy mutation cost). +//! 5. `persist_path` — `EntityEvents::push` on an `event_context` event +//! (the real `data_for_storing()` call, one per persisted event; with +//! `tracing-context` enabled this includes the `"tracing"` insert) and +//! `serde_json::to_value` (the sqlx `Encode` serialization shape). + +use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use serde::{Deserialize, Serialize}; + +use std::{ + future::Future, + hint::black_box, + pin::{Pin, pin}, + task::{Context, Poll, Waker}, + time::Duration, +}; + +use es_entity::{ContextData, EntityEvents, EventContext, WithEventContext, *}; + +es_entity::entity_id! { BenchEntityId } + +#[derive(EsEvent, Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +#[es_event(id = "BenchEntityId", event_context)] +enum BenchEvent { + Happened { value: u64 }, +} + +const KEYS: [&str; 8] = [ + "key_0", "key_1", "key_2", "key_3", "key_4", "key_5", "key_6", "key_7", +]; + +/// Realistic context payloads: what `#[es_event_context]` inserts at request +/// setup (ids / small strings). +fn context_data(entries: usize) -> ContextData { + let mut ctx = EventContext::fork(); + for key in KEYS.iter().take(entries) { + ctx.insert(key, &"01996a2b-3c4d-5e6f-7a8b-9c0d1e2f3a4b") + .unwrap(); + } + ctx.data() +} + +/// Number of `Poll::Pending` returns before completion; a request that +/// suspends 64 times pays the wrapper 65 polls. +const POLLS: u32 = 64; + +struct YieldN(u32); + +impl Future for YieldN { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + if self.0 == 0 { + Poll::Ready(()) + } else { + self.0 -= 1; + Poll::Pending + } + } +} + +/// Drive a future to completion with a no-op waker — isolates poll cost from +/// any scheduler. +fn drive(fut: F) -> F::Output { + let mut fut = pin!(fut); + let mut cx = Context::from_waker(Waker::noop()); + loop { + if let Poll::Ready(out) = fut.as_mut().poll(&mut cx) { + return out; + } + } +} + +fn per_poll_overhead(c: &mut Criterion) { + let mut g = c.benchmark_group("1. per_poll_overhead"); + g.throughput(Throughput::Elements(POLLS as u64 + 1)); + + g.bench_function("bare_future", |b| { + b.iter(|| drive(black_box(YieldN(POLLS)))) + }); + + for entries in [0usize, 3] { + let data = context_data(entries); + g.bench_function( + BenchmarkId::new("with_event_context", format!("{entries}_entries")), + |b| b.iter(|| drive(black_box(YieldN(POLLS)).with_event_context(data.clone()))), + ); + } + g.finish(); +} + +async fn yield_many() { + for _ in 0..POLLS { + tokio::task::yield_now().await; + } +} + +fn tokio_yield(c: &mut Criterion) { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap(); + + let mut g = c.benchmark_group("2. tokio_yield"); + g.throughput(Throughput::Elements(POLLS as u64 + 1)); + + g.bench_function("bare_future", |b| b.to_async(&rt).iter(yield_many)); + + let data = context_data(3); + g.bench_function("with_event_context/3_entries", |b| { + b.to_async(&rt) + .iter(|| yield_many().with_event_context(data.clone())) + }); + g.finish(); +} + +fn context_data_clone(c: &mut Criterion) { + let mut g = c.benchmark_group("3. context_data_clone"); + for entries in [0usize, 3, 8] { + let data = context_data(entries); + g.bench_function( + BenchmarkId::from_parameter(format!("{entries}_entries")), + |b| b.iter(|| black_box(data.clone())), + ); + } + g.finish(); +} + +fn context_lifecycle(c: &mut Criterion) { + let mut g = c.benchmark_group("4. context_lifecycle"); + + // TLS push + Rc alloc + reverse-scan removal — per-poll cost the PR keeps. + let data = context_data(3); + g.bench_function("seed_drop/3_entries", |b| { + b.iter(|| { + let ctx = EventContext::seed(black_box(data.clone())); + drop(ctx); + }) + }); + + // Adds one insert into shared data: Arc::make_mut copy-on-write on the PR + // branch, HAMT path-copy (chunk clone) on the im branch. + g.bench_function("seed_insert_drop/3_entries", |b| { + b.iter(|| { + let mut ctx = EventContext::seed(black_box(data.clone())); + ctx.insert("inserted_key", &"inserted-value").unwrap(); + drop(ctx); + }) + }); + g.finish(); +} + +fn persist_path(c: &mut Criterion) { + let mut g = c.benchmark_group("5. persist_path"); + + // Keep an ambient context (as a request handler would) alive across the + // whole group so `data_for_storing()` clones realistic data. + let mut ambient = EventContext::current(); + for key in KEYS.iter().take(3) { + ambient + .insert(key, &"01996a2b-3c4d-5e6f-7a8b-9c0d1e2f3a4b") + .unwrap(); + } + + // One `data_for_storing()` per pushed event — ~10M calls in the stress run. + g.bench_function("entity_events_push/3_ambient_entries", |b| { + b.iter_batched( + || EntityEvents::::init(BenchEntityId::new(), std::iter::empty()), + |mut events| { + events.push(BenchEvent::Happened { value: 42 }); + events + }, + BatchSize::SmallInput, + ) + }); + + // The serialization the sqlx `Encode` path performs per stored context. + let data = context_data(3); + g.bench_function("serde_to_value/3_entries", |b| { + b.iter(|| serde_json::to_value(black_box(&data)).unwrap()) + }); + g.finish(); + + drop(ambient); +} + +criterion_group!( + name = benches; + config = Criterion::default() + .warm_up_time(Duration::from_millis(500)) + .measurement_time(Duration::from_secs(2)); + targets = per_poll_overhead, tokio_yield, context_data_clone, context_lifecycle, persist_path +); +criterion_main!(benches); diff --git a/src/context/mod.rs b/src/context/mod.rs index 38f0f7b7..a458e103 100644 --- a/src/context/mod.rs +++ b/src/context/mod.rs @@ -203,6 +203,7 @@ impl EventContext { /// // Context is now available for the current thread /// ``` pub fn current() -> Self { + with_event_context::materialize_pending_seeds(); CONTEXT_STACK.with(|c| { let mut stack = c.borrow_mut(); if let Some(last) = stack.last() { @@ -242,6 +243,14 @@ impl EventContext { /// // new_ctx now has its own independent context stack /// ``` pub fn seed(data: ContextData) -> Self { + with_event_context::materialize_pending_seeds(); + Self::push_entry(data) + } + + /// Pushes a new stack entry without materializing pending seeds first. + /// Only [`seed`](Self::seed) (after materializing) and + /// `with_event_context::materialize_pending_seeds` itself may call this. + fn push_entry(data: ContextData) -> Self { CONTEXT_STACK.with(|c| { let mut stack = c.borrow_mut(); let id = Rc::new(()); @@ -503,12 +512,17 @@ mod tests { let handle = tokio::spawn( async { - assert_eq!(stack_depth(), 2); + // Seeding is lazy: only the outer test context exists until + // the wrapped future observes the context. + assert_eq!(stack_depth(), 1); EventContext::current() .insert("spawned", &serde_json::json!("value")) .unwrap(); + // Observing the context materialized the wrapper's entry. + assert_eq!(stack_depth(), 2); + assert_eq!( current_json(), serde_json::json!({ "parent": "context", "spawned": "value" }) @@ -535,12 +549,17 @@ mod tests { let handle = tokio::spawn( async { - assert_eq!(stack_depth(), 1); + // Seeding is lazy: the worker thread's stack stays empty + // until the wrapped future observes the context. + assert_eq!(stack_depth(), 0); EventContext::current() .insert("spawned", &serde_json::json!("value")) .unwrap(); + // Observing the context materialized the wrapper's entry. + assert_eq!(stack_depth(), 1); + assert_eq!( current_json(), serde_json::json!({ "parent": "context", "spawned": "value" }) @@ -560,4 +579,77 @@ mod tests { assert_eq!(current_json(), serde_json::json!({ "parent": "context" })); } + + fn pending_depth() -> usize { + with_event_context::pending_seed_depth() + } + + #[tokio::test] + async fn with_event_context_untouched_poll_leaves_stacks_alone() { + let mut ctx = EventContext::current(); + ctx.insert("parent", &serde_json::json!("context")).unwrap(); + let before = current_json(); + + async { + // The wrapper parked a pending seed but no stack entry exists. + assert_eq!(stack_depth(), 1); + assert_eq!(pending_depth(), 1); + tokio::task::yield_now().await; + assert_eq!(stack_depth(), 1); + } + .with_event_context(ctx.data()) + .await; + + assert_eq!(current_json(), before); + assert_eq!(stack_depth(), 1); + assert_eq!(pending_depth(), 0); + } + + #[tokio::test] + async fn with_event_context_fork_materializes_wrapper_seed() { + let mut ctx = EventContext::current(); + ctx.insert("parent", &serde_json::json!("context")).unwrap(); + + async { + // fork() = current() + seed(): materializes the wrapper's entry + // and pushes the fork above it. + let mut forked = EventContext::fork(); + forked.insert("forked", &serde_json::json!("data")).unwrap(); + assert_eq!(stack_depth(), 3); + assert_eq!( + current_json(), + serde_json::json!({ "parent": "context", "forked": "data" }) + ); + drop(forked); + assert_eq!(stack_depth(), 2); + // Fork isolation: the wrapper's entry never saw "forked". + assert_eq!(current_json(), serde_json::json!({ "parent": "context" })); + } + .with_event_context(ctx.data()) + .await; + + assert_eq!(current_json(), serde_json::json!({ "parent": "context" })); + } + + #[tokio::test] + async fn with_event_context_panic_leaves_stacks_balanced() { + let ctx = EventContext::current(); + let data = ctx.data(); + + let res = tokio::spawn( + async { + // Materialize the wrapper's entry, then panic mid-poll. + let _ = EventContext::current(); + panic!("boom"); + } + .with_event_context(data), + ) + .await; + + assert!(res.is_err()); + // The guard's unwind path removed both the pending seed and the + // materialized entry (current-thread runtime: same thread). + assert_eq!(stack_depth(), 1); + assert_eq!(pending_depth(), 0); + } } diff --git a/src/context/with_event_context.rs b/src/context/with_event_context.rs index b1d0b4fc..b1de3b39 100644 --- a/src/context/with_event_context.rs +++ b/src/context/with_event_context.rs @@ -1,6 +1,7 @@ use pin_project::pin_project; use std::{ + cell::RefCell, future::Future, pin::Pin, task::{Context, Poll}, @@ -54,7 +55,7 @@ pub trait WithEventContext: Future { { EventContextFuture { future: self, - context_data, + context_data: Some(context_data), } } } @@ -70,11 +71,24 @@ impl WithEventContext for F {} /// The future maintains context isolation - the context is only active /// during the polling of the wrapped future and does not leak to other /// concurrent operations. +/// +/// Seeding is lazy: on each poll the context data is parked on a +/// thread-local pending-seed stack, and a real [`EventContext`] entry is +/// only materialized if the inner future actually observes the context +/// (via [`EventContext::current`] or [`EventContext::seed`]) during that +/// poll. Polls that never touch the context — the overwhelming majority — +/// pay neither an allocation nor a clone. +/// +/// [`EventContext`]: super::EventContext +/// [`EventContext::current`]: super::EventContext::current +/// [`EventContext::seed`]: super::EventContext::seed #[pin_project] pub struct EventContextFuture { #[pin] future: F, - context_data: ContextData, + /// `Some` between polls; taken for the duration of each poll while the + /// data is parked on the pending-seed stack. + context_data: Option, } impl Future for EventContextFuture { @@ -82,9 +96,106 @@ impl Future for EventContextFuture { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); - let ctx = EventContext::seed(this.context_data.clone()); + let data = this + .context_data + .take() + .expect("EventContextFuture must not be polled after a panic"); + let guard = SeedGuard::enter(data); let res = this.future.poll(cx); - *this.context_data = ctx.data(); + *this.context_data = Some(guard.finish()); res } } + +/// A context seed deferred by [`EventContextFuture`]: the wrapper's +/// [`ContextData`] is parked here on poll entry instead of eagerly becoming +/// a context-stack entry. A real entry is only materialized — see +/// [`materialize_pending_seeds`] — if the inner future actually observes the +/// context during that poll. Most polls of a request/job future never do, so +/// the common case skips the entry allocation, the push/pop, and the clone +/// entirely. +struct PendingSeed { + /// `Some` until materialized; then moved into the stack entry. + data: Option, + /// Handle to the materialized entry, if any. Dropping it removes the + /// entry via the normal [`EventContext`] drop logic. + ctx: Option, +} + +thread_local! { + static PENDING_SEEDS: RefCell> = const { RefCell::new(Vec::new()) }; +} + +/// Materializes every pending seed into a real context-stack entry, in push +/// order. +/// +/// [`EventContext`] calls this before any operation that observes or pushes +/// onto the stack ([`EventContext::current`] / [`EventContext::seed`]). This +/// preserves the invariant that pending seeds always sit logically *above* +/// every stack entry that existed when they were parked: any new entry is +/// pushed after the seeds have been materialized beneath it. +pub(super) fn materialize_pending_seeds() { + PENDING_SEEDS.with(|p| { + let mut pending = p.borrow_mut(); + for seed in pending.iter_mut() { + if seed.ctx.is_none() { + let data = seed.data.take().expect("unmaterialized seed retains data"); + seed.ctx = Some(EventContext::push_entry(data)); + } + } + }) +} + +/// RAII token for one [`EventContextFuture`] poll invocation. +/// +/// [`enter`](Self::enter) parks the wrapper's data as a pending seed; +/// [`finish`](Self::finish) (normal exit) pops it and returns the data to +/// store back into the wrapper — either untouched (fast path: the poll never +/// observed the context) or harvested from the materialized entry. If the +/// inner poll panics, `Drop` pops the record and discards any materialized +/// handle so the thread-local stacks stay balanced. +#[must_use] +struct SeedGuard(()); + +impl SeedGuard { + fn enter(data: ContextData) -> Self { + PENDING_SEEDS.with(|p| { + p.borrow_mut().push(PendingSeed { + data: Some(data), + ctx: None, + }) + }); + SeedGuard(()) + } + + fn finish(self) -> ContextData { + let seed = PENDING_SEEDS + .with(|p| p.borrow_mut().pop()) + .expect("SeedGuard::finish: pending-seed stack is empty"); + std::mem::forget(self); + match seed { + // Materialized: harvest the (possibly mutated) data. Dropping the + // handle removes the entry unless the inner future kept its own + // handle across the poll — same lifecycle as an eager seed. + PendingSeed { ctx: Some(ctx), .. } => ctx.data(), + // Fast path: the poll never observed the context; hand the data + // back unchanged without ever having touched the context stack. + PendingSeed { data, .. } => data.expect("unmaterialized seed retains data"), + } + } +} + +impl Drop for SeedGuard { + fn drop(&mut self) { + // Unwind path: the wrapped poll panicked. Pop the record; dropping a + // materialized handle removes its stack entry via EventContext::drop. + // Write-back is skipped, matching the eager-seeding behavior. + let _ = PENDING_SEEDS.with(|p| p.borrow_mut().pop()); + } +} + +/// Test-only visibility into the pending-seed stack for the context tests. +#[cfg(test)] +pub(super) fn pending_seed_depth() -> usize { + PENDING_SEEDS.with(|p| p.borrow().len()) +}