Skip to main content

keyhog_profile/
runtime.rs

1use crate::schema::Stage;
2use crate::schema_v2::{
3    AnnotationV2, BatchRouteV2, BlockedWaitRecordV2, CacheEffectivenessV2, DistributionBucketV2,
4    Evidence, EvidenceGap, IndexedCounterRecordV2, LatencyBucketV2, LatencyDistributionV2,
5    MetricDistributionV2, PointEventV2, QueueDepthV2, QueueLinkV2, RetryRecordV2, SpanRecordV2,
6    StageConcurrencyV2, TypedMetricRecordV2, WorkOrigin, WorkerImbalanceV2, WorkerLoadV2,
7    WorkerOccupancyRowV2, WorkerOccupancyV2,
8};
9use std::cell::{Cell, RefCell};
10use std::collections::HashMap;
11use std::future::{poll_fn, Future};
12use std::marker::PhantomData;
13use std::rc::Rc;
14use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
15use std::sync::{Arc, Mutex, Weak};
16use std::time::Instant;
17
18pub(crate) const STAGE_COUNT: usize = Stage::ALL.len();
19
20/// Maximum number of causal span records retained by one profiling runtime.
21pub const MAX_RECORDED_SPANS: usize = 65_536;
22const MAX_NESTED_SPANS: usize = 64;
23const LATENCY_BUCKET_COUNT: usize = 65;
24pub const MAX_POINT_EVENTS: usize = 16_384;
25pub const MAX_ANNOTATIONS: usize = 16_384;
26/// Maximum pending enqueues and completed links retained per runtime.
27pub const MAX_QUEUE_LINKS: usize = 16_384;
28/// Hard cap on retained batch-route records; further routes count as drops.
29pub const MAX_BATCH_ROUTES: usize = 16_384;
30
31/// Exact reasons queue causality records were not retained.
32#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
33pub struct QueueLinkLossCounts {
34    /// Pending enqueues dropped at capacity or displaced by a duplicate sequence.
35    pub dropped_enqueues: u64,
36    /// Completed links dropped at capacity.
37    pub dropped_links: u64,
38    /// Dequeues with no recorded matching enqueue.
39    pub unmatched_dequeues: u64,
40    /// Pending enqueues never matched by a dequeue before the drain.
41    pub unconsumed_enqueues: u64,
42}
43
44/// Portable causal parent captured from a runtime's current span context.
45///
46/// The token is plain data: pass it across crate, thread, or spawn boundaries
47/// and attach it with [`span_with_parent`] or [`instrument_future_with_parent`]
48/// where thread-local propagation cannot reach. A token only applies inside
49/// the runtime it was captured from; elsewhere the span records as a root.
50#[derive(Clone, Copy, Debug, Eq, PartialEq)]
51pub struct CausalParent {
52    context_id: u64,
53    span_id: u64,
54}
55
56impl CausalParent {
57    /// Process-local identity of the runtime this token was captured from.
58    pub const fn context_id(self) -> u64 {
59        self.context_id
60    }
61
62    /// Parent span identifier, or zero when captured outside any span.
63    pub const fn span_id(self) -> u64 {
64        self.span_id
65    }
66
67    /// Whether this token names the runtime root instead of a live span.
68    pub const fn is_root(self) -> bool {
69        self.span_id == 0
70    }
71}
72
73/// Deterministic bounded policy for retaining expensive detail events.
74#[derive(Clone, Copy, Debug, Eq, PartialEq)]
75pub struct SamplingPolicy {
76    initial_events: u64,
77    every_nth_after_initial: u64,
78    maximum_retained: u64,
79}
80
81impl SamplingPolicy {
82    /// Retain `initial_events`, then every Nth observation, up to an absolute bound.
83    pub const fn bounded(
84        initial_events: u64,
85        every_nth_after_initial: u64,
86        maximum_retained: u64,
87    ) -> Self {
88        Self {
89            initial_events,
90            every_nth_after_initial: if every_nth_after_initial == 0 {
91                1
92            } else {
93                every_nth_after_initial
94            },
95            maximum_retained,
96        }
97    }
98
99    fn selects(self, observation: u64) -> bool {
100        observation < self.initial_events
101            || (observation - self.initial_events) % self.every_nth_after_initial == 0
102    }
103}
104
105/// Exact reasons typed timeline records were not retained.
106#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
107pub struct EventLossCounts {
108    pub point_events: u64,
109    pub annotations: u64,
110    pub sampled_out_events: u64,
111}
112
113impl EventLossCounts {
114    /// Total capacity loss. Policy sampling is reported separately.
115    pub const fn capacity_drops(self) -> u64 {
116        self.point_events.saturating_add(self.annotations)
117    }
118}
119
120#[derive(Clone, Copy)]
121struct ActiveSpan {
122    runtime_key: usize,
123    span_id: u64,
124}
125
126#[derive(Clone, Copy)]
127struct SpanTrace {
128    record_index: usize,
129    span_id: u64,
130    stack_slot: Option<usize>,
131}
132
133/// One completed span's recording payload, assembled on the guard's drop path.
134#[derive(Clone, Copy)]
135struct SpanOutcome {
136    start_offset_ns: u64,
137    elapsed_ns: u64,
138    attributed: bool,
139    blocked: bool,
140    serial: bool,
141    outermost: bool,
142}
143
144struct RawSpanRecord {
145    span_id: u64,
146    parent_span_id: u64,
147    metric_id: crate::MetricId,
148    start_ns: u64,
149    inclusive_ns: u64,
150    thread_id: u64,
151    worker_id: u64,
152    task_id: u64,
153    work_origin: WorkOrigin,
154    completed: bool,
155    hardware: RawSpanHardware,
156}
157
158/// Raw counter readings captured at span edges; the cold path joins them.
159#[derive(Clone, Copy, Default)]
160struct RawSpanHardware {
161    cycles_begin: u64,
162    cycles_end: u64,
163    instructions_begin: u64,
164    instructions_end: u64,
165    has_cycles: bool,
166    has_instructions: bool,
167    finished: bool,
168}
169
170impl RawSpanHardware {
171    fn begin() -> Self {
172        let mut hardware = Self::default();
173        if let Some(reading) = crate::hardware::span_counter_reading() {
174            if let Some(cycles) = reading.cycles {
175                hardware.cycles_begin = cycles;
176                hardware.has_cycles = true;
177            }
178            if let Some(instructions) = reading.instructions {
179                hardware.instructions_begin = instructions;
180                hardware.has_instructions = true;
181            }
182        }
183        hardware
184    }
185
186    fn finish(&mut self) {
187        if let Some(reading) = crate::hardware::span_counter_reading() {
188            if self.has_cycles {
189                if let Some(cycles) = reading.cycles {
190                    self.cycles_end = cycles;
191                }
192            }
193            if self.has_instructions {
194                if let Some(instructions) = reading.instructions {
195                    self.instructions_end = instructions;
196                }
197            }
198        }
199        self.finished = true;
200    }
201
202    fn into_evidence(self) -> Evidence<crate::hardware::SpanHardwareV2> {
203        if !self.finished || (!self.has_cycles && !self.has_instructions) {
204            return Evidence::unavailable(EvidenceGap::Unavailable);
205        }
206        let pair = |present: bool, begin: u64, end: u64| {
207            if present {
208                (Evidence::recorded(begin), Evidence::recorded(end))
209            } else {
210                (
211                    Evidence::unavailable(EvidenceGap::Unsupported),
212                    Evidence::unavailable(EvidenceGap::Unsupported),
213                )
214            }
215        };
216        let (cycles_begin, cycles_end) = pair(self.has_cycles, self.cycles_begin, self.cycles_end);
217        let (instructions_begin, instructions_end) = pair(
218            self.has_instructions,
219            self.instructions_begin,
220            self.instructions_end,
221        );
222        Evidence::recorded(crate::hardware::SpanHardwareV2 {
223            version: crate::hardware::SPAN_HARDWARE_V2_VERSION,
224            cycles_begin,
225            cycles_end,
226            instructions_begin,
227            instructions_end,
228        })
229    }
230}
231
232struct PendingQueueEnqueue {
233    thread_id: u64,
234    elapsed_ns: u64,
235}
236
237struct WorkerShard {
238    sequence: u64,
239    elapsed_ns: [AtomicU64; STAGE_COUNT],
240    calls: [AtomicU64; STAGE_COUNT],
241    attributed_ns: [AtomicU64; STAGE_COUNT],
242    blocked_ns: [AtomicU64; STAGE_COUNT],
243    blocked_calls: [AtomicU64; STAGE_COUNT],
244    legacy_elapsed_ns: [AtomicU64; STAGE_COUNT],
245    legacy_calls: [AtomicU64; STAGE_COUNT],
246    legacy_attributed_ns: [AtomicU64; STAGE_COUNT],
247    latency_buckets: [[AtomicU64; LATENCY_BUCKET_COUNT]; STAGE_COUNT],
248    latency_min_ns: [AtomicU64; STAGE_COUNT],
249    latency_max_ns: [AtomicU64; STAGE_COUNT],
250    counter_values: [AtomicU64; crate::MetricId::COUNT],
251    input_bytes: AtomicU64,
252    input_units: AtomicU64,
253    legacy_input_bytes: AtomicU64,
254    legacy_input_units: AtomicU64,
255    derived_decoder_bytes: AtomicU64,
256    backend_dispatched_bytes: AtomicU64,
257    stage_first_start_ns: [AtomicU64; STAGE_COUNT],
258    stage_last_end_ns: [AtomicU64; STAGE_COUNT],
259    stage_bytes: [AtomicU64; STAGE_COUNT],
260    serial_ns: [AtomicU64; STAGE_COUNT],
261    serial_calls: [AtomicU64; STAGE_COUNT],
262    top_level_busy_ns: AtomicU64,
263    top_level_blocked_ns: AtomicU64,
264    top_level_calls: AtomicU64,
265    cache_hits: [AtomicU64; crate::CacheId::COUNT],
266    cache_misses: [AtomicU64; crate::CacheId::COUNT],
267    indexed_counters: [[AtomicU64; crate::INDEXED_COUNTER_SLOTS]; crate::IndexedCounterId::COUNT],
268    indexed_counter_dropped: AtomicU64,
269    retries: [AtomicU64; crate::RetryCause::COUNT],
270}
271
272const fn zero_counters() -> [AtomicU64; STAGE_COUNT] {
273    [const { AtomicU64::new(0) }; STAGE_COUNT]
274}
275
276const fn zero_indexed_counters(
277) -> [[AtomicU64; crate::INDEXED_COUNTER_SLOTS]; crate::IndexedCounterId::COUNT] {
278    [const { [const { AtomicU64::new(0) }; crate::INDEXED_COUNTER_SLOTS] };
279        crate::IndexedCounterId::COUNT]
280}
281
282const fn zero_event_values() -> [AtomicU64; crate::EventId::COUNT] {
283    [const { AtomicU64::new(0) }; crate::EventId::COUNT]
284}
285
286const fn zero_metric_values() -> [AtomicU64; crate::MetricId::COUNT] {
287    [const { AtomicU64::new(0) }; crate::MetricId::COUNT]
288}
289
290const fn zero_latency_buckets() -> [[AtomicU64; LATENCY_BUCKET_COUNT]; STAGE_COUNT] {
291    [const { [const { AtomicU64::new(0) }; LATENCY_BUCKET_COUNT] }; STAGE_COUNT]
292}
293
294const fn max_counters() -> [AtomicU64; STAGE_COUNT] {
295    [const { AtomicU64::new(u64::MAX) }; STAGE_COUNT]
296}
297
298const fn zero_cache_counters() -> [AtomicU64; crate::CacheId::COUNT] {
299    [const { AtomicU64::new(0) }; crate::CacheId::COUNT]
300}
301
302impl WorkerShard {
303    fn new(sequence: u64) -> Self {
304        Self {
305            sequence,
306            elapsed_ns: zero_counters(),
307            calls: zero_counters(),
308            attributed_ns: zero_counters(),
309            blocked_ns: zero_counters(),
310            blocked_calls: zero_counters(),
311            legacy_elapsed_ns: zero_counters(),
312            legacy_calls: zero_counters(),
313            legacy_attributed_ns: zero_counters(),
314            latency_buckets: zero_latency_buckets(),
315            latency_min_ns: max_counters(),
316            latency_max_ns: zero_counters(),
317            counter_values: zero_metric_values(),
318            input_bytes: AtomicU64::new(0),
319            input_units: AtomicU64::new(0),
320            legacy_input_bytes: AtomicU64::new(0),
321            legacy_input_units: AtomicU64::new(0),
322            derived_decoder_bytes: AtomicU64::new(0),
323            backend_dispatched_bytes: AtomicU64::new(0),
324            stage_first_start_ns: max_counters(),
325            stage_last_end_ns: zero_counters(),
326            stage_bytes: zero_counters(),
327            serial_ns: zero_counters(),
328            serial_calls: zero_counters(),
329            top_level_busy_ns: AtomicU64::new(0),
330            top_level_blocked_ns: AtomicU64::new(0),
331            top_level_calls: AtomicU64::new(0),
332            cache_hits: zero_cache_counters(),
333            cache_misses: zero_cache_counters(),
334            indexed_counters: zero_indexed_counters(),
335            indexed_counter_dropped: AtomicU64::new(0),
336            retries: [const { AtomicU64::new(0) }; crate::RetryCause::COUNT],
337        }
338    }
339}
340
341fn percentile_upper_bound(
342    buckets: &[LatencyBucketV2],
343    call_count: u64,
344    percentile: u64,
345    maximum_ns: u64,
346) -> u64 {
347    let rank = (u128::from(call_count) * u128::from(percentile)).div_ceil(100);
348    let mut cumulative = 0_u128;
349    for bucket in buckets {
350        cumulative += u128::from(bucket.count);
351        if cumulative >= rank {
352            return bucket.upper_bound_ns.min(maximum_ns);
353        }
354    }
355    maximum_ns
356}
357
358#[inline]
359fn latency_bucket_index(elapsed_ns: u64) -> usize {
360    if elapsed_ns == 0 {
361        0
362    } else {
363        (u64::BITS - elapsed_ns.leading_zeros()) as usize
364    }
365}
366
367fn latency_bucket_bounds(index: usize) -> (u64, u64) {
368    match index {
369        0 => (0, 0),
370        64 => (1_u64 << 63, u64::MAX),
371        _ => (1_u64 << (index - 1), (1_u64 << index) - 1),
372    }
373}
374
375pub(crate) struct RawStageCounters {
376    pub(crate) elapsed_ns: [u64; STAGE_COUNT],
377    pub(crate) calls: [u64; STAGE_COUNT],
378    pub(crate) attributed_ns: [u64; STAGE_COUNT],
379}
380
381static ACTIVE_CONTEXTS: AtomicUsize = AtomicUsize::new(0);
382static NEXT_THREAD_ID: AtomicU64 = AtomicU64::new(1);
383static NEXT_CONTEXT_ID: AtomicU64 = AtomicU64::new(1);
384
385struct RuntimeInner {
386    context_id: u64,
387    next_shard_sequence: AtomicU64,
388    elapsed_ns: [AtomicU64; STAGE_COUNT],
389    calls: [AtomicU64; STAGE_COUNT],
390    attributed_ns: [AtomicU64; STAGE_COUNT],
391    session_shards: Mutex<Vec<Arc<WorkerShard>>>,
392    session_gauge_values: [AtomicU64; crate::MetricId::COUNT],
393    session_gauge_present: [AtomicU64; 2],
394    input_bytes: AtomicU64,
395    input_units: AtomicU64,
396    session_recording: bool,
397    session_route_sequence: AtomicU64,
398    session_batch_routes: Mutex<Vec<BatchRouteV2>>,
399    session_dropped_batch_routes: AtomicU64,
400    started: Instant,
401    session_span_sequence: AtomicU64,
402    session_span_reservations: AtomicUsize,
403    session_spans: Mutex<Vec<RawSpanRecord>>,
404    session_dropped_spans: AtomicU64,
405    session_event_sequence: AtomicU64,
406    session_point_events: Mutex<Vec<PointEventV2>>,
407    session_annotations: Mutex<Vec<AnnotationV2>>,
408    session_dropped_point_events: AtomicU64,
409    session_dropped_annotations: AtomicU64,
410    session_sample_observations: [AtomicU64; crate::EventId::COUNT],
411    session_sample_retained: [AtomicU64; crate::EventId::COUNT],
412    session_sampled_out_events: AtomicU64,
413    queue_pending: Mutex<HashMap<(u8, u64), PendingQueueEnqueue>>,
414    queue_links: Mutex<Vec<QueueLinkV2>>,
415    queue_dropped_enqueues: AtomicU64,
416    queue_dropped_links: AtomicU64,
417    queue_unmatched_dequeues: AtomicU64,
418    queue_depth_current: [AtomicU64; crate::QueueId::COUNT],
419    queue_depth_high_water: [AtomicU64; crate::QueueId::COUNT],
420    queue_depth_enqueues: [AtomicU64; crate::QueueId::COUNT],
421    queue_depth_dequeues: [AtomicU64; crate::QueueId::COUNT],
422    legacy_typed_counters: [AtomicU64; crate::MetricId::COUNT],
423    distribution_buckets: [[AtomicU64; LATENCY_BUCKET_COUNT]; crate::MetricId::COUNT],
424    distribution_min: [AtomicU64; crate::MetricId::COUNT],
425    distribution_max: [AtomicU64; crate::MetricId::COUNT],
426}
427
428const fn zero_queue_depths() -> [AtomicU64; crate::QueueId::COUNT] {
429    [const { AtomicU64::new(0) }; crate::QueueId::COUNT]
430}
431
432const fn zero_distribution_buckets() -> [[AtomicU64; LATENCY_BUCKET_COUNT]; crate::MetricId::COUNT]
433{
434    [const { [const { AtomicU64::new(0) }; LATENCY_BUCKET_COUNT] }; crate::MetricId::COUNT]
435}
436
437const fn zero_distribution_mins() -> [AtomicU64; crate::MetricId::COUNT] {
438    [const { AtomicU64::new(u64::MAX) }; crate::MetricId::COUNT]
439}
440
441const fn zero_distribution_maxes() -> [AtomicU64; crate::MetricId::COUNT] {
442    [const { AtomicU64::new(0) }; crate::MetricId::COUNT]
443}
444
445impl RuntimeInner {
446    fn new(session_recording: bool, started: Instant) -> Self {
447        Self {
448            context_id: NEXT_CONTEXT_ID.fetch_add(1, Ordering::Relaxed),
449            next_shard_sequence: AtomicU64::new(1),
450            elapsed_ns: zero_counters(),
451            calls: zero_counters(),
452            attributed_ns: zero_counters(),
453            session_shards: Mutex::new(Vec::new()),
454            session_gauge_values: zero_metric_values(),
455            session_gauge_present: [const { AtomicU64::new(0) }; 2],
456            input_bytes: AtomicU64::new(0),
457            input_units: AtomicU64::new(0),
458            session_recording,
459            session_route_sequence: AtomicU64::new(0),
460            session_batch_routes: Mutex::new(Vec::new()),
461            session_dropped_batch_routes: AtomicU64::new(0),
462            started,
463            session_span_sequence: AtomicU64::new(0),
464            session_span_reservations: AtomicUsize::new(0),
465            session_spans: Mutex::new(Vec::new()),
466            session_dropped_spans: AtomicU64::new(0),
467            session_event_sequence: AtomicU64::new(0),
468            session_point_events: Mutex::new(Vec::new()),
469            session_annotations: Mutex::new(Vec::new()),
470            session_dropped_point_events: AtomicU64::new(0),
471            session_dropped_annotations: AtomicU64::new(0),
472            session_sample_observations: zero_event_values(),
473            session_sample_retained: zero_event_values(),
474            session_sampled_out_events: AtomicU64::new(0),
475            queue_pending: Mutex::new(HashMap::new()),
476            queue_links: Mutex::new(Vec::new()),
477            queue_dropped_enqueues: AtomicU64::new(0),
478            queue_dropped_links: AtomicU64::new(0),
479            queue_unmatched_dequeues: AtomicU64::new(0),
480            queue_depth_current: zero_queue_depths(),
481            queue_depth_high_water: zero_queue_depths(),
482            queue_depth_enqueues: zero_queue_depths(),
483            queue_depth_dequeues: zero_queue_depths(),
484            legacy_typed_counters: zero_metric_values(),
485            distribution_buckets: zero_distribution_buckets(),
486            distribution_min: zero_distribution_mins(),
487            distribution_max: zero_distribution_maxes(),
488        }
489    }
490
491    fn sorted_shards(&self) -> Vec<Arc<WorkerShard>> {
492        let shards = match self.session_shards.lock() {
493            Ok(shards) => shards,
494            Err(poisoned) => poisoned.into_inner(),
495        };
496        let mut sorted: Vec<Arc<WorkerShard>> = shards.to_vec();
497        sorted.sort_by_key(|shard| shard.sequence);
498        sorted
499    }
500}
501
502struct ThreadShardAssignment {
503    runtime: Weak<RuntimeInner>,
504    shard: Arc<WorkerShard>,
505}
506
507/// Owned fixed-stage metric storage that can be propagated across worker boundaries.
508#[derive(Clone)]
509pub struct Runtime {
510    inner: Arc<RuntimeInner>,
511}
512
513impl Runtime {
514    /// Create an isolated runtime whose measurements can back one profiling session.
515    pub fn new() -> Self {
516        Self::new_at(Instant::now())
517    }
518
519    /// Process-local identity that distinguishes this runtime from every peer.
520    pub fn context_id(&self) -> u64 {
521        self.inner.context_id
522    }
523
524    pub(crate) fn new_at(started: Instant) -> Self {
525        Self {
526            inner: Arc::new(RuntimeInner::new(true, started)),
527        }
528    }
529
530    fn legacy() -> Self {
531        Self {
532            inner: Arc::new(RuntimeInner::new(false, Instant::now())),
533        }
534    }
535
536    fn worker_shard(&self) -> Option<Arc<WorkerShard>> {
537        if !self.inner.session_recording {
538            return None;
539        }
540        THREAD_SHARDS.with(|assignments| {
541            let mut assignments = assignments.borrow_mut();
542            assignments.retain(|assignment| assignment.runtime.strong_count() != 0);
543            if let Some(assignment) = assignments.iter().find(|assignment| {
544                std::ptr::eq(assignment.runtime.as_ptr(), Arc::as_ptr(&self.inner))
545            }) {
546                return Some(assignment.shard.clone());
547            }
548            let shard = Arc::new(WorkerShard::new(
549                self.inner
550                    .next_shard_sequence
551                    .fetch_add(1, Ordering::Relaxed),
552            ));
553            match self.inner.session_shards.lock() {
554                Ok(mut shards) => shards.push(shard.clone()),
555                Err(poisoned) => poisoned.into_inner().push(shard.clone()),
556            }
557            assignments.push(ThreadShardAssignment {
558                runtime: Arc::downgrade(&self.inner),
559                shard: shard.clone(),
560            });
561            Some(shard)
562        })
563    }
564
565    /// Number of isolated worker counter shards registered by this runtime.
566    pub fn worker_shard_count(&self) -> usize {
567        match self.inner.session_shards.lock() {
568            Ok(shards) => shards.len(),
569            Err(poisoned) => poisoned.into_inner().len(),
570        }
571    }
572
573    /// Shard sequence of the calling thread without registering a new shard.
574    fn peek_worker_shard_sequence(&self) -> u64 {
575        THREAD_SHARDS.with(|assignments| {
576            assignments
577                .borrow()
578                .iter()
579                .find(|assignment| {
580                    std::ptr::eq(assignment.runtime.as_ptr(), Arc::as_ptr(&self.inner))
581                })
582                .map_or(0, |assignment| assignment.shard.sequence)
583        })
584    }
585
586    /// Capture a portable token naming this runtime's current causal parent.
587    pub fn causal_parent(&self) -> CausalParent {
588        CausalParent {
589            context_id: self.inner.context_id,
590            span_id: self.current_parent_span_id(),
591        }
592    }
593
594    /// Make this runtime current on the calling thread until the guard is dropped.
595    pub fn enter(&self) -> ContextGuard {
596        let _ = self.worker_shard();
597        CURRENT.with(|stack| stack.borrow_mut().push(self.clone()));
598        ACTIVE_CONTEXTS.fetch_add(1, Ordering::Relaxed);
599        ContextGuard {
600            runtime: self.clone(),
601            not_send: PhantomData,
602        }
603    }
604
605    fn enter_async_parent(&self, span_id: u64) -> Option<AsyncParentGuard> {
606        let stack_slot =
607            ASYNC_PARENT_SPANS.with(|stack| stack.borrow().iter().position(Option::is_none))?;
608        let runtime_key = Arc::as_ptr(&self.inner) as usize;
609        ASYNC_PARENT_SPANS.with(|stack| {
610            stack.borrow_mut()[stack_slot] = Some(ActiveSpan {
611                runtime_key,
612                span_id,
613            });
614        });
615        Some(AsyncParentGuard {
616            runtime_key,
617            span_id,
618            stack_slot,
619            not_send: PhantomData,
620        })
621    }
622
623    /// Run one synchronous closure with this runtime as its current context.
624    pub fn scope<T>(&self, operation: impl FnOnce() -> T) -> T {
625        let _guard = self.enter();
626        operation()
627    }
628
629    fn elapsed_ns(&self) -> u64 {
630        u64::try_from(self.inner.started.elapsed().as_nanos()).unwrap_or(u64::MAX)
631    }
632
633    pub(crate) fn add_counter(&self, counter: crate::CounterId, delta: u64) {
634        if let Some(shard) = self.worker_shard() {
635            shard.counter_values[counter.metric_id() as usize].fetch_add(delta, Ordering::Relaxed);
636        } else if !self.inner.session_recording {
637            self.inner.legacy_typed_counters[counter.metric_id() as usize]
638                .fetch_add(delta, Ordering::Relaxed);
639        }
640    }
641
642    fn record_distribution(&self, metric_id: crate::MetricId, value: u64) {
643        let index = metric_id as usize;
644        let bucket = latency_bucket_index(value);
645        self.inner.distribution_buckets[index][bucket].fetch_add(1, Ordering::Relaxed);
646        self.inner.distribution_min[index].fetch_min(value, Ordering::Relaxed);
647        self.inner.distribution_max[index].fetch_max(value, Ordering::Relaxed);
648    }
649
650    /// Count retry annotations recorded so far without draining them.
651    pub(crate) fn retry_annotation_count(&self) -> u64 {
652        if !self.inner.session_recording {
653            return 0;
654        }
655        let annotations = match self.inner.session_annotations.lock() {
656            Ok(annotations) => annotations,
657            Err(poisoned) => poisoned.into_inner(),
658        };
659        annotations
660            .iter()
661            .filter(|annotation| annotation.annotation_id == crate::AnnotationId::RetryAttempt)
662            .count() as u64
663    }
664
665    /// Read one session gauge without clearing it; `None` when never set.
666    pub(crate) fn session_gauge(&self, gauge: crate::GaugeId) -> Option<u64> {
667        if !self.inner.session_recording {
668            return None;
669        }
670        let index = gauge.metric_id() as usize;
671        let present = self.inner.session_gauge_present[index / 64].load(Ordering::Relaxed);
672        (present & (1_u64 << (index % 64)) != 0)
673            .then(|| self.inner.session_gauge_values[index].load(Ordering::Relaxed))
674    }
675
676    /// Record the current retained-buffer level and its running high water.
677    pub(crate) fn record_retained_buffer_bytes(&self, bytes: u64) {
678        if !self.inner.session_recording {
679            return;
680        }
681        self.set_gauge(crate::GaugeId::RetainedBufferBytes, bytes);
682        let peak = crate::GaugeId::RetainedBufferPeakBytes.metric_id() as usize;
683        self.inner.session_gauge_values[peak].fetch_max(bytes, Ordering::Relaxed);
684        self.inner.session_gauge_present[peak / 64]
685            .fetch_or(1_u64 << (peak % 64), Ordering::Relaxed);
686    }
687
688    /// Drain caller-recorded value distributions in stable metric order.
689    pub fn take_metric_distributions(&self) -> Vec<MetricDistributionV2> {
690        let mut records = Vec::new();
691        for index in 0..crate::MetricId::COUNT {
692            let mut call_count = 0_u64;
693            let buckets: Vec<DistributionBucketV2> = (0..LATENCY_BUCKET_COUNT)
694                .filter_map(|bucket| {
695                    let count =
696                        self.inner.distribution_buckets[index][bucket].swap(0, Ordering::Relaxed);
697                    call_count = call_count.saturating_add(count);
698                    if count == 0 {
699                        return None;
700                    }
701                    let (lower_bound, upper_bound) = latency_bucket_bounds(bucket);
702                    Some(DistributionBucketV2 {
703                        version: 1,
704                        lower_bound,
705                        upper_bound,
706                        count,
707                    })
708                })
709                .collect();
710            if call_count == 0 {
711                continue;
712            }
713            let minimum = self.inner.distribution_min[index].swap(u64::MAX, Ordering::Relaxed);
714            let maximum = self.inner.distribution_max[index].swap(0, Ordering::Relaxed);
715            records.push(MetricDistributionV2 {
716                version: 1,
717                metric_id: crate::METRICS[index].id,
718                call_count,
719                minimum,
720                maximum,
721                buckets,
722            });
723        }
724        records
725    }
726
727    /// Drain typed counters recorded by the standalone (non-session) runtime.
728    pub fn take_legacy_typed_metrics(&self) -> Vec<TypedMetricRecordV2> {
729        if self.inner.session_recording {
730            return self.take_session_typed_metrics();
731        }
732        let mut records = Vec::new();
733        for counter in crate::CounterId::ALL {
734            let metric_id = counter.metric_id();
735            let value =
736                self.inner.legacy_typed_counters[metric_id as usize].swap(0, Ordering::Relaxed);
737            if value != 0 {
738                records.push(TypedMetricRecordV2 {
739                    version: 1,
740                    metric_id,
741                    kind: crate::MetricKind::Counter,
742                    value,
743                });
744            }
745        }
746        records
747    }
748
749    pub(crate) fn set_gauge(&self, gauge: crate::GaugeId, value: u64) {
750        if !self.inner.session_recording {
751            return;
752        }
753        let index = gauge.metric_id() as usize;
754        self.inner.session_gauge_values[index].store(value, Ordering::Relaxed);
755        self.inner.session_gauge_present[index / 64]
756            .fetch_or(1_u64 << (index % 64), Ordering::Relaxed);
757    }
758
759    fn record_event(&self, event_id: crate::EventId, value: u64) -> bool {
760        if !self.inner.session_recording {
761            return false;
762        }
763        let sequence = self
764            .inner
765            .session_event_sequence
766            .fetch_add(1, Ordering::Relaxed);
767        let event = PointEventV2 {
768            version: 2,
769            sequence,
770            event_id,
771            elapsed_ns: self.elapsed_ns(),
772            thread_id: numeric_thread_id(),
773            value,
774            task_id: evidence_or_unavailable(current_task_id()),
775            worker_id: evidence_or_unavailable(self.peek_worker_shard_sequence()),
776        };
777        let mut events = match self.inner.session_point_events.lock() {
778            Ok(events) => events,
779            Err(poisoned) => poisoned.into_inner(),
780        };
781        if events.len() == MAX_POINT_EVENTS {
782            self.inner
783                .session_dropped_point_events
784                .fetch_add(1, Ordering::Relaxed);
785            false
786        } else {
787            events.push(event);
788            true
789        }
790    }
791
792    fn record_sampled_event(
793        &self,
794        event_id: crate::EventId,
795        value: u64,
796        policy: SamplingPolicy,
797    ) -> bool {
798        if !self.inner.session_recording {
799            return false;
800        }
801        let index = event_id.index();
802        let observation =
803            self.inner.session_sample_observations[index].fetch_add(1, Ordering::Relaxed);
804        let retained = policy.selects(observation)
805            && self.inner.session_sample_retained[index]
806                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |retained| {
807                    (retained < policy.maximum_retained).then_some(retained + 1)
808                })
809                .is_ok();
810        if retained {
811            self.record_event(event_id, value)
812        } else {
813            self.inner
814                .session_sampled_out_events
815                .fetch_add(1, Ordering::Relaxed);
816            false
817        }
818    }
819
820    fn record_annotation(&self, annotation_id: crate::AnnotationId, value: u64) {
821        if !self.inner.session_recording {
822            return;
823        }
824        let sequence = self
825            .inner
826            .session_event_sequence
827            .fetch_add(1, Ordering::Relaxed);
828        let annotation = AnnotationV2 {
829            version: 2,
830            sequence,
831            annotation_id,
832            elapsed_ns: self.elapsed_ns(),
833            thread_id: numeric_thread_id(),
834            value,
835            task_id: evidence_or_unavailable(current_task_id()),
836            worker_id: evidence_or_unavailable(self.peek_worker_shard_sequence()),
837        };
838        let mut annotations = match self.inner.session_annotations.lock() {
839            Ok(annotations) => annotations,
840            Err(poisoned) => poisoned.into_inner(),
841        };
842        if annotations.len() == MAX_ANNOTATIONS {
843            self.inner
844                .session_dropped_annotations
845                .fetch_add(1, Ordering::Relaxed);
846        } else {
847            annotations.push(annotation);
848        }
849    }
850
851    /// Drain typed counter and gauge records in stable metric order.
852    pub fn take_session_typed_metrics(&self) -> Vec<TypedMetricRecordV2> {
853        let mut records =
854            Vec::with_capacity(crate::CounterId::ALL.len() + crate::GaugeId::ALL.len());
855        let shards = self.inner.sorted_shards();
856        for counter in crate::CounterId::ALL {
857            let metric_id = counter.metric_id();
858            let value = shards.iter().fold(0_u64, |total, shard| {
859                total.saturating_add(
860                    shard.counter_values[metric_id as usize].swap(0, Ordering::Relaxed),
861                )
862            });
863            if value != 0 {
864                records.push(TypedMetricRecordV2 {
865                    version: 1,
866                    metric_id,
867                    kind: crate::MetricKind::Counter,
868                    value,
869                });
870            }
871        }
872        drop(shards);
873        let present = [
874            self.inner.session_gauge_present[0].swap(0, Ordering::Relaxed),
875            self.inner.session_gauge_present[1].swap(0, Ordering::Relaxed),
876        ];
877        for gauge in crate::GaugeId::ALL {
878            let metric_id = gauge.metric_id();
879            let index = metric_id as usize;
880            if present[index / 64] & (1_u64 << (index % 64)) != 0 {
881                records.push(TypedMetricRecordV2 {
882                    version: 1,
883                    metric_id,
884                    kind: crate::MetricKind::Gauge,
885                    value: self.inner.session_gauge_values[metric_id as usize]
886                        .swap(0, Ordering::Relaxed),
887                });
888            }
889        }
890        records.sort_unstable_by_key(|record| record.metric_id);
891        records
892    }
893
894    /// Drain typed timeline records and exact loss counts by cause.
895    pub fn take_session_typed_events(
896        &self,
897    ) -> (Vec<PointEventV2>, Vec<AnnotationV2>, EventLossCounts) {
898        let mut events = match self.inner.session_point_events.lock() {
899            Ok(events) => events,
900            Err(poisoned) => poisoned.into_inner(),
901        };
902        let mut events = std::mem::take(&mut *events);
903        let mut annotations = match self.inner.session_annotations.lock() {
904            Ok(annotations) => annotations,
905            Err(poisoned) => poisoned.into_inner(),
906        };
907        let mut annotations = std::mem::take(&mut *annotations);
908        events.sort_unstable_by_key(|event| event.sequence);
909        annotations.sort_unstable_by_key(|annotation| annotation.sequence);
910        let loss = EventLossCounts {
911            point_events: self
912                .inner
913                .session_dropped_point_events
914                .swap(0, Ordering::Relaxed),
915            annotations: self
916                .inner
917                .session_dropped_annotations
918                .swap(0, Ordering::Relaxed),
919            sampled_out_events: self
920                .inner
921                .session_sampled_out_events
922                .swap(0, Ordering::Relaxed),
923        };
924        (events, annotations, loss)
925    }
926
927    fn current_parent_span_id(&self) -> u64 {
928        let runtime_key = Arc::as_ptr(&self.inner) as usize;
929        ACTIVE_SPANS
930            .with(|stack| {
931                stack
932                    .borrow()
933                    .iter()
934                    .rev()
935                    .flatten()
936                    .find(|active| active.runtime_key == runtime_key)
937                    .map(|active| active.span_id)
938            })
939            .or_else(|| {
940                ASYNC_PARENT_SPANS.with(|stack| {
941                    stack
942                        .borrow()
943                        .iter()
944                        .rev()
945                        .flatten()
946                        .find(|active| active.runtime_key == runtime_key)
947                        .map(|active| active.span_id)
948                })
949            })
950            .unwrap_or(0)
951    }
952
953    fn reserve_span(
954        &self,
955        stage: Stage,
956        started: Instant,
957        parent_span_id: u64,
958        stack_slot: Option<usize>,
959        worker_id: u64,
960    ) -> Option<SpanTrace> {
961        let reservation = self
962            .inner
963            .session_span_reservations
964            .fetch_add(1, Ordering::Relaxed);
965        if reservation >= MAX_RECORDED_SPANS {
966            self.inner
967                .session_dropped_spans
968                .fetch_add(1, Ordering::Relaxed);
969            return None;
970        }
971        let mut records = match self.inner.session_spans.lock() {
972            Ok(records) => records,
973            Err(poisoned) => poisoned.into_inner(),
974        };
975        let span_id = self
976            .inner
977            .session_span_sequence
978            .fetch_add(1, Ordering::Relaxed)
979            .saturating_add(1);
980        let record_index = records.len();
981        records.push(RawSpanRecord {
982            span_id,
983            parent_span_id,
984            metric_id: stage.into(),
985            start_ns: u64::try_from(
986                started
987                    .checked_duration_since(self.inner.started)
988                    .unwrap_or_default()
989                    .as_nanos(),
990            )
991            .unwrap_or(u64::MAX),
992            inclusive_ns: 0,
993            thread_id: numeric_thread_id(),
994            worker_id,
995            task_id: current_task_id(),
996            work_origin: current_work_origin(),
997            completed: false,
998            hardware: RawSpanHardware::begin(),
999        });
1000        Some(SpanTrace {
1001            record_index,
1002            span_id,
1003            stack_slot,
1004        })
1005    }
1006
1007    fn begin_span_with(
1008        &self,
1009        stage: Stage,
1010        started: Instant,
1011        parent_span_id: u64,
1012        worker_id: u64,
1013    ) -> Option<SpanTrace> {
1014        if !self.inner.session_recording {
1015            return None;
1016        }
1017        let stack_slot = ACTIVE_SPANS.with(|stack| stack.borrow().iter().position(Option::is_none));
1018        let Some(stack_slot) = stack_slot else {
1019            self.inner
1020                .session_dropped_spans
1021                .fetch_add(1, Ordering::Relaxed);
1022            return None;
1023        };
1024        let trace =
1025            self.reserve_span(stage, started, parent_span_id, Some(stack_slot), worker_id)?;
1026        let runtime_key = Arc::as_ptr(&self.inner) as usize;
1027        ACTIVE_SPANS.with(|stack| {
1028            stack.borrow_mut()[stack_slot] = Some(ActiveSpan {
1029                runtime_key,
1030                span_id: trace.span_id,
1031            });
1032        });
1033        Some(trace)
1034    }
1035
1036    fn begin_span(&self, stage: Stage, started: Instant, worker_id: u64) -> Option<SpanTrace> {
1037        let parent_span_id = self.current_parent_span_id();
1038        self.begin_span_with(stage, started, parent_span_id, worker_id)
1039    }
1040
1041    fn begin_async_span(
1042        &self,
1043        stage: Stage,
1044        started: Instant,
1045        parent_span_id: u64,
1046        worker_id: u64,
1047    ) -> Option<SpanTrace> {
1048        if !self.inner.session_recording {
1049            return None;
1050        }
1051        self.reserve_span(stage, started, parent_span_id, None, worker_id)
1052    }
1053
1054    fn finish_span(&self, trace: SpanTrace, inclusive_ns: u64) {
1055        if let Some(stack_slot) = trace.stack_slot {
1056            ACTIVE_SPANS.with(|stack| {
1057                let mut stack = stack.borrow_mut();
1058                if stack[stack_slot].is_some_and(|active| active.span_id == trace.span_id) {
1059                    stack[stack_slot] = None;
1060                }
1061            });
1062        }
1063        let mut records = match self.inner.session_spans.lock() {
1064            Ok(records) => records,
1065            Err(poisoned) => poisoned.into_inner(),
1066        };
1067        if let Some(record) = records.get_mut(trace.record_index) {
1068            record.inclusive_ns = inclusive_ns;
1069            record.completed = true;
1070            record.hardware.finish();
1071        }
1072    }
1073
1074    /// Drain bounded causal spans and the exact count omitted or unfinished.
1075    pub fn take_session_span_records(&self) -> (Vec<SpanRecordV2>, u64) {
1076        let mut records = match self.inner.session_spans.lock() {
1077            Ok(records) => records,
1078            Err(poisoned) => poisoned.into_inner(),
1079        };
1080        let mut raw = std::mem::take(&mut *records);
1081        drop(records);
1082        raw.sort_unstable_by_key(|record| record.span_id);
1083        let unfinished = raw.iter().filter(|record| !record.completed).count() as u64;
1084        raw.retain(|record| record.completed);
1085        let positions: HashMap<u64, usize> = raw
1086            .iter()
1087            .enumerate()
1088            .map(|(index, record)| (record.span_id, index))
1089            .collect();
1090        let mut exclusive: Vec<u64> = raw.iter().map(|record| record.inclusive_ns).collect();
1091        for record in &raw {
1092            if let Some(parent_index) = positions.get(&record.parent_span_id).copied() {
1093                exclusive[parent_index] =
1094                    exclusive[parent_index].saturating_sub(record.inclusive_ns);
1095            }
1096        }
1097        let spans = raw
1098            .into_iter()
1099            .enumerate()
1100            .map(|(index, record)| SpanRecordV2 {
1101                version: 3,
1102                span_id: record.span_id,
1103                parent_span_id: if positions.contains_key(&record.parent_span_id) {
1104                    Evidence::recorded(record.parent_span_id)
1105                } else {
1106                    Evidence::unavailable(EvidenceGap::Unavailable)
1107                },
1108                metric_id: record.metric_id,
1109                start_ns: record.start_ns,
1110                inclusive_ns: record.inclusive_ns,
1111                exclusive_ns: exclusive[index],
1112                thread_id: record.thread_id,
1113                task_id: evidence_or_unavailable(record.task_id),
1114                worker_id: evidence_or_unavailable(record.worker_id),
1115                work_origin: record.work_origin,
1116                hardware: record.hardware.into_evidence(),
1117            })
1118            .collect();
1119        let dropped = self
1120            .inner
1121            .session_dropped_spans
1122            .swap(0, Ordering::Relaxed)
1123            .saturating_add(unfinished);
1124        (spans, dropped)
1125    }
1126
1127    /// Drain per-micro-function logarithmic latency distributions.
1128    pub fn take_session_latency_distributions(&self) -> Vec<LatencyDistributionV2> {
1129        let shards = self.inner.sorted_shards();
1130        Stage::ALL
1131            .into_iter()
1132            .filter_map(|stage| {
1133                let mut call_count = 0_u64;
1134                let buckets: Vec<LatencyBucketV2> = (0..LATENCY_BUCKET_COUNT)
1135                    .filter_map(|index| {
1136                        let count = shards.iter().fold(0_u64, |total, shard| {
1137                            total.saturating_add(
1138                                shard.latency_buckets[stage.index()][index]
1139                                    .swap(0, Ordering::Relaxed),
1140                            )
1141                        });
1142                        call_count = call_count.saturating_add(count);
1143                        (count != 0).then(|| {
1144                            let (lower_bound_ns, upper_bound_ns) = latency_bucket_bounds(index);
1145                            LatencyBucketV2 {
1146                                version: 1,
1147                                lower_bound_ns,
1148                                upper_bound_ns,
1149                                count,
1150                            }
1151                        })
1152                    })
1153                    .collect();
1154                if call_count == 0 {
1155                    return None;
1156                }
1157                let minimum_ns = shards.iter().fold(u64::MAX, |minimum, shard| {
1158                    minimum
1159                        .min(shard.latency_min_ns[stage.index()].swap(u64::MAX, Ordering::Relaxed))
1160                });
1161                let maximum_ns = shards.iter().fold(0_u64, |maximum, shard| {
1162                    maximum.max(shard.latency_max_ns[stage.index()].swap(0, Ordering::Relaxed))
1163                });
1164                Some(LatencyDistributionV2 {
1165                    version: 2,
1166                    metric_id: stage.metric_id(),
1167                    macro_stage_id: stage.macro_stage_id(),
1168                    call_count,
1169                    minimum_ns,
1170                    maximum_ns,
1171                    p50_ns: percentile_upper_bound(&buckets, call_count, 50, maximum_ns),
1172                    p90_ns: percentile_upper_bound(&buckets, call_count, 90, maximum_ns),
1173                    p95_ns: percentile_upper_bound(&buckets, call_count, 95, maximum_ns),
1174                    p99_ns: percentile_upper_bound(&buckets, call_count, 99, maximum_ns),
1175                    buckets,
1176                })
1177            })
1178            .collect()
1179    }
1180
1181    fn record(&self, shard: Option<&WorkerShard>, stage: Stage, outcome: SpanOutcome) {
1182        let index = stage.index();
1183        let elapsed_ns = outcome.elapsed_ns;
1184        if self.inner.session_recording {
1185            let Some(shard) = shard else {
1186                return;
1187            };
1188            shard.elapsed_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1189            shard.calls[index].fetch_add(1, Ordering::Relaxed);
1190            shard.legacy_elapsed_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1191            shard.legacy_calls[index].fetch_add(1, Ordering::Relaxed);
1192            let bucket = latency_bucket_index(elapsed_ns);
1193            shard.latency_buckets[index][bucket].fetch_add(1, Ordering::Relaxed);
1194            shard.latency_min_ns[index].fetch_min(elapsed_ns, Ordering::Relaxed);
1195            shard.latency_max_ns[index].fetch_max(elapsed_ns, Ordering::Relaxed);
1196            shard.stage_first_start_ns[index].fetch_min(outcome.start_offset_ns, Ordering::Relaxed);
1197            shard.stage_last_end_ns[index].fetch_max(
1198                outcome.start_offset_ns.saturating_add(elapsed_ns),
1199                Ordering::Relaxed,
1200            );
1201            if outcome.attributed {
1202                shard.attributed_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1203                shard.legacy_attributed_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1204            }
1205            if outcome.blocked {
1206                shard.blocked_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1207                shard.blocked_calls[index].fetch_add(1, Ordering::Relaxed);
1208            }
1209            if outcome.serial {
1210                shard.serial_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1211                shard.serial_calls[index].fetch_add(1, Ordering::Relaxed);
1212            }
1213            // Only the outermost span on a thread contributes occupancy, so a
1214            // nested span never counts its parent's time a second time.
1215            if outcome.outermost {
1216                shard.top_level_calls.fetch_add(1, Ordering::Relaxed);
1217                if outcome.blocked {
1218                    shard
1219                        .top_level_blocked_ns
1220                        .fetch_add(elapsed_ns, Ordering::Relaxed);
1221                } else {
1222                    shard
1223                        .top_level_busy_ns
1224                        .fetch_add(elapsed_ns, Ordering::Relaxed);
1225                }
1226            }
1227            return;
1228        }
1229        self.inner.elapsed_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1230        self.inner.calls[index].fetch_add(1, Ordering::Relaxed);
1231        if outcome.attributed {
1232            self.inner.attributed_ns[index].fetch_add(elapsed_ns, Ordering::Relaxed);
1233        }
1234    }
1235
1236    fn add_stage_bytes(&self, stage: Stage, bytes: u64) {
1237        if let Some(shard) = self.worker_shard() {
1238            shard.stage_bytes[stage.index()].fetch_add(bytes, Ordering::Relaxed);
1239        }
1240    }
1241
1242    fn record_cache_outcome(&self, cache: crate::CacheId, hit: bool) {
1243        let Some(shard) = self.worker_shard() else {
1244            return;
1245        };
1246        let slot = if hit {
1247            &shard.cache_hits[cache.index()]
1248        } else {
1249            &shard.cache_misses[cache.index()]
1250        };
1251        slot.fetch_add(1, Ordering::Relaxed);
1252    }
1253
1254    fn record_retry(&self, cause: crate::RetryCause) {
1255        if let Some(shard) = self.worker_shard() {
1256            shard.retries[cause.index()].fetch_add(1, Ordering::Relaxed);
1257        }
1258    }
1259
1260    fn add_indexed_counter(&self, counter: crate::IndexedCounterId, slot: u16, delta: u64) {
1261        let Some(shard) = self.worker_shard() else {
1262            return;
1263        };
1264        // Folding an out-of-range slot into the last one would attribute one
1265        // caller's cost to another. Count it as dropped instead.
1266        match shard.indexed_counters[counter.index()].get(usize::from(slot)) {
1267            Some(cell) => {
1268                cell.fetch_add(delta, Ordering::Relaxed);
1269            }
1270            None => {
1271                shard
1272                    .indexed_counter_dropped
1273                    .fetch_add(1, Ordering::Relaxed);
1274            }
1275        }
1276    }
1277
1278    fn add_input_bytes(&self, bytes: u64) {
1279        if let Some(shard) = self.worker_shard() {
1280            shard.input_bytes.fetch_add(bytes, Ordering::Relaxed);
1281            shard.legacy_input_bytes.fetch_add(bytes, Ordering::Relaxed);
1282            shard.counter_values[crate::MetricId::InputBytes as usize]
1283                .fetch_add(bytes, Ordering::Relaxed);
1284        } else {
1285            self.inner.input_bytes.fetch_add(bytes, Ordering::Relaxed);
1286            self.inner.legacy_typed_counters[crate::MetricId::InputBytes as usize]
1287                .fetch_add(bytes, Ordering::Relaxed);
1288        }
1289    }
1290
1291    fn add_input_units(&self, units: u64) {
1292        if let Some(shard) = self.worker_shard() {
1293            shard.input_units.fetch_add(units, Ordering::Relaxed);
1294            shard.legacy_input_units.fetch_add(units, Ordering::Relaxed);
1295            shard.counter_values[crate::MetricId::InputUnits as usize]
1296                .fetch_add(units, Ordering::Relaxed);
1297        } else {
1298            self.inner.input_units.fetch_add(units, Ordering::Relaxed);
1299            self.inner.legacy_typed_counters[crate::MetricId::InputUnits as usize]
1300                .fetch_add(units, Ordering::Relaxed);
1301        }
1302    }
1303
1304    fn add_derived_decoder_bytes(&self, bytes: u64) {
1305        if let Some(shard) = self.worker_shard() {
1306            shard
1307                .derived_decoder_bytes
1308                .fetch_add(bytes, Ordering::Relaxed);
1309        }
1310    }
1311
1312    fn add_backend_dispatched_bytes(&self, bytes: u64) {
1313        if let Some(shard) = self.worker_shard() {
1314            shard
1315                .backend_dispatched_bytes
1316                .fetch_add(bytes, Ordering::Relaxed);
1317        }
1318    }
1319
1320    pub(crate) fn drain_stage_counters(&self, session: bool) -> RawStageCounters {
1321        if !session && !self.inner.session_recording {
1322            return RawStageCounters {
1323                elapsed_ns: std::array::from_fn(|index| {
1324                    self.inner.elapsed_ns[index].swap(0, Ordering::Relaxed)
1325                }),
1326                calls: std::array::from_fn(|index| {
1327                    self.inner.calls[index].swap(0, Ordering::Relaxed)
1328                }),
1329                attributed_ns: std::array::from_fn(|index| {
1330                    self.inner.attributed_ns[index].swap(0, Ordering::Relaxed)
1331                }),
1332            };
1333        }
1334        let shards = self.inner.sorted_shards();
1335        if !session {
1336            return RawStageCounters {
1337                elapsed_ns: std::array::from_fn(|index| {
1338                    shards.iter().fold(0_u64, |total, shard| {
1339                        total.saturating_add(
1340                            shard.legacy_elapsed_ns[index].swap(0, Ordering::Relaxed),
1341                        )
1342                    })
1343                }),
1344                calls: std::array::from_fn(|index| {
1345                    shards.iter().fold(0_u64, |total, shard| {
1346                        total.saturating_add(shard.legacy_calls[index].swap(0, Ordering::Relaxed))
1347                    })
1348                }),
1349                attributed_ns: std::array::from_fn(|index| {
1350                    shards.iter().fold(0_u64, |total, shard| {
1351                        total.saturating_add(
1352                            shard.legacy_attributed_ns[index].swap(0, Ordering::Relaxed),
1353                        )
1354                    })
1355                }),
1356            };
1357        }
1358        RawStageCounters {
1359            elapsed_ns: std::array::from_fn(|index| {
1360                shards.iter().fold(0_u64, |total, shard| {
1361                    total.saturating_add(shard.elapsed_ns[index].swap(0, Ordering::Relaxed))
1362                })
1363            }),
1364            calls: std::array::from_fn(|index| {
1365                shards.iter().fold(0_u64, |total, shard| {
1366                    total.saturating_add(shard.calls[index].swap(0, Ordering::Relaxed))
1367                })
1368            }),
1369            attributed_ns: std::array::from_fn(|index| {
1370                shards.iter().fold(0_u64, |total, shard| {
1371                    total.saturating_add(shard.attributed_ns[index].swap(0, Ordering::Relaxed))
1372                })
1373            }),
1374        }
1375    }
1376
1377    fn take_input_totals(&self) -> (u64, u64) {
1378        if !self.inner.session_recording {
1379            return (
1380                self.inner.input_bytes.swap(0, Ordering::Relaxed),
1381                self.inner.input_units.swap(0, Ordering::Relaxed),
1382            );
1383        }
1384        let shards = self.inner.sorted_shards();
1385        shards.iter().fold((0_u64, 0_u64), |totals, shard| {
1386            (
1387                totals
1388                    .0
1389                    .saturating_add(shard.legacy_input_bytes.swap(0, Ordering::Relaxed)),
1390                totals
1391                    .1
1392                    .saturating_add(shard.legacy_input_units.swap(0, Ordering::Relaxed)),
1393            )
1394        })
1395    }
1396
1397    pub(crate) fn take_session_input_totals(&self) -> (u64, u64) {
1398        let shards = self.inner.sorted_shards();
1399        shards.iter().fold((0_u64, 0_u64), |totals, shard| {
1400            (
1401                totals
1402                    .0
1403                    .saturating_add(shard.input_bytes.swap(0, Ordering::Relaxed)),
1404                totals
1405                    .1
1406                    .saturating_add(shard.input_units.swap(0, Ordering::Relaxed)),
1407            )
1408        })
1409    }
1410
1411    pub(crate) fn take_session_workload_totals(&self) -> (u64, u64) {
1412        let shards = self.inner.sorted_shards();
1413        shards.iter().fold((0_u64, 0_u64), |totals, shard| {
1414            (
1415                totals
1416                    .0
1417                    .saturating_add(shard.derived_decoder_bytes.swap(0, Ordering::Relaxed)),
1418                totals
1419                    .1
1420                    .saturating_add(shard.backend_dispatched_bytes.swap(0, Ordering::Relaxed)),
1421            )
1422        })
1423    }
1424
1425    fn record_batch_route(
1426        &self,
1427        workload_key_digest: &str,
1428        requested_backend: &str,
1429        selected_backend: &str,
1430        completed_backend: &str,
1431        recovered_from_backend: Option<&str>,
1432    ) {
1433        if !self.inner.session_recording {
1434            return;
1435        }
1436        let batch_sequence = self
1437            .inner
1438            .session_route_sequence
1439            .fetch_add(1, Ordering::Relaxed);
1440        self.record_event(crate::EventId::BackendBatchCompleted, batch_sequence);
1441        if recovered_from_backend.is_some() {
1442            self.record_event(crate::EventId::BackendRecovered, batch_sequence);
1443        }
1444        let record = BatchRouteV2 {
1445            version: 1,
1446            batch_sequence,
1447            workload_key_digest: workload_key_digest.to_owned(),
1448            requested_backend: requested_backend.to_owned(),
1449            selected_backend: selected_backend.to_owned(),
1450            completed_backend: completed_backend.to_owned(),
1451            recovered_from_backend: recovered_from_backend.map_or_else(
1452                || Evidence::unavailable(crate::schema_v2::EvidenceGap::Unavailable),
1453                |backend| Evidence::recorded(backend.to_owned()),
1454            ),
1455        };
1456        let mut records = match self.inner.session_batch_routes.lock() {
1457            Ok(records) => records,
1458            Err(poisoned) => poisoned.into_inner(),
1459        };
1460        if records.len() >= MAX_BATCH_ROUTES {
1461            self.inner
1462                .session_dropped_batch_routes
1463                .fetch_add(1, Ordering::Relaxed);
1464            return;
1465        }
1466        records.push(record);
1467    }
1468
1469    /// Drain completed batch-route records from this profiling runtime.
1470    pub fn take_session_batch_routes(&self) -> Vec<BatchRouteV2> {
1471        let mut records = match self.inner.session_batch_routes.lock() {
1472            Ok(records) => records,
1473            Err(poisoned) => poisoned.into_inner(),
1474        };
1475        let mut drained = std::mem::take(&mut *records);
1476        drained.sort_unstable_by_key(|record| record.batch_sequence);
1477        drained
1478    }
1479
1480    /// Drain the count of batch routes dropped after [`MAX_BATCH_ROUTES`].
1481    pub fn take_session_dropped_batch_routes(&self) -> u64 {
1482        self.inner
1483            .session_dropped_batch_routes
1484            .swap(0, Ordering::Relaxed)
1485    }
1486
1487    fn record_queue_enqueue(&self, queue: crate::QueueId, sequence: u64) {
1488        if !self.inner.session_recording {
1489            return;
1490        }
1491        let enqueue = PendingQueueEnqueue {
1492            thread_id: numeric_thread_id(),
1493            elapsed_ns: self.elapsed_ns(),
1494        };
1495        let mut pending = match self.inner.queue_pending.lock() {
1496            Ok(pending) => pending,
1497            Err(poisoned) => poisoned.into_inner(),
1498        };
1499        if pending.len() == MAX_QUEUE_LINKS && !pending.contains_key(&(queue as u8, sequence)) {
1500            self.inner
1501                .queue_dropped_enqueues
1502                .fetch_add(1, Ordering::Relaxed);
1503            return;
1504        }
1505        if pending.insert((queue as u8, sequence), enqueue).is_some() {
1506            // A duplicate (queue, sequence) enqueue displaces the earlier record.
1507            self.inner
1508                .queue_dropped_enqueues
1509                .fetch_add(1, Ordering::Relaxed);
1510        }
1511    }
1512
1513    fn record_queue_dequeue(&self, queue: crate::QueueId, sequence: u64) {
1514        if !self.inner.session_recording {
1515            return;
1516        }
1517        let enqueue = {
1518            let mut pending = match self.inner.queue_pending.lock() {
1519                Ok(pending) => pending,
1520                Err(poisoned) => poisoned.into_inner(),
1521            };
1522            pending.remove(&(queue as u8, sequence))
1523        };
1524        let Some(enqueue) = enqueue else {
1525            self.inner
1526                .queue_unmatched_dequeues
1527                .fetch_add(1, Ordering::Relaxed);
1528            return;
1529        };
1530        let link = QueueLinkV2 {
1531            version: 1,
1532            queue,
1533            sequence,
1534            producer_thread_id: enqueue.thread_id,
1535            producer_elapsed_ns: enqueue.elapsed_ns,
1536            consumer_thread_id: numeric_thread_id(),
1537            consumer_elapsed_ns: self.elapsed_ns(),
1538        };
1539        let mut links = match self.inner.queue_links.lock() {
1540            Ok(links) => links,
1541            Err(poisoned) => poisoned.into_inner(),
1542        };
1543        if links.len() == MAX_QUEUE_LINKS {
1544            self.inner
1545                .queue_dropped_links
1546                .fetch_add(1, Ordering::Relaxed);
1547        } else {
1548            links.push(link);
1549        }
1550    }
1551
1552    /// Drain matched queue links in stable (queue, sequence) order with exact loss counts.
1553    pub fn take_session_queue_links(&self) -> (Vec<QueueLinkV2>, QueueLinkLossCounts) {
1554        let mut links = match self.inner.queue_links.lock() {
1555            Ok(links) => links,
1556            Err(poisoned) => poisoned.into_inner(),
1557        };
1558        let mut links = std::mem::take(&mut *links);
1559        links.sort_unstable_by_key(|link| (link.queue, link.sequence));
1560        let unconsumed = {
1561            let mut pending = match self.inner.queue_pending.lock() {
1562                Ok(pending) => pending,
1563                Err(poisoned) => poisoned.into_inner(),
1564            };
1565            let pending = std::mem::take(&mut *pending);
1566            u64::try_from(pending.len()).unwrap_or(u64::MAX)
1567        };
1568        let loss = QueueLinkLossCounts {
1569            dropped_enqueues: self.inner.queue_dropped_enqueues.swap(0, Ordering::Relaxed),
1570            dropped_links: self.inner.queue_dropped_links.swap(0, Ordering::Relaxed),
1571            unmatched_dequeues: self
1572                .inner
1573                .queue_unmatched_dequeues
1574                .swap(0, Ordering::Relaxed),
1575            unconsumed_enqueues: unconsumed,
1576        };
1577        (links, loss)
1578    }
1579
1580    fn queue_depth_enqueue(&self, queue: crate::QueueId) {
1581        if !self.inner.session_recording {
1582            return;
1583        }
1584        let index = queue.index();
1585        let depth = self.inner.queue_depth_current[index].fetch_add(1, Ordering::Relaxed) + 1;
1586        self.inner.queue_depth_high_water[index].fetch_max(depth, Ordering::Relaxed);
1587        self.inner.queue_depth_enqueues[index].fetch_add(1, Ordering::Relaxed);
1588    }
1589
1590    fn queue_depth_dequeue(&self, queue: crate::QueueId) {
1591        if !self.inner.session_recording {
1592            return;
1593        }
1594        let index = queue.index();
1595        let _ = self.inner.queue_depth_current[index].fetch_update(
1596            Ordering::Relaxed,
1597            Ordering::Relaxed,
1598            |depth| Some(depth.saturating_sub(1)),
1599        );
1600        self.inner.queue_depth_dequeues[index].fetch_add(1, Ordering::Relaxed);
1601    }
1602
1603    fn set_queue_depth(&self, queue: crate::QueueId, depth: u64) {
1604        if !self.inner.session_recording {
1605            return;
1606        }
1607        let index = queue.index();
1608        self.inner.queue_depth_current[index].store(depth, Ordering::Relaxed);
1609        self.inner.queue_depth_high_water[index].fetch_max(depth, Ordering::Relaxed);
1610    }
1611
1612    /// Drain queue depth occupancy and high-water records in stable queue order.
1613    ///
1614    /// The current depth is reported without being reset; the high-water mark
1615    /// restarts from the current depth and the enqueue/dequeue totals reset.
1616    pub fn take_session_queue_depths(&self) -> Vec<QueueDepthV2> {
1617        if !self.inner.session_recording {
1618            return Vec::new();
1619        }
1620        crate::QueueId::ALL
1621            .into_iter()
1622            .filter_map(|queue| {
1623                let index = queue.index();
1624                let current = self.inner.queue_depth_current[index].load(Ordering::Relaxed);
1625                let high_water =
1626                    self.inner.queue_depth_high_water[index].swap(current, Ordering::Relaxed);
1627                let enqueues = self.inner.queue_depth_enqueues[index].swap(0, Ordering::Relaxed);
1628                let dequeues = self.inner.queue_depth_dequeues[index].swap(0, Ordering::Relaxed);
1629                (current != 0 || high_water != 0 || enqueues != 0 || dequeues != 0).then_some(
1630                    QueueDepthV2 {
1631                        version: 1,
1632                        queue,
1633                        current,
1634                        high_water,
1635                        enqueues,
1636                        dequeues,
1637                    },
1638                )
1639            })
1640            .collect()
1641    }
1642
1643    /// Drain per-worker load and imbalance evidence merged in sorted shard order.
1644    ///
1645    /// This drains the same per-stage call and elapsed counters as
1646    /// `take_session_stage_measurements`; call it once at session drain.
1647    pub fn take_session_worker_imbalance(&self) -> WorkerImbalanceV2 {
1648        let shards = self.inner.sorted_shards();
1649        let workers: Vec<WorkerLoadV2> = shards
1650            .iter()
1651            .map(|shard| {
1652                let mut calls = 0_u64;
1653                let mut elapsed_ns = 0_u64;
1654                for index in 0..STAGE_COUNT {
1655                    calls = calls.saturating_add(shard.calls[index].swap(0, Ordering::Relaxed));
1656                    elapsed_ns = elapsed_ns
1657                        .saturating_add(shard.elapsed_ns[index].swap(0, Ordering::Relaxed));
1658                }
1659                WorkerLoadV2 {
1660                    version: 1,
1661                    worker_id: shard.sequence,
1662                    calls,
1663                    elapsed_ns,
1664                }
1665            })
1666            .collect();
1667        let worker_count = u64::try_from(workers.len()).unwrap_or(u64::MAX);
1668        let total_calls = workers
1669            .iter()
1670            .fold(0_u64, |total, worker| total.saturating_add(worker.calls));
1671        let total_elapsed_ns = workers.iter().fold(0_u64, |total, worker| {
1672            total.saturating_add(worker.elapsed_ns)
1673        });
1674        let ppm = |part: u64, whole: u64| -> u64 {
1675            if whole == 0 {
1676                0
1677            } else {
1678                u64::try_from((u128::from(part) * 1_000_000) / u128::from(whole))
1679                    .unwrap_or(u64::MAX)
1680            }
1681        };
1682        let max_calls = workers.iter().map(|worker| worker.calls).max().unwrap_or(0);
1683        let mut sorted_calls: Vec<u64> = workers.iter().map(|worker| worker.calls).collect();
1684        sorted_calls.sort_unstable();
1685        let median_calls = if sorted_calls.is_empty() {
1686            0
1687        } else {
1688            sorted_calls[sorted_calls.len() / 2]
1689        };
1690        let idle_workers = workers.iter().filter(|worker| worker.calls == 0).count() as u64;
1691        WorkerImbalanceV2 {
1692            version: 1,
1693            worker_count,
1694            total_calls,
1695            total_elapsed_ns,
1696            max_share_ppm: ppm(max_calls, total_calls),
1697            median_share_ppm: ppm(median_calls, total_calls),
1698            idle_share_ppm: ppm(idle_workers, worker_count),
1699            workers,
1700        }
1701    }
1702
1703    /// Drain blocked-wait records merged in sorted shard order.
1704    pub fn take_session_blocked_waits(&self) -> Vec<BlockedWaitRecordV2> {
1705        let shards = self.inner.sorted_shards();
1706        Stage::ALL
1707            .into_iter()
1708            .filter_map(|stage| {
1709                let index = stage.index();
1710                let calls = shards.iter().fold(0_u64, |total, shard| {
1711                    total.saturating_add(shard.blocked_calls[index].swap(0, Ordering::Relaxed))
1712                });
1713                if calls == 0 {
1714                    return None;
1715                }
1716                let blocked_ns = shards.iter().fold(0_u64, |total, shard| {
1717                    total.saturating_add(shard.blocked_ns[index].swap(0, Ordering::Relaxed))
1718                });
1719                Some(BlockedWaitRecordV2 {
1720                    version: 1,
1721                    metric_id: stage.metric_id(),
1722                    macro_stage_id: stage.macro_stage_id(),
1723                    calls,
1724                    blocked_ns,
1725                })
1726            })
1727            .collect()
1728    }
1729
1730    /// Offset of one instant from this runtime's start, saturating at zero.
1731    fn offset_ns(&self, at: Instant) -> u64 {
1732        u64::try_from(
1733            at.checked_duration_since(self.inner.started)
1734                .unwrap_or_default()
1735                .as_nanos(),
1736        )
1737        .unwrap_or(u64::MAX)
1738    }
1739
1740    /// Drain per-micro-function wall-clock occupancy merged across workers.
1741    ///
1742    /// Call this before `Session::finish`, which drains the shared per-stage
1743    /// call and elapsed counters this record reads.
1744    pub fn take_session_stage_concurrency(&self) -> Vec<StageConcurrencyV2> {
1745        let shards = self.inner.sorted_shards();
1746        Stage::ALL
1747            .into_iter()
1748            .filter_map(|stage| {
1749                let index = stage.index();
1750                let mut calls = 0_u64;
1751                let mut elapsed_ns = 0_u64;
1752                let mut max_worker_elapsed_ns = 0_u64;
1753                let mut worker_count = 0_u64;
1754                let mut first_start_ns = u64::MAX;
1755                let mut last_end_ns = 0_u64;
1756                let mut declared_serial_ns = 0_u64;
1757                let mut declared_serial_calls = 0_u64;
1758                let mut bytes = 0_u64;
1759                for shard in &shards {
1760                    let shard_calls = shard.calls[index].load(Ordering::Relaxed);
1761                    let shard_elapsed = shard.elapsed_ns[index].load(Ordering::Relaxed);
1762                    if shard_calls != 0 {
1763                        worker_count += 1;
1764                    }
1765                    calls = calls.saturating_add(shard_calls);
1766                    elapsed_ns = elapsed_ns.saturating_add(shard_elapsed);
1767                    max_worker_elapsed_ns = max_worker_elapsed_ns.max(shard_elapsed);
1768                    first_start_ns = first_start_ns
1769                        .min(shard.stage_first_start_ns[index].load(Ordering::Relaxed));
1770                    last_end_ns =
1771                        last_end_ns.max(shard.stage_last_end_ns[index].load(Ordering::Relaxed));
1772                    declared_serial_ns = declared_serial_ns
1773                        .saturating_add(shard.serial_ns[index].swap(0, Ordering::Relaxed));
1774                    declared_serial_calls = declared_serial_calls
1775                        .saturating_add(shard.serial_calls[index].swap(0, Ordering::Relaxed));
1776                    bytes =
1777                        bytes.saturating_add(shard.stage_bytes[index].swap(0, Ordering::Relaxed));
1778                }
1779                if calls == 0 {
1780                    return None;
1781                }
1782                let first_start_ns = if first_start_ns == u64::MAX {
1783                    0
1784                } else {
1785                    first_start_ns
1786                };
1787                let window_ns = last_end_ns.saturating_sub(first_start_ns);
1788                // A window of zero means every call fell inside one clock tick;
1789                // report the calls as serial rather than inventing concurrency.
1790                // A stage entered recursively on one thread sums its nested
1791                // time, so raw elapsed can exceed the wall the thread spent
1792                // there. Average concurrency can never exceed the number of
1793                // workers that entered the stage, so cap it there rather than
1794                // report a single-threaded recursion as parallel.
1795                let ceiling_milli = worker_count.saturating_mul(1_000);
1796                let concurrency_milli = if window_ns == 0 {
1797                    1_000
1798                } else {
1799                    u64::try_from((u128::from(elapsed_ns) * 1_000) / u128::from(window_ns))
1800                        .unwrap_or(u64::MAX)
1801                        .min(ceiling_milli)
1802                };
1803                Some(StageConcurrencyV2 {
1804                    version: crate::schema_v2::STAGE_CONCURRENCY_V2_VERSION,
1805                    metric_id: stage.metric_id(),
1806                    macro_stage_id: stage.macro_stage_id(),
1807                    calls,
1808                    elapsed_ns,
1809                    window_ns,
1810                    first_start_ns,
1811                    last_end_ns,
1812                    worker_count,
1813                    max_worker_elapsed_ns,
1814                    concurrency_milli,
1815                    declared_serial_ns,
1816                    declared_serial_calls,
1817                    bytes,
1818                })
1819            })
1820            .collect()
1821    }
1822
1823    /// Drain per-worker busy and blocked time merged in sorted shard order.
1824    pub fn take_session_worker_occupancy(&self) -> WorkerOccupancyV2 {
1825        let shards = self.inner.sorted_shards();
1826        let workers: Vec<WorkerOccupancyRowV2> = shards
1827            .iter()
1828            .map(|shard| WorkerOccupancyRowV2 {
1829                version: crate::schema_v2::WORKER_OCCUPANCY_V2_VERSION,
1830                worker_id: shard.sequence,
1831                busy_ns: shard.top_level_busy_ns.swap(0, Ordering::Relaxed),
1832                blocked_ns: shard.top_level_blocked_ns.swap(0, Ordering::Relaxed),
1833                calls: shard.top_level_calls.swap(0, Ordering::Relaxed),
1834            })
1835            .collect();
1836        let busy_ns = workers
1837            .iter()
1838            .fold(0_u64, |total, worker| total.saturating_add(worker.busy_ns));
1839        let blocked_ns = workers.iter().fold(0_u64, |total, worker| {
1840            total.saturating_add(worker.blocked_ns)
1841        });
1842        let calls = workers
1843            .iter()
1844            .fold(0_u64, |total, worker| total.saturating_add(worker.calls));
1845        let mut sorted_busy: Vec<u64> = workers.iter().map(|worker| worker.busy_ns).collect();
1846        sorted_busy.sort_unstable();
1847        let median_busy_ns = sorted_busy
1848            .get(sorted_busy.len() / 2)
1849            .copied()
1850            .unwrap_or_default();
1851        WorkerOccupancyV2 {
1852            version: crate::schema_v2::WORKER_OCCUPANCY_V2_VERSION,
1853            worker_count: u64::try_from(workers.len()).unwrap_or(u64::MAX),
1854            active_worker_count: workers.iter().filter(|worker| worker.calls != 0).count() as u64,
1855            busy_ns,
1856            blocked_ns,
1857            calls,
1858            busiest_busy_ns: sorted_busy.last().copied().unwrap_or_default(),
1859            median_busy_ns,
1860            workers,
1861        }
1862    }
1863
1864    /// Drain reuse-cache hit and miss counts merged in sorted shard order.
1865    pub fn take_session_cache_effectiveness(&self) -> Vec<CacheEffectivenessV2> {
1866        let shards = self.inner.sorted_shards();
1867        crate::CacheId::ALL
1868            .into_iter()
1869            .filter_map(|cache| {
1870                let index = cache.index();
1871                let mut hits = 0_u64;
1872                let mut misses = 0_u64;
1873                for shard in &shards {
1874                    hits = hits.saturating_add(shard.cache_hits[index].swap(0, Ordering::Relaxed));
1875                    misses =
1876                        misses.saturating_add(shard.cache_misses[index].swap(0, Ordering::Relaxed));
1877                }
1878                let total = hits.saturating_add(misses);
1879                if total == 0 {
1880                    return None;
1881                }
1882                Some(CacheEffectivenessV2 {
1883                    version: crate::schema_v2::CACHE_EFFECTIVENESS_V2_VERSION,
1884                    cache,
1885                    hits,
1886                    misses,
1887                    hit_rate_ppm: u64::try_from((u128::from(hits) * 1_000_000) / u128::from(total))
1888                        .unwrap_or(u64::MAX),
1889                })
1890            })
1891            .collect()
1892    }
1893
1894    /// Drain retry attempts by cause, merged in sorted shard order.
1895    pub fn take_session_retries(&self) -> Vec<RetryRecordV2> {
1896        let shards = self.inner.sorted_shards();
1897        crate::RetryCause::ALL
1898            .into_iter()
1899            .filter_map(|cause| {
1900                let attempts = shards.iter().fold(0_u64, |total, shard| {
1901                    total.saturating_add(shard.retries[cause.index()].swap(0, Ordering::Relaxed))
1902                });
1903                (attempts != 0).then_some(RetryRecordV2 {
1904                    version: crate::schema_v2::RETRY_RECORD_V2_VERSION,
1905                    cause,
1906                    attempts,
1907                })
1908            })
1909            .collect()
1910    }
1911
1912    /// Drain indexed counter families merged in sorted shard order.
1913    pub fn take_session_indexed_counters(&self) -> Vec<IndexedCounterRecordV2> {
1914        let shards = self.inner.sorted_shards();
1915        let dropped_out_of_range = shards.iter().fold(0_u64, |total, shard| {
1916            total.saturating_add(shard.indexed_counter_dropped.swap(0, Ordering::Relaxed))
1917        });
1918        crate::IndexedCounterId::ALL
1919            .into_iter()
1920            .filter_map(|counter| {
1921                let index = counter.index();
1922                let mut slots = [0_u64; crate::INDEXED_COUNTER_SLOTS];
1923                for shard in &shards {
1924                    for (slot, total) in slots.iter_mut().enumerate() {
1925                        *total = total.saturating_add(
1926                            shard.indexed_counters[index][slot].swap(0, Ordering::Relaxed),
1927                        );
1928                    }
1929                }
1930                if slots.iter().all(|value| *value == 0) && dropped_out_of_range == 0 {
1931                    return None;
1932                }
1933                Some(IndexedCounterRecordV2 {
1934                    version: crate::schema_v2::INDEXED_COUNTER_V2_VERSION,
1935                    counter,
1936                    slots: slots.to_vec(),
1937                    dropped_out_of_range,
1938                })
1939            })
1940            .collect()
1941    }
1942
1943    /// Discard every per-run accumulator this runtime owns.
1944    ///
1945    /// Benchmarks call this between measured rounds to drop warm-up, so
1946    /// anything left behind is reported as part of the next round. That makes
1947    /// a partial reset a wrong number presented as a right one, which is why
1948    /// this clears the per-worker shards as well as the runtime-level stores.
1949    fn reset(&self) {
1950        // Session drains clear their own storage; the legacy drain clears the
1951        // legacy mirrors. Both are needed because a session runtime keeps two.
1952        let _ = self.drain_stage_counters(false);
1953        let _ = self.drain_stage_counters(true);
1954        let _ = self.take_session_worker_occupancy();
1955        let _ = self.take_session_blocked_waits();
1956        let _ = self.take_session_stage_concurrency();
1957        let _ = self.take_session_cache_effectiveness();
1958        let _ = self.take_session_indexed_counters();
1959        let _ = self.take_session_retries();
1960        let _ = self.take_session_queue_depths();
1961        self.inner.input_bytes.store(0, Ordering::Relaxed);
1962        self.inner.input_units.store(0, Ordering::Relaxed);
1963        for index in 0..crate::MetricId::COUNT {
1964            self.inner.legacy_typed_counters[index].store(0, Ordering::Relaxed);
1965            self.inner.distribution_min[index].store(u64::MAX, Ordering::Relaxed);
1966            self.inner.distribution_max[index].store(0, Ordering::Relaxed);
1967            for bucket in 0..LATENCY_BUCKET_COUNT {
1968                self.inner.distribution_buckets[index][bucket].store(0, Ordering::Relaxed);
1969            }
1970        }
1971        for shard in self.inner.sorted_shards() {
1972            shard.input_bytes.store(0, Ordering::Relaxed);
1973            shard.input_units.store(0, Ordering::Relaxed);
1974            shard.legacy_input_bytes.store(0, Ordering::Relaxed);
1975            shard.legacy_input_units.store(0, Ordering::Relaxed);
1976            shard.derived_decoder_bytes.store(0, Ordering::Relaxed);
1977            shard.backend_dispatched_bytes.store(0, Ordering::Relaxed);
1978            for index in 0..crate::MetricId::COUNT {
1979                shard.counter_values[index].store(0, Ordering::Relaxed);
1980            }
1981            for index in 0..STAGE_COUNT {
1982                shard.latency_min_ns[index].store(u64::MAX, Ordering::Relaxed);
1983                shard.latency_max_ns[index].store(0, Ordering::Relaxed);
1984                shard.stage_first_start_ns[index].store(u64::MAX, Ordering::Relaxed);
1985                shard.stage_last_end_ns[index].store(0, Ordering::Relaxed);
1986                for bucket in 0..LATENCY_BUCKET_COUNT {
1987                    shard.latency_buckets[index][bucket].store(0, Ordering::Relaxed);
1988                }
1989            }
1990        }
1991    }
1992}
1993
1994impl Default for Runtime {
1995    fn default() -> Self {
1996        Self::new()
1997    }
1998}
1999
2000/// Thread context guard returned by [`Runtime::enter`].
2001pub struct ContextGuard {
2002    runtime: Runtime,
2003    not_send: PhantomData<Rc<()>>,
2004}
2005
2006impl Drop for ContextGuard {
2007    fn drop(&mut self) {
2008        CURRENT.with(|stack| {
2009            let mut stack = stack.borrow_mut();
2010            if stack
2011                .last()
2012                .is_some_and(|runtime| Arc::ptr_eq(&runtime.inner, &self.runtime.inner))
2013            {
2014                stack.pop();
2015            } else if let Some(position) = stack
2016                .iter()
2017                .rposition(|runtime| Arc::ptr_eq(&runtime.inner, &self.runtime.inner))
2018            {
2019                stack.remove(position);
2020            }
2021        });
2022        ACTIVE_CONTEXTS.fetch_sub(1, Ordering::Relaxed);
2023    }
2024}
2025
2026struct AsyncParentGuard {
2027    runtime_key: usize,
2028    span_id: u64,
2029    stack_slot: usize,
2030    not_send: PhantomData<Rc<()>>,
2031}
2032
2033impl Drop for AsyncParentGuard {
2034    fn drop(&mut self) {
2035        ASYNC_PARENT_SPANS.with(|stack| {
2036            let mut stack = stack.borrow_mut();
2037            if stack[self.stack_slot].is_some_and(|active| {
2038                active.runtime_key == self.runtime_key && active.span_id == self.span_id
2039            }) {
2040                stack[self.stack_slot] = None;
2041            }
2042        });
2043    }
2044}
2045
2046struct WorkOriginGuard {
2047    previous: WorkOrigin,
2048}
2049
2050impl WorkOriginGuard {
2051    fn enter(origin: WorkOrigin) -> Self {
2052        Self {
2053            previous: WORK_ORIGIN.with(|slot| slot.replace(origin)),
2054        }
2055    }
2056}
2057
2058impl Drop for WorkOriginGuard {
2059    fn drop(&mut self) {
2060        WORK_ORIGIN.with(|slot| slot.set(self.previous));
2061    }
2062}
2063
2064struct TaskGuard {
2065    previous: u64,
2066}
2067
2068impl TaskGuard {
2069    fn enter(task_id: u64) -> Self {
2070        Self {
2071            previous: TASK_ID.with(|slot| slot.replace(task_id)),
2072        }
2073    }
2074}
2075
2076impl Drop for TaskGuard {
2077    fn drop(&mut self) {
2078        TASK_ID.with(|slot| slot.set(self.previous));
2079    }
2080}
2081
2082struct LegacyRuntime {
2083    runtime: Runtime,
2084    enabled: bool,
2085}
2086
2087impl LegacyRuntime {
2088    fn new() -> Self {
2089        Self {
2090            runtime: Runtime::legacy(),
2091            enabled: false,
2092        }
2093    }
2094}
2095
2096/// Optional attribution for work performed inside a derived input.
2097#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
2098#[repr(u8)]
2099pub enum Attribution {
2100    #[default]
2101    Root = 0,
2102    Decoded = 1,
2103}
2104
2105impl From<Attribution> for WorkOrigin {
2106    fn from(attribution: Attribution) -> Self {
2107        match attribution {
2108            Attribution::Root => Self::Root,
2109            Attribution::Decoded => Self::Decoded,
2110        }
2111    }
2112}
2113
2114thread_local! {
2115    static CURRENT: RefCell<Vec<Runtime>> = const { RefCell::new(Vec::new()) };
2116    static WORK_ORIGIN: Cell<WorkOrigin> = const { Cell::new(WorkOrigin::Root) };
2117    static TASK_ID: Cell<u64> = const { Cell::new(0) };
2118    static SPAN_DEPTH: Cell<u32> = const { Cell::new(0) };
2119    static LEGACY: RefCell<LegacyRuntime> = RefCell::new(LegacyRuntime::new());
2120    static ACTIVE_SPANS: RefCell<[Option<ActiveSpan>; MAX_NESTED_SPANS]> =
2121        const { RefCell::new([None; MAX_NESTED_SPANS]) };
2122    static ASYNC_PARENT_SPANS: RefCell<[Option<ActiveSpan>; MAX_NESTED_SPANS]> =
2123        const { RefCell::new([None; MAX_NESTED_SPANS]) };
2124    static NUMERIC_THREAD_ID: u64 = NEXT_THREAD_ID.fetch_add(1, Ordering::Relaxed);
2125    static THREAD_SHARDS: RefCell<Vec<ThreadShardAssignment>> = const { RefCell::new(Vec::new()) };
2126}
2127
2128fn numeric_thread_id() -> u64 {
2129    NUMERIC_THREAD_ID.with(|thread_id| *thread_id)
2130}
2131
2132/// Return a clone of the runtime current on this thread.
2133pub fn current_runtime() -> Option<Runtime> {
2134    if ACTIVE_CONTEXTS.load(Ordering::Relaxed) == 0 {
2135        return None;
2136    }
2137    CURRENT
2138        .with(|stack| stack.borrow().last().cloned())
2139        .or_else(|| {
2140            LEGACY.with(|legacy| {
2141                let legacy = legacy.borrow();
2142                legacy.enabled.then(|| legacy.runtime.clone())
2143            })
2144        })
2145}
2146
2147pub(crate) fn runtime_for_drain() -> Runtime {
2148    CURRENT
2149        .with(|stack| stack.borrow().last().cloned())
2150        .unwrap_or_else(|| LEGACY.with(|legacy| legacy.borrow().runtime.clone()))
2151}
2152
2153/// Replace this thread's attribution and return its previous value.
2154///
2155/// Legacy projection of [`set_work_origin`]: the returned value is `Decoded`
2156/// whenever the previous origin was any attributed (non-root) work.
2157pub fn set_attribution(attribution: Attribution) -> Attribution {
2158    let previous = set_work_origin(attribution.into());
2159    if previous.is_attributed_work() {
2160        Attribution::Decoded
2161    } else {
2162        Attribution::Root
2163    }
2164}
2165
2166/// Replace this thread's causal work origin and return its previous value.
2167pub fn set_work_origin(origin: WorkOrigin) -> WorkOrigin {
2168    WORK_ORIGIN.with(|slot| slot.replace(origin))
2169}
2170
2171/// Current thread's causal work origin.
2172pub fn current_work_origin() -> WorkOrigin {
2173    WORK_ORIGIN.with(|slot| slot.get())
2174}
2175
2176/// Replace this thread's caller-assigned task identity and return the previous.
2177///
2178/// Zero clears the identity. The value is a caller-managed task or name index;
2179/// [`instrument_future`] propagates it across polls and worker threads.
2180pub fn set_task_id(task_id: u64) -> u64 {
2181    TASK_ID.with(|slot| slot.replace(task_id))
2182}
2183
2184/// Current thread's caller-assigned task identity, or zero when unset.
2185pub fn current_task_id() -> u64 {
2186    TASK_ID.with(|slot| slot.get())
2187}
2188
2189fn evidence_or_unavailable(value: u64) -> Evidence<u64> {
2190    if value == 0 {
2191        Evidence::unavailable(EvidenceGap::Unavailable)
2192    } else {
2193        Evidence::recorded(value)
2194    }
2195}
2196
2197/// Capture a portable token naming the current runtime's causal parent.
2198///
2199/// Returns `None` when no runtime is current on this thread.
2200pub fn current_causal_parent() -> Option<CausalParent> {
2201    current_runtime().map(|runtime| runtime.causal_parent())
2202}
2203
2204/// Return whether fixed-stage profiling is active on the calling thread.
2205#[inline]
2206pub fn enabled() -> bool {
2207    if ACTIVE_CONTEXTS.load(Ordering::Relaxed) == 0 {
2208        return false;
2209    }
2210    CURRENT.with(|stack| !stack.borrow().is_empty())
2211        || LEGACY.with(|legacy| legacy.borrow().enabled)
2212}
2213
2214/// Enable or disable the calling thread's standalone profiling runtime.
2215///
2216/// Prefer [`crate::Session::start`] for operator runs because it also captures
2217/// identity, resources, and state transitions. This switch remains available
2218/// to libraries and microbenchmarks that only need stage counters.
2219pub fn set_enabled(enabled: bool) {
2220    LEGACY.with(|legacy| {
2221        let mut legacy = legacy.borrow_mut();
2222        if legacy.enabled == enabled {
2223            return;
2224        }
2225        legacy.enabled = enabled;
2226        if enabled {
2227            ACTIVE_CONTEXTS.fetch_add(1, Ordering::Relaxed);
2228        } else {
2229            ACTIVE_CONTEXTS.fetch_sub(1, Ordering::Relaxed);
2230        }
2231    });
2232}
2233
2234struct AsyncSpan {
2235    runtime: Option<Runtime>,
2236    shard: Option<Arc<WorkerShard>>,
2237    stage: Stage,
2238    started: Option<Instant>,
2239    trace: Option<SpanTrace>,
2240    attributed: bool,
2241}
2242
2243impl Drop for AsyncSpan {
2244    fn drop(&mut self) {
2245        let (Some(runtime), Some(started)) = (&self.runtime, self.started) else {
2246            return;
2247        };
2248        let elapsed_ns = u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX);
2249        runtime.record(
2250            self.shard.as_deref(),
2251            self.stage,
2252            SpanOutcome {
2253                start_offset_ns: runtime.offset_ns(started),
2254                elapsed_ns,
2255                attributed: self.attributed,
2256                blocked: false,
2257                serial: false,
2258                // An async span can be polled on any thread, so it never
2259                // owns a worker's outermost interval.
2260                outermost: false,
2261            },
2262        );
2263        if let Some(trace) = self.trace {
2264            runtime.finish_span(trace, elapsed_ns);
2265        }
2266    }
2267}
2268
2269fn instrument_impl<F>(
2270    stage: Stage,
2271    future: F,
2272    parent_override: Option<CausalParent>,
2273) -> impl Future<Output = F::Output>
2274where
2275    F: Future,
2276{
2277    let runtime = current_runtime();
2278    let shard = runtime.as_ref().and_then(Runtime::worker_shard);
2279    let worker_id = shard.as_ref().map_or(0, |shard| shard.sequence);
2280    let started = runtime.as_ref().map(|_| Instant::now());
2281    let parent_span_id = match (runtime.as_ref(), parent_override) {
2282        (Some(runtime), Some(parent)) if runtime.context_id() == parent.context_id() => {
2283            parent.span_id()
2284        }
2285        (Some(_), Some(_)) => 0,
2286        (Some(runtime), None) => runtime.current_parent_span_id(),
2287        (None, _) => 0,
2288    };
2289    let trace = runtime
2290        .as_ref()
2291        .zip(started)
2292        .and_then(|(runtime, started)| {
2293            runtime.begin_async_span(stage, started, parent_span_id, worker_id)
2294        });
2295    let span_id = trace.map(|trace| trace.span_id);
2296    let origin = current_work_origin();
2297    let task_id = current_task_id();
2298    let attributed = runtime.is_some() && origin.is_attributed_work();
2299    let poll_runtime = runtime.clone();
2300
2301    async move {
2302        let _span = AsyncSpan {
2303            runtime,
2304            shard,
2305            stage,
2306            started,
2307            trace,
2308            attributed,
2309        };
2310        let mut future = std::pin::pin!(future);
2311        poll_fn(|context| {
2312            let _context_guard = poll_runtime.as_ref().map(Runtime::enter);
2313            let _parent_guard = poll_runtime
2314                .as_ref()
2315                .zip(span_id)
2316                .and_then(|(runtime, span_id)| runtime.enter_async_parent(span_id));
2317            let _origin_guard = WorkOriginGuard::enter(origin);
2318            let _task_guard = TaskGuard::enter(task_id);
2319            if span_id.is_some() {
2320                crate::allocation::stage_context_push(stage);
2321            }
2322            let result = future.as_mut().poll(context);
2323            if span_id.is_some() {
2324                crate::allocation::stage_context_pop();
2325            }
2326            result
2327        })
2328        .await
2329    }
2330}
2331
2332/// Propagate the current runtime and causal parent while polling one future.
2333///
2334/// The returned future records one wall-time span from wrapper construction
2335/// through completion or cancellation. Runtime and parent guards exist only
2336/// during each poll, so the returned future remains `Send` when `F` is `Send`.
2337pub fn instrument_future<F>(stage: Stage, future: F) -> impl Future<Output = F::Output>
2338where
2339    F: Future,
2340{
2341    instrument_impl(stage, future, None)
2342}
2343
2344/// Propagate the current runtime with an explicit portable causal parent.
2345///
2346/// Use this when the future crosses a spawn boundary where thread-local
2347/// parentage cannot reach. A token captured from another runtime records the
2348/// future's span as a root of the current runtime instead.
2349pub fn instrument_future_with_parent<F>(
2350    parent: CausalParent,
2351    stage: Stage,
2352    future: F,
2353) -> impl Future<Output = F::Output>
2354where
2355    F: Future,
2356{
2357    instrument_impl(stage, future, Some(parent))
2358}
2359
2360/// Allocation-free stage guard. It contains no start timestamp while disabled.
2361#[must_use]
2362pub struct Span {
2363    runtime: Option<Runtime>,
2364    shard: Option<Arc<WorkerShard>>,
2365    stage: Stage,
2366    started: Option<Instant>,
2367    trace: Option<SpanTrace>,
2368    attributed: bool,
2369    blocked: bool,
2370    serial: bool,
2371    outermost: bool,
2372}
2373
2374impl Span {
2375    /// Whether this span reads and will record a clock measurement.
2376    pub fn is_recording(&self) -> bool {
2377        self.started.is_some()
2378    }
2379}
2380
2381fn span_impl(
2382    stage: Stage,
2383    parent_override: Option<CausalParent>,
2384    blocked: bool,
2385    serial: bool,
2386) -> Span {
2387    if ACTIVE_CONTEXTS.load(Ordering::Relaxed) == 0 {
2388        return Span {
2389            runtime: None,
2390            shard: None,
2391            stage,
2392            started: None,
2393            trace: None,
2394            attributed: false,
2395            blocked: false,
2396            serial: false,
2397            outermost: false,
2398        };
2399    }
2400    let runtime = current_runtime();
2401    let shard = runtime.as_ref().and_then(Runtime::worker_shard);
2402    let worker_id = shard.as_ref().map_or(0, |shard| shard.sequence);
2403    let started = runtime.as_ref().map(|_| Instant::now());
2404    // Depth is tracked for every recording guard, including guards whose span
2405    // record was dropped for capacity, so occupancy never double-counts a
2406    // nested region just because the forest was truncated.
2407    let outermost =
2408        started.is_some() && SPAN_DEPTH.with(|depth| depth.replace(depth.get() + 1) == 0);
2409    let trace = match (runtime.as_ref(), started, parent_override) {
2410        (Some(runtime), Some(started), Some(parent)) => {
2411            let parent_span_id = if runtime.context_id() == parent.context_id() {
2412                parent.span_id()
2413            } else {
2414                0
2415            };
2416            runtime.begin_span_with(stage, started, parent_span_id, worker_id)
2417        }
2418        (Some(runtime), Some(started), None) => runtime.begin_span(stage, started, worker_id),
2419        _ => None,
2420    };
2421    // Blocked wait is never attributed execution.
2422    let attributed = !blocked && runtime.is_some() && current_work_origin().is_attributed_work();
2423    if trace.is_some() {
2424        crate::allocation::stage_context_push(stage);
2425    }
2426    Span {
2427        runtime,
2428        shard,
2429        stage,
2430        started,
2431        trace,
2432        attributed,
2433        blocked,
2434        serial,
2435        outermost,
2436    }
2437}
2438
2439/// Start one fixed-stage measurement.
2440#[inline]
2441pub fn span(stage: Stage) -> Span {
2442    span_impl(stage, None, false, false)
2443}
2444
2445/// Start one fixed-stage measurement with an explicit portable causal parent.
2446///
2447/// The token replaces thread-local parent lookup, so the caller can carry
2448/// parentage across crate, thread, or spawn boundaries. A token captured from
2449/// another runtime records this span as a root of the current runtime.
2450pub fn span_with_parent(parent: CausalParent, stage: Stage) -> Span {
2451    span_impl(stage, Some(parent), false, false)
2452}
2453
2454/// Declare that this region runs with the worker pool idle.
2455///
2456/// Use it for barriers such as enumeration, plan compilation, or final merge,
2457/// where no other worker can make progress. The guard measures wall time like
2458/// [`span`] and additionally accumulates declared-serial time, which the
2459/// profiler reports as the Amdahl floor on any speedup from more threads.
2460/// The profiler also derives serial phases from observed concurrency, so a
2461/// region left undeclared is still detected; declaring it makes the report
2462/// state the intent rather than infer it.
2463pub fn serial_span(stage: Stage) -> Span {
2464    span_impl(stage, None, false, true)
2465}
2466
2467/// Record one blocked wait interval separately from runnable execution.
2468///
2469/// The guard measures wall time like [`span`] but additionally accumulates
2470/// per-stage blocked time drained by `Runtime::take_session_blocked_waits`,
2471/// and it never counts as attributed (decoded or derived) execution. Reuse
2472/// wait stages such as [`Stage::SourceQueueWait`] and [`Stage::ScannerQueueWait`].
2473pub fn blocked(stage: Stage) -> Span {
2474    span_impl(stage, None, true, false)
2475}
2476
2477/// Time a region whose measurement drives a decision, profiled or not.
2478///
2479/// [`span`] deliberately measures nothing while profiling is off, which is
2480/// right for reporting and wrong for a measurement the product acts on. A
2481/// timer whose value picks a backend must produce the same value whether or
2482/// not an operator passed `--profile`, or the flag would change routing.
2483///
2484/// This guard therefore reads the clock unconditionally, returns the elapsed
2485/// duration to the caller, and additionally records it as a [`span`] would
2486/// when a runtime is current. The name says "decision" so a reader knows at
2487/// the call site that it costs a clock read even when profiling is off. On a
2488/// hot path where nothing acts on the value, use [`span`].
2489#[must_use = "a decision timer only measures when finished"]
2490pub struct DecisionTimer {
2491    stage: Stage,
2492    started: Instant,
2493}
2494
2495/// Start a decision-driving measurement of one micro-function.
2496pub fn decision_timer(stage: Stage) -> DecisionTimer {
2497    DecisionTimer {
2498        stage,
2499        started: Instant::now(),
2500    }
2501}
2502
2503impl DecisionTimer {
2504    /// Stop the timer, record it when profiling is on, and return the elapsed time.
2505    pub fn finish(self) -> std::time::Duration {
2506        let elapsed = self.started.elapsed();
2507        if ACTIVE_CONTEXTS.load(Ordering::Relaxed) != 0 {
2508            if let Some(runtime) = current_runtime() {
2509                let shard = runtime.worker_shard();
2510                runtime.record(
2511                    shard.as_deref(),
2512                    self.stage,
2513                    SpanOutcome {
2514                        start_offset_ns: runtime.offset_ns(self.started),
2515                        elapsed_ns: u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX),
2516                        attributed: false,
2517                        blocked: false,
2518                        serial: false,
2519                        // The caller owns the enclosing span, if any; a decision
2520                        // timer never claims a worker's outermost interval.
2521                        outermost: false,
2522                    },
2523                );
2524            }
2525        }
2526        elapsed
2527    }
2528}
2529
2530impl Drop for Span {
2531    #[inline]
2532    fn drop(&mut self) {
2533        let (Some(runtime), Some(started)) = (&self.runtime, self.started) else {
2534            return;
2535        };
2536        SPAN_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
2537        let elapsed_ns = u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX);
2538        runtime.record(
2539            self.shard.as_deref(),
2540            self.stage,
2541            SpanOutcome {
2542                start_offset_ns: runtime.offset_ns(started),
2543                elapsed_ns,
2544                attributed: self.attributed,
2545                blocked: self.blocked,
2546                serial: self.serial,
2547                outermost: self.outermost,
2548            },
2549        );
2550        if let Some(trace) = self.trace {
2551            runtime.finish_span(trace, elapsed_ns);
2552            crate::allocation::stage_context_pop();
2553        }
2554    }
2555}
2556
2557/// Time a sub-stage region into a [`crate::CounterId`] instead of a stage.
2558///
2559/// Some measurements sit strictly inside a stage leaf. Recording them as
2560/// spans would add their time to that leaf's inclusive total a second time,
2561/// so they belong in a counter. This guard is [`span`] with a counter sink:
2562/// same enabled gate, same clock, no allocation on drop.
2563#[must_use = "a counter span measures nothing until it is dropped"]
2564pub struct CounterSpan {
2565    counter: crate::CounterId,
2566    started: Option<Instant>,
2567}
2568
2569/// Start a sub-stage measurement that accumulates into one counter.
2570#[inline]
2571pub fn counter_span(counter: crate::CounterId) -> CounterSpan {
2572    CounterSpan {
2573        counter,
2574        started: (ACTIVE_CONTEXTS.load(Ordering::Relaxed) != 0).then(Instant::now),
2575    }
2576}
2577
2578impl Drop for CounterSpan {
2579    #[inline]
2580    fn drop(&mut self) {
2581        let Some(started) = self.started else {
2582            return;
2583        };
2584        add_counter(
2585            self.counter,
2586            u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX),
2587        );
2588    }
2589}
2590
2591/// Count one retry attempt, whether or not the retry eventually succeeded.
2592///
2593/// Count every attempt, not every operation: a path that retries a thousand
2594/// times must read as a thousand. The profiler reports retries as a finding
2595/// because a retry that fires means a failure the product did not design out.
2596#[inline]
2597pub fn record_retry(cause: crate::RetryCause) {
2598    if let Some(runtime) = current_runtime() {
2599        runtime.record_retry(cause);
2600    }
2601}
2602
2603/// Add to one slot of an indexed counter family.
2604///
2605/// A slot at or beyond [`crate::INDEXED_COUNTER_SLOTS`] is counted as dropped
2606/// on the drained record rather than folded into another slot.
2607#[inline]
2608pub fn add_indexed_counter(counter: crate::IndexedCounterId, slot: u16, delta: u64) {
2609    if let Some(runtime) = current_runtime() {
2610        runtime.add_indexed_counter(counter, slot, delta);
2611    }
2612}
2613
2614/// Attribute bytes to one micro-function so its throughput can be reported.
2615///
2616/// Record the bytes the named stage actually moved or examined, not the run's
2617/// input size. A stage that sees each byte twice reports twice the bytes, and
2618/// that is the honest number for its own throughput.
2619#[inline]
2620pub fn add_stage_bytes(stage: crate::Stage, bytes: u64) {
2621    if let Some(runtime) = current_runtime() {
2622        runtime.add_stage_bytes(stage, bytes);
2623    }
2624}
2625
2626/// Count one consultation of a reuse cache that was served from the cache.
2627#[inline]
2628pub fn record_cache_hit(cache: crate::CacheId) {
2629    if let Some(runtime) = current_runtime() {
2630        runtime.record_cache_outcome(cache, true);
2631    }
2632}
2633
2634/// Count one consultation of a reuse cache that had to recompute or refetch.
2635#[inline]
2636pub fn record_cache_miss(cache: crate::CacheId) {
2637    if let Some(runtime) = current_runtime() {
2638        runtime.record_cache_outcome(cache, false);
2639    }
2640}
2641
2642/// Add source bytes processed by the current profile.
2643#[inline]
2644pub fn add_input_bytes(bytes: u64) {
2645    if let Some(runtime) = current_runtime() {
2646        runtime.add_input_bytes(bytes);
2647    }
2648}
2649
2650/// Add source units such as files, objects, responses, or chunks.
2651#[inline]
2652pub fn add_input_units(units: u64) {
2653    if let Some(runtime) = current_runtime() {
2654        runtime.add_input_units(units);
2655    }
2656}
2657
2658/// Record one expensive detail event under a deterministic bounded sampling policy.
2659#[inline]
2660pub fn record_sampled_event(event: crate::EventId, value: u64, policy: SamplingPolicy) -> bool {
2661    current_runtime().is_some_and(|runtime| runtime.record_sampled_event(event, value, policy))
2662}
2663/// Add bytes produced by accepted decode-through work in the current profile.
2664#[inline]
2665pub fn add_derived_decoder_bytes(bytes: u64) {
2666    if let Some(runtime) = current_runtime() {
2667        runtime.add_derived_decoder_bytes(bytes);
2668    }
2669}
2670
2671/// Add bytes submitted once to the completed backend route in the current profile.
2672#[inline]
2673pub fn add_backend_dispatched_bytes(bytes: u64) {
2674    if let Some(runtime) = current_runtime() {
2675        runtime.add_backend_dispatched_bytes(bytes);
2676    }
2677}
2678
2679/// Increment one typed monotonic counter in the current profiling runtime.
2680#[inline]
2681pub fn add_counter(counter: crate::CounterId, delta: u64) {
2682    if let Some(runtime) = current_runtime() {
2683        runtime.add_counter(counter, delta);
2684    }
2685}
2686
2687/// Replace one typed latest-value gauge in the current profiling runtime.
2688#[inline]
2689pub fn set_gauge(gauge: crate::GaugeId, value: u64) {
2690    if let Some(runtime) = current_runtime() {
2691        runtime.set_gauge(gauge, value);
2692    }
2693}
2694
2695/// Record one typed instantaneous event with a numeric payload.
2696#[inline]
2697pub fn record_event(event: crate::EventId, value: u64) {
2698    if let Some(runtime) = current_runtime() {
2699        runtime.record_event(event, value);
2700    }
2701}
2702
2703/// Record one producer enqueue for later matching by [`record_queue_dequeue`].
2704///
2705/// The `(queue, sequence)` pair must be unique per in-flight item. Retention
2706/// is bounded by [`MAX_QUEUE_LINKS`]; loss is counted explicitly.
2707pub fn record_queue_enqueue(queue: crate::QueueId, sequence: u64) {
2708    if let Some(runtime) = current_runtime() {
2709        runtime.record_queue_enqueue(queue, sequence);
2710    }
2711}
2712
2713/// Record the consumer dequeue matching one earlier [`record_queue_enqueue`].
2714///
2715/// A dequeue with no recorded pending enqueue increments the unmatched count.
2716pub fn record_queue_dequeue(queue: crate::QueueId, sequence: u64) {
2717    if let Some(runtime) = current_runtime() {
2718        runtime.record_queue_dequeue(queue, sequence);
2719    }
2720}
2721
2722/// Increment one queue's depth gauge and refresh its high-water mark.
2723pub fn record_queue_depth_enqueue(queue: crate::QueueId) {
2724    if let Some(runtime) = current_runtime() {
2725        runtime.queue_depth_enqueue(queue);
2726    }
2727}
2728
2729/// Decrement one queue's depth gauge, saturating at zero.
2730pub fn record_queue_depth_dequeue(queue: crate::QueueId) {
2731    if let Some(runtime) = current_runtime() {
2732        runtime.queue_depth_dequeue(queue);
2733    }
2734}
2735
2736/// Replace one queue's depth gauge and refresh its high-water mark.
2737pub fn set_queue_depth(queue: crate::QueueId, depth: u64) {
2738    if let Some(runtime) = current_runtime() {
2739        runtime.set_queue_depth(queue, depth);
2740    }
2741}
2742
2743/// Record one observed value into a metric's bounded logarithmic distribution.
2744#[inline]
2745pub fn record_distribution(metric: crate::MetricId, value: u64) {
2746    if let Some(runtime) = current_runtime() {
2747        runtime.record_distribution(metric, value);
2748    }
2749}
2750
2751/// Record one filesystem open latency inside a [`Stage::SourceWalk`] or
2752/// [`Stage::SourceRead`] instrumented path.
2753#[inline]
2754pub fn record_fs_open_latency_ns(elapsed_ns: u64) {
2755    record_distribution(crate::MetricId::FsOpenLatencyNs, elapsed_ns);
2756}
2757
2758/// Record one filesystem read latency inside a [`Stage::SourceRead`]
2759/// instrumented path.
2760#[inline]
2761pub fn record_fs_read_latency_ns(elapsed_ns: u64) {
2762    record_distribution(crate::MetricId::FsReadLatencyNs, elapsed_ns);
2763}
2764
2765/// Record one filesystem metadata (stat/readdir) latency inside a
2766/// [`Stage::SourceWalk`] instrumented path.
2767#[inline]
2768pub fn record_fs_metadata_latency_ns(elapsed_ns: u64) {
2769    record_distribution(crate::MetricId::FsMetadataLatencyNs, elapsed_ns);
2770}
2771
2772/// Record one network request latency observed by a caller.
2773#[inline]
2774pub fn record_network_latency_ns(elapsed_ns: u64) {
2775    record_distribution(crate::MetricId::NetworkLatencyNs, elapsed_ns);
2776}
2777
2778/// Add network bytes a caller read and wrote; process-level counters are not
2779/// visible to the profiler on every host, so callers report their own IO.
2780#[inline]
2781pub fn record_network_bytes(read_bytes: u64, written_bytes: u64) {
2782    if read_bytes > 0 {
2783        add_counter(crate::CounterId::NetworkBytesRead, read_bytes);
2784    }
2785    if written_bytes > 0 {
2786        add_counter(crate::CounterId::NetworkBytesWritten, written_bytes);
2787    }
2788}
2789
2790/// Count one completed network request.
2791#[inline]
2792pub fn record_network_request() {
2793    add_counter(crate::CounterId::NetworkRequests, 1);
2794}
2795
2796/// Record one explicitly observed page-cache state for IO work.
2797///
2798/// The observation becomes one [`crate::AnnotationId::IoCacheState`] timeline
2799/// record and increments the matching observation counter. The profiler
2800/// never infers cache state from latency; only caller knowledge is recorded.
2801#[inline]
2802pub fn record_io_cache_state(state: crate::IoCacheStateV2) {
2803    record_annotation(crate::AnnotationId::IoCacheState, state.as_value());
2804    let counter = match state {
2805        crate::IoCacheStateV2::Cold => crate::CounterId::PageCacheColdObservations,
2806        crate::IoCacheStateV2::Warm => crate::CounterId::PageCacheWarmObservations,
2807        crate::IoCacheStateV2::Direct => crate::CounterId::PageCacheDirectObservations,
2808    };
2809    add_counter(counter, 1);
2810}
2811
2812/// Record the current retained-buffer level in bytes; the runtime keeps the
2813/// running high water alongside the latest value.
2814#[inline]
2815pub fn record_retained_buffer_bytes(bytes: u64) {
2816    if let Some(runtime) = current_runtime() {
2817        runtime.record_retained_buffer_bytes(bytes);
2818    }
2819}
2820
2821/// Drain typed counters from the current session or standalone runtime.
2822///
2823/// Under a [`Session`] this equals `Runtime::take_session_typed_metrics`;
2824/// under the standalone `set_enabled` runtime it drains the legacy store.
2825pub fn take_typed_metrics() -> Vec<TypedMetricRecordV2> {
2826    runtime_for_drain().take_legacy_typed_metrics()
2827}
2828
2829/// Drain caller-recorded value distributions from the current runtime.
2830///
2831/// Works under both a [`Session`] and the standalone `set_enabled` runtime.
2832pub fn take_metric_distributions() -> Vec<MetricDistributionV2> {
2833    runtime_for_drain().take_metric_distributions()
2834}
2835
2836/// Record one typed numeric annotation on the current run timeline.
2837#[inline]
2838pub fn record_annotation(annotation: crate::AnnotationId, value: u64) {
2839    if let Some(runtime) = current_runtime() {
2840        runtime.record_annotation(annotation, value);
2841    }
2842}
2843
2844/// Record the requested, selected, and completed route for one completed batch.
2845///
2846/// A recovered batch records the failed selected backend in
2847/// `recovered_from_backend` and the replay backend in `completed_backend`.
2848pub fn record_batch_route(
2849    workload_key_digest: &str,
2850    requested_backend: &str,
2851    selected_backend: &str,
2852    completed_backend: &str,
2853    recovered_from_backend: Option<&str>,
2854) {
2855    if let Some(runtime) = current_runtime() {
2856        runtime.record_batch_route(
2857            workload_key_digest,
2858            requested_backend,
2859            selected_backend,
2860            completed_backend,
2861            recovered_from_backend,
2862        );
2863    }
2864}
2865
2866/// Atomically read and clear aggregate input bytes and units.
2867pub fn take_input_totals() -> (u64, u64) {
2868    runtime_for_drain().take_input_totals()
2869}
2870
2871/// Discard fixed-stage counters and input totals in the current runtime.
2872pub fn reset() {
2873    runtime_for_drain().reset();
2874}