perf(context): eliminate per-poll context clone overhead

[nicolasburtey]

Motivation

This came out of a 5-hour DB-bottleneck stress test of lana-bank
(sandbox sb-str3lps, stress-testing preset at a 3 loans/sec target,
lana-bank @ 314fbd8, es-entity pinned at 0.11.7). The headline numbers
from that run:

• The server was CPU-bound, not DB-bound: 5.2 cores average / 6.6
  max (Honeycomb kubelet stats) the entire run, while the database sat
  largely idle behind advisory-lock waits.
• Achieved throughput plateaued at ~2.1-2.4 loans/sec (~75% of target).
  The direct cap is cala's per-balance pg_advisory_xact_lock (mean
  wait ~700ms), but the lock holds are stretched app-side: holders
  sit idle in transaction between statements because the app has no
  spare CPU. Framework CPU overhead therefore converts directly into
  lock queueing and lost throughput — every percent of wasted app CPU
  is a percent of pipeline throughput.

The 60s /debug/pprof CPU capture showed no single hot function; the
burn is framework-level per-operation overhead, and es-entity's event
context is a measurable part of it:

• <EventContextFuture as Future>::poll is an ancestor of 51% of
  all samples (it wraps every request/job root future).
• std::thread::local::LocalKey::try_with is on the stack of 63%
  of samples.
• sized_chunks::Chunk::clone (the im persistent-map internals
  behind ContextData::clone) shows up as its own profile row.
• Flat self-time leaders are allocator traffic (malloc/free ~15%)
  and atomic refcount ops (~2%) — exactly the cost shape of the
  clone-heavy context machinery.

The problem, mechanically

EventContextFuture::poll runs on every poll of every request/job
future — futures are re-polled at every .await resumption, so a
request that suspends 50 times pays this 50 times. Each poll did:

1. context_data.clone() — ContextData was
   im::HashMap<Cow<'static, str>, serde_json::Value>. Cloning an
   im map walks the persistent chunk tree bumping Arc refcounts.
2. EventContext::seed(clone) — TLS push + Rc allocation.
3. ctx.data() — a second TLS walk and a second full clone,
   unconditionally written back into the wrapper.

Plus EventContext::drop: third TLS access, reverse stack scan, and
a Vec::remove. All of this to support a write-back that is only
meaningful when the inner future mutated the context — which in
practice happens once at request setup (the #[es_event_context]
macro's inserts), i.e. in a vanishing fraction of polls.

Separately, every event persist calls data_for_storing(), which also
pays the im clone — the stress run wrote on the order of 10M event
rows.

The fix

Two contained changes, no public API or behavior change:

1. ContextData is now copy-on-write Arc<Vec<(Cow, Value)>>.
   Clone = one atomic refcount bump. Mutation clones the entry vector
   only when shared (Arc::make_mut). Context maps hold 0-3 entries
   in practice, so a linear-scan vector beats both hashing and a
   persistent tree on every operation this type actually performs.
   Custom Serialize/Deserialize impls preserve the exact stored
   JSON shape (a plain object — what #[serde(transparent)] over the
   map produced), so existing context columns and event payloads
   round-trip identically. This also makes data(), fork() and
   data_for_storing() cheap everywhere, not just in the poll path.
2. Write-back is skipped unless the poll mutated the context.
   StackEntry gains a dirty flag set by EventContext::insert;
   the wrapper only clones data back out (new crate-private
   data_if_dirty()) when an insert actually happened during that
   poll. Semantics are unchanged: unchanged context propagates
   unchanged.

Drops the im dependency (it was used nowhere else).

Preserved invariants: the per-poll re-seed itself stays — on
tokio's multi-thread scheduler a task's polls hop worker threads and
TLS doesn't travel with the task, so seeding at each poll is what
guarantees EventContext::current() inside the future resolves
correctly. Only the redundant clones around it are gone.

Expected effect

Direct profile attribution for this machinery is ~2-5% of app CPU
(chunk-tree clones, Rc traffic, TLS walks, plus the per-persist clone
across ~10M event writes) — modest in isolation, but on this workload
it sits directly upstream of the throughput-capping lock holds, and it
compounds with the other allocation-rate reducers (tracing span
volume, allocator swap) rather than overlapping them.

Validation plan

• Re-run the same stress-testing sandbox at 3 loans/sec with a
  lana-bank build on this es-entity version; compare Honeycomb
  container.cpu.usage AVG, cala find_for_update mean advisory-lock
  wait, and achieved loans/sec against the baseline from the stress
  run (5.2 cores avg / ~700ms lock wait / ~2.2 loans/s).

Testing

• 3 new unit tests: serde JSON-object shape + round-trip, insert key
  replacement, dirty-flag write-back behavior.
• All 37 lib tests + 10 context doctests pass;
  cargo clippy --workspace --all-features and cargo fmt --check
  are clean.
• DB-backed integration tests were not run locally (no postgres in
  this environment); the change touches no SQL, and ContextData's
  sqlx Encode/Decode go through the preserved serde shape.

---

Note

Medium Risk
Touches thread-local context propagation on every async poll and persisted event context JSON shape; behavior is intended to be unchanged but regressions would affect audit metadata and cross-task context.

Overview
Reduces hot-path CPU and allocation cost in event context propagation by changing how ContextData is stored and when async wrappers write context back after each poll.

ContextData no longer uses im::HashMap (and the im crate is removed). It is now Arc<Vec<(Cow, Value)>> with copy-on-write on insert, so clone() is a cheap refcount bump instead of walking persistent map chunks. Custom Serialize/Deserialize keep the stored JSON as a plain object for DB/event payloads. len / is_empty are added on ContextData.

EventContextFuture::poll still re-seeds TLS each poll, but write-back uses new data_if_dirty() instead of unconditional ctx.data(). StackEntry tracks a dirty flag set on insert, so read-only polls skip cloning context back into the wrapper.

Unit tests cover serde round-trip, key replacement, and dirty write-back behavior.

^{Reviewed by Cursor Bugbot for commit 81527d2. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bodymindarts]

I’d want to see local profiling as proof that this is better.

The im HashMap should already be optimized for cloning due to structural sharing.

[bodymindarts]

I'll take this over

[bodymindarts]

Benchmarking this change in isolation (criterion suite in #164, cherry-pickable onto either branch) pointed us to a different fix, now up as #164: an im::HashMap clone is a flat ~3.4ns regardless of size — structural sharing already makes it O(1) — so of the wrapper's ~21ns/poll, ~18ns is the seed/drop machinery itself (three thread-local accesses, a per-poll Rc::new(()) heap allocation, and a stack push/pop), which is independent of the container and survives this PR unchanged (this PR lands at ~18.6ns/poll, −13%). #164 instead makes seeding lazy — the wrapper parks its ContextData on a thread-local pending-seed stack (a move: no clone, no allocation) and only materializes a real stack entry if the inner future actually observes the context during that poll, which almost none do — landing at 12.6ns/poll (−41%), ~5.8ns/poll added in a real tokio harness (−54% on the wrapped task), with zero heap allocations and zero clones on untouched polls; that also addresses the two signals actually visible in the stress-run profile (malloc/free traffic and LocalKey::try_with), which the container swap does not. It achieves this while keeping im and without touching ContextData at all — no dependency churn and no change to the persisted context JSONB serialization path, so nothing to re-validate there. The remaining case for this PR is the persist/serialize-path improvement (−29% on EntityEvents::push with tracing-context) and dropping four transitive deps — worth weighing on hygiene grounds, but that path is not hot either: the insert/serialize work only runs when an entity actually mutates (one data_for_storing() + Encode per persisted event, nothing on reads or queries), and mutations are rare — even the DB-bottleneck stress run peaked at ~555 events/sec, so the −29% amounts to ~0.5s of CPU per 10M events and should be argued as cleanup rather than throughput recovery; note the Encode shape issue flagged separately would need fixing first either way.

🤖 Generated with Claude Code