http-streams-core 0.2.0

Framework-neutral HTTP body streaming formats: JSON/CSV/Protobuf/Arrow encode and decode
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
//! Accounting and observability for one streamed body, in either direction.
//!
//! This is the merge of two implementations that had independently converged on the same
//! design (`axum-streams`' `progress.rs` and `reqwest-streams`' `observability.rs`), and it
//! takes the union of their behaviour. Both dependents drive it; neither owns a copy.
//!
//! Unlike an earlier design, the tracing callsites live **here** rather than in the binding
//! crates. `tracing` bakes a span's name and target into a `static Metadata`, so they cannot be
//! passed in at runtime; keeping them here means one target, `http_streams_core`, with the
//! direction carried as a span field instead.

use crate::error::StreamError;
use bytes::Bytes;
use futures::stream::{Stream, TryStreamExt};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

/// Reported about once a second unless overridden.
pub const DEFAULT_PROGRESS_INTERVAL: Duration = Duration::from_secs(1);

/// Which HTTP message the body belongs to.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum Direction {
    /// A request body: uploaded by a client, received by a server.
    Request,
    /// A response body: produced by a server, read by a client.
    #[default]
    Response,
}

impl Direction {
    /// A short, stable name, reported as the `direction` tracing field.
    pub fn as_str(&self) -> &'static str {
        match self {
            Direction::Request => "request",
            Direction::Response => "response",
        }
    }
}

/// Which end of the connection this code is running on.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum Side {
    /// Running in an HTTP client.
    #[default]
    Client,
    /// Running in an HTTP server.
    Server,
}

impl Side {
    /// A short, stable name, reported as the `side` tracing field.
    pub fn as_str(&self) -> &'static str {
        match self {
            Side::Client => "client",
            Side::Server => "server",
        }
    }
}

/// How a stream ended.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum StreamOutcome {
    /// Still running; reported by interim progress only.
    InProgress,
    /// Reached the end of the body with no errors.
    Completed,
    /// Ended early because the other side of the stream went away.
    Aborted,
    /// Reached its end, but at least one error was reported along the way.
    Failed,
}

impl StreamOutcome {
    /// A short, stable name, reported as the `outcome` tracing field.
    pub fn as_str(&self) -> &'static str {
        match self {
            StreamOutcome::InProgress => "in_progress",
            StreamOutcome::Completed => "completed",
            StreamOutcome::Aborted => "aborted",
            StreamOutcome::Failed => "failed",
        }
    }
}

/// A snapshot of one stream's accounting.
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub struct StreamProgress {
    /// Items encoded or decoded so far.
    pub items: u64,
    /// Body bytes transferred so far.
    pub bytes: u64,
    /// Errors reported so far. Not every error is terminal.
    pub errors: u64,
    /// Time since the stream was created.
    pub elapsed: Duration,
    /// How the stream ended, or [`InProgress`] if it has not.
    ///
    /// [`InProgress`]: StreamOutcome::InProgress
    pub outcome: StreamOutcome,
}

/// Called for every error reported by a stream.
pub type StreamErrorHandler = Arc<dyn Fn(&StreamError) + Send + Sync + 'static>;

/// Called for every progress report, interim and terminal.
pub type StreamProgressHandler = Arc<dyn Fn(&StreamProgress) + Send + Sync + 'static>;

/// The direction-neutral subset of the dependents' options structs.
///
/// Both `StreamBodyAsOptions` and `ReqwestStreamOptions` keep their own definitions, because they
/// carry direction-specific fields, and neither could add inherent builder methods to a type
/// defined here (E0116). Each builds one of these internally instead.
#[derive(Clone)]
#[non_exhaustive]
pub struct ProgressOptions {
    /// Invoked for every error reported by the stream.
    pub on_error: Option<StreamErrorHandler>,
    /// Invoked for every progress report.
    pub on_progress: Option<StreamProgressHandler>,
    /// How often to report interim progress. `None` disables the time trigger.
    pub progress_interval: Option<Duration>,
    /// Report interim progress every N items. `None` disables the item trigger.
    pub progress_items: Option<u64>,
}

impl ProgressOptions {
    /// Default options: report about once a second, no item trigger, no callbacks.
    pub fn new() -> Self {
        Self {
            on_error: None,
            on_progress: None,
            progress_interval: Some(DEFAULT_PROGRESS_INTERVAL),
            progress_items: None,
        }
    }

    /// Set the error callback.
    pub fn on_error(mut self, handler: StreamErrorHandler) -> Self {
        self.on_error = Some(handler);
        self
    }

    /// Set the progress callback.
    pub fn on_progress(mut self, handler: StreamProgressHandler) -> Self {
        self.on_progress = Some(handler);
        self
    }

    /// Set the interim reporting interval.
    pub fn progress_interval(mut self, interval: Duration) -> Self {
        self.progress_interval = Some(interval);
        self
    }

    /// Report interim progress every `items` items.
    pub fn progress_items(mut self, items: u64) -> Self {
        self.progress_items = Some(items);
        self
    }
}

impl Default for ProgressOptions {
    fn default() -> Self {
        Self::new()
    }
}

/// What a stream is about, for the span covering it.
///
/// Fields the caller cannot know are simply left `None` and are omitted from the span rather
/// than reported as a placeholder.
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct StreamContext {
    /// The format's [`format_name`](crate::StreamFormat::format_name).
    ///
    /// `Cow` rather than `&'static str` because `axum-streams`' long-standing public
    /// `StreamingFormat::format_name` returns a borrowed `&str`, and that signature cannot
    /// change without breaking the third-party implementations it was made public for.
    /// Core's own formats all return `&'static str` and so allocate nothing.
    pub format: std::borrow::Cow<'static, str>,
    /// Request or response body.
    pub direction: Direction,
    /// Client or server.
    pub side: Side,
    /// Response status, where one is known.
    pub status: Option<u16>,
    /// Declared body length, where one is known.
    pub content_length: Option<u64>,
    /// The negotiated content type, where one is known.
    pub content_type: Option<String>,
    /// The per-object limit in force, where one applies.
    pub max_obj_len: Option<usize>,
    /// The read-buffer size in force, where one applies.
    pub buf_capacity: Option<usize>,
}

impl StreamContext {
    /// A context for `format`, in the given direction and on the given side.
    pub fn new(
        format: impl Into<std::borrow::Cow<'static, str>>,
        direction: Direction,
        side: Side,
    ) -> Self {
        Self {
            format: format.into(),
            direction,
            side,
            ..Default::default()
        }
    }

    /// Record the response status.
    pub fn status(mut self, status: u16) -> Self {
        self.status = Some(status);
        self
    }

    /// Record the declared body length.
    pub fn content_length(mut self, len: Option<u64>) -> Self {
        self.content_length = len;
        self
    }

    /// Record the content type.
    pub fn content_type(mut self, ct: impl Into<String>) -> Self {
        self.content_type = Some(ct.into());
        self
    }

    /// Record the per-object limit.
    pub fn max_obj_len(mut self, len: usize) -> Self {
        self.max_obj_len = Some(len);
        self
    }

    /// Record the read-buffer size.
    pub fn buf_capacity(mut self, cap: usize) -> Self {
        self.buf_capacity = Some(cap);
        self
    }
}

/// What the accounting needs to know about an error passing through a stream.
///
/// Deliberately not `&StreamError`: a binding crate's pipeline carries that crate's own error
/// type, `axum::Error` say, and it could not implement a trait from here for it anyway,
/// since both the trait and the type would be foreign to it (orphan rule). Everything the
/// accounting actually needs is a `Display` and, when available, the typed error.
pub struct ErrorInfo<'a> {
    display: &'a dyn std::fmt::Display,
    stream_error: Option<&'a StreamError>,
}

impl<'a> ErrorInfo<'a> {
    /// The error's `Display`, for the `error` tracing field.
    pub fn display(&self) -> &dyn std::fmt::Display {
        self.display
    }

    /// The typed error, when this stream carries [`StreamError`]s.
    pub fn stream_error(&self) -> Option<&'a StreamError> {
        self.stream_error
    }

    /// A short, stable kind name for the `error_kind` tracing field.
    pub fn kind_str(&self) -> &'static str {
        self.stream_error.map_or("unknown", |e| e.kind().as_str())
    }
}

/// Lets [`instrument`] classify items without being generic over the item type.
///
/// A `T` that appeared only in a `where` clause would be an unconstrained type parameter
/// (E0207), and a `PhantomData<T>` would drag `T`'s auto traits into the stream's type, which
/// would break any binding whose `T` carries no `Send + 'b` bound even though the method it
/// backs promises a `Send` stream. `reqwest-streams`' CSV reader is one such.
pub trait ProgressItem {
    /// The error this item carries, if it is one.
    fn progress_error(&self) -> Option<ErrorInfo<'_>>;
}

/// Covers every `Result` whose error is a standard error, so binding crates get this for their
/// own error types without implementing anything.
impl<T, E> ProgressItem for Result<T, E>
where
    E: std::error::Error + 'static,
{
    fn progress_error(&self) -> Option<ErrorInfo<'_>> {
        self.as_ref().err().map(|err| ErrorInfo {
            display: err,
            // Resolves at compile time for any concrete `E`, so this costs nothing when the
            // stream does not carry `StreamError`s.
            stream_error: (err as &dyn std::any::Any).downcast_ref::<StreamError>(),
        })
    }
}

/// Checked at `ERROR`, the least verbose level the accounting can produce: a failed stream
/// reports there, so gating any higher would mean `RUST_LOG=http_streams_core=error` silently
/// loses the totals of the very streams it asked about. Every more verbose filter enables
/// `ERROR` too, so this can never suppress wanted output.
#[cfg(feature = "tracing")]
fn tracing_enabled() -> bool {
    tracing::enabled!(target: "http_streams_core", tracing::Level::ERROR)
}

#[cfg(not(feature = "tracing"))]
fn tracing_enabled() -> bool {
    false
}

/// Shared accounting for one streamed body.
///
/// Bytes are counted on the byte stream and items on the item stream, so the two counters live
/// in different combinators and share this state. The ordering is `Relaxed` throughout: these
/// are counters, not synchronisation.
struct ProgressState {
    items: AtomicU64,
    bytes: AtomicU64,
    errors: AtomicU64,
    last_emit_micros: AtomicU64,
    next_item_step: AtomicU64,
    polled: AtomicBool,
    finalized: AtomicBool,
    start: Instant,
    interval_micros: Option<u64>,
    item_step: Option<u64>,
    on_error: Option<StreamErrorHandler>,
    on_progress: Option<StreamProgressHandler>,
    #[cfg(feature = "tracing")]
    span: tracing::Span,
}

impl ProgressState {
    /// Returns `None` when nobody is listening, in which case every accounting call below
    /// short-circuits on a single `Option` check.
    #[cfg_attr(not(feature = "tracing"), allow(unused_variables))]
    fn maybe_new(context: &StreamContext, options: &ProgressOptions) -> Option<Arc<Self>> {
        if options.on_progress.is_none() && options.on_error.is_none() && !tracing_enabled() {
            return None;
        }

        // A step of zero would never advance, so treat it as "disabled" rather than looping.
        let item_step = options.progress_items.filter(|step| *step > 0);

        Some(Arc::new(Self {
            items: AtomicU64::new(0),
            bytes: AtomicU64::new(0),
            errors: AtomicU64::new(0),
            last_emit_micros: AtomicU64::new(0),
            next_item_step: AtomicU64::new(item_step.unwrap_or(u64::MAX)),
            polled: AtomicBool::new(false),
            finalized: AtomicBool::new(false),
            start: Instant::now(),
            interval_micros: options
                .progress_interval
                .map(|interval| interval.as_micros() as u64),
            item_step,
            on_error: options.on_error.clone(),
            on_progress: options.on_progress.clone(),
            #[cfg(feature = "tracing")]
            span: Self::new_span(context),
        }))
    }

    /// The span covering the whole stream, created by the caller while its own span is still
    /// current, so collectors nest it under their request rather than orphaning it. The stream
    /// itself is polled later, potentially from an entirely different task.
    ///
    /// Every counter is declared up front as an empty field so it can be filled in later with
    /// [`tracing::Span::record`]: collectors that read span attributes (OpenTelemetry and
    /// friends) then see `items`/`bytes`/`outcome` as structured values on a span whose
    /// duration is the streaming duration, instead of having to parse log messages.
    ///
    /// No URL is recorded, deliberately: it carries query strings and userinfo, which
    /// routinely means presigned-URL signatures and `?api_key=`.
    #[cfg(feature = "tracing")]
    fn new_span(context: &StreamContext) -> tracing::Span {
        let span = tracing::info_span!(
            target: "http_streams_core",
            "http_streams_core::stream",
            format = context.format.as_ref(),
            direction = context.direction.as_str(),
            side = context.side.as_str(),
            // `Option` is a `Value` that simply skips the field when it is empty.
            status = context.status,
            content_length = context.content_length,
            content_type = context.content_type.as_deref(),
            max_obj_len = tracing::field::Empty,
            buf_capacity = tracing::field::Empty,
            items = tracing::field::Empty,
            bytes = tracing::field::Empty,
            errors = tracing::field::Empty,
            elapsed_ms = tracing::field::Empty,
            outcome = tracing::field::Empty,
        );

        // `usize::MAX` means "no limit", which is noise rather than information.
        if let Some(max) = context.max_obj_len.filter(|m| *m != usize::MAX) {
            span.record("max_obj_len", max as u64);
        }
        if let Some(cap) = context.buf_capacity {
            span.record("buf_capacity", cap as u64);
        }

        span
    }

    fn record_bytes(&self, len: u64) {
        let bytes = self.bytes.fetch_add(len, Ordering::Relaxed) + len;
        let items = self.items.load(Ordering::Relaxed);

        #[cfg(feature = "tracing")]
        tracing::trace!(
            target: "http_streams_core",
            parent: &self.span,
            chunk_bytes = len,
            items,
            bytes,
            "Transferred an HTTP body chunk"
        );

        // Progress is driven from transferred bytes as well as from items, because a single
        // item can take a long time: one large Arrow batch, or a JSON array streamed slowly,
        // would otherwise report nothing at all until it completed. Emitting resets the
        // interval, so a chunk and an item cannot both report for the same tick.
        if !self.finalized.load(Ordering::Relaxed) && self.should_emit_elapsed() {
            self.emit(
                StreamOutcome::InProgress,
                items,
                bytes,
                self.errors.load(Ordering::Relaxed),
            );
        }
    }

    fn record_item(&self) {
        let items = self.items.fetch_add(1, Ordering::Relaxed) + 1;

        // Nothing may be reported after the summary, or the final snapshot would no longer be
        // final. A consumer is free to keep polling a stream past its end.
        if !self.finalized.load(Ordering::Relaxed) && self.should_emit_items(items) {
            self.emit(
                StreamOutcome::InProgress,
                items,
                self.bytes.load(Ordering::Relaxed),
                self.errors.load(Ordering::Relaxed),
            );
        }
    }

    /// Errors are reported as they happen but are deliberately **not** terminal.
    ///
    /// Only some of them are: `FramedRead` latches its own error state and ends the stream,
    /// but the JSON Lines and CSV formats produce their decoding errors from a successfully
    /// framed line, and the stream carries on to the next one. Finalising here would stop
    /// counting the remaining items of a stream that is still perfectly healthy, so the
    /// terminal outcome is decided at the end instead, from this counter.
    fn record_error(&self, info: &ErrorInfo<'_>) {
        self.errors.fetch_add(1, Ordering::Relaxed);

        #[cfg(feature = "tracing")]
        tracing::error!(
            target: "http_streams_core",
            parent: &self.span,
            error = %info.display(),
            error_kind = info.kind_str(),
            "An error occurred while streaming an HTTP body"
        );

        // Only fires for streams that carry `StreamError`s. A binding whose pipeline uses its
        // own error type reports through its own typed callback instead, so that the closure
        // its users already wrote keeps compiling.
        if let (Some(handler), Some(err)) = (&self.on_error, info.stream_error()) {
            handler(err);
        }
    }

    /// The time trigger, checked when bytes move.
    ///
    /// The two triggers are deliberately driven from *different* signals rather than both
    /// being checked everywhere. Bytes are what keeps flowing regardless of how the payload
    /// divides into items (a single large item, or a slow one, still reports), so elapsed time
    /// is checked here. Items drive the item-step trigger below.
    ///
    /// Checking both from both places double-reports: every pipeline counts the same payload
    /// twice, once as items and once as bytes, so a zero interval would emit two events per
    /// unit rather than one.
    fn should_emit_elapsed(&self) -> bool {
        let Some(interval) = self.interval_micros else {
            return false;
        };

        let elapsed = self.start.elapsed().as_micros() as u64;
        let since_last = elapsed.saturating_sub(self.last_emit_micros.load(Ordering::Relaxed));
        if since_last >= interval {
            self.last_emit_micros.store(elapsed, Ordering::Relaxed);
            return true;
        }

        false
    }

    /// The item-step trigger, checked when items are counted.
    fn should_emit_items(&self, items: u64) -> bool {
        let Some(step) = self.item_step else {
            return false;
        };

        if items < self.next_item_step.load(Ordering::Relaxed) {
            return false;
        }

        // Skip past every step the current count already crossed, so a single poll carrying
        // many items cannot queue up a burst of events.
        self.next_item_step
            .store(items - (items % step) + step, Ordering::Relaxed);

        // Emitting resets the time trigger as well, so a step and a tick that fall together
        // produce one event rather than two.
        self.last_emit_micros
            .store(self.start.elapsed().as_micros() as u64, Ordering::Relaxed);

        true
    }

    fn mark_polled(&self) {
        self.polled.store(true, Ordering::Relaxed);
    }

    /// Emits the terminal snapshot, exactly once per stream.
    ///
    /// A stream that was never polled reports nothing at all. Building one and dropping it
    /// unconsumed is routine (a `?` short-circuits, a handler returns early), and reporting
    /// those as aborted would bury the real ones in `items=0 bytes=0` noise.
    fn finalize(&self, aborted: bool) {
        if !self.polled.load(Ordering::Relaxed) || self.finalized.swap(true, Ordering::Relaxed) {
            return;
        }

        let items = self.items.load(Ordering::Relaxed);
        let bytes = self.bytes.load(Ordering::Relaxed);
        let errors = self.errors.load(Ordering::Relaxed);

        let outcome = if errors > 0 {
            StreamOutcome::Failed
        } else if aborted {
            StreamOutcome::Aborted
        } else {
            StreamOutcome::Completed
        };

        // Recorded once, here rather than on every progress report: subscribers are free to
        // treat `record` as append-only (`tracing-subscriber`'s formatter does), so writing a
        // field repeatedly makes the rendered span grow with every tick. Once per span also
        // means the values a collector reads are the final ones.
        #[cfg(feature = "tracing")]
        {
            self.span.record("items", items);
            self.span.record("bytes", bytes);
            self.span.record("errors", errors);
            self.span
                .record("elapsed_ms", self.start.elapsed().as_millis() as u64);
            self.span.record("outcome", outcome.as_str());
        }

        self.emit(outcome, items, bytes, errors);
    }

    fn emit(&self, outcome: StreamOutcome, items: u64, bytes: u64, errors: u64) {
        let progress = StreamProgress {
            items,
            bytes,
            errors,
            elapsed: self.start.elapsed(),
            outcome,
        };

        #[cfg(feature = "tracing")]
        {
            let elapsed_ms = progress.elapsed.as_millis() as u64;

            match outcome {
                // Interim progress is chatter; the summary is the line worth keeping, and a
                // truncated stream is worth an operator's attention.
                StreamOutcome::InProgress => tracing::debug!(
                    target: "http_streams_core",
                    parent: &self.span,
                    items,
                    bytes,
                    elapsed_ms,
                    "Streaming an HTTP body"
                ),
                StreamOutcome::Failed => tracing::error!(
                    target: "http_streams_core",
                    parent: &self.span,
                    items,
                    bytes,
                    errors,
                    elapsed_ms,
                    outcome = outcome.as_str(),
                    "Failed streaming an HTTP body"
                ),
                // Completed, and aborted: an end that stops early is ordinary.
                _ => tracing::info!(
                    target: "http_streams_core",
                    parent: &self.span,
                    items,
                    bytes,
                    errors,
                    elapsed_ms,
                    outcome = outcome.as_str(),
                    "Finished streaming an HTTP body"
                ),
            }
        }

        if let Some(handler) = &self.on_progress {
            handler(&progress);
        }
    }
}

/// The accounting handle threaded through one stream's pipeline.
///
/// `None` inside means nobody is listening and every method is a no-op. Cheap to clone.
#[derive(Clone)]
pub struct Progress(Option<Arc<ProgressState>>);

impl Progress {
    /// Build a handle for one stream.
    ///
    /// Call this while the caller's own tracing span is still current, and before the body is
    /// consumed, so the span nests correctly and the context fields are still available.
    pub fn new(context: &StreamContext, options: &ProgressOptions) -> Self {
        Progress(ProgressState::maybe_new(context, options))
    }

    /// A handle that reports nothing.
    pub fn disabled() -> Self {
        Progress(None)
    }

    /// Whether anything is listening. Useful to skip work that only feeds reporting.
    pub fn is_enabled(&self) -> bool {
        self.0.is_some()
    }

    /// Count one item.
    ///
    /// A no-op when nobody is listening. Provided for bindings that cannot use
    /// [`count_items`] because their pipeline's item type carries no `'static`-ish bound:
    /// applying a combinator would force one onto their public signature.
    pub fn record_item(&self) {
        if let Some(state) = &self.0 {
            state.record_item();
        }
    }

    /// Count `len` transferred bytes.
    ///
    /// A no-op when nobody is listening. The combinator form is [`count_bytes`].
    pub fn record_bytes(&self, len: u64) {
        if let Some(state) = &self.0 {
            state.record_bytes(len);
        }
    }
}

/// Whether [`instrument`] should count the items it sees.
///
/// The outermost stream of an encode pipeline yields `Bytes` chunks, not items, so counting
/// its `Ok`s would report a JSON array's `[` and `]` as items. In that case items are counted
/// upstream by [`count_items`] and this is set to [`Bytes`].
///
/// [`Bytes`]: Counting::Bytes
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Counting {
    /// Each `Ok` is one item.
    Items,
    /// Each `Ok` is a byte chunk; items are counted elsewhere.
    Bytes,
}

/// Counts the bytes flowing through a byte stream.
///
/// Applied to the byte stream rather than the item stream so that `bytes` is what actually
/// crossed the wire, independently of how many objects that turned into.
pub fn count_bytes<'b, S, E>(
    stream: S,
    progress: &Progress,
) -> impl Stream<Item = Result<Bytes, E>> + Send + 'b
where
    S: Stream<Item = Result<Bytes, E>> + Send + 'b,
    E: 'b,
{
    let progress = progress.clone();
    stream.inspect_ok(move |chunk| {
        if let Some(state) = &progress.0 {
            state.record_bytes(chunk.len() as u64);
        }
    })
}

/// Counts items flowing through an item stream, without touching errors or the outcome.
///
/// Used on the encode side, where items exist only upstream of the encoder. This is the one
/// place they are still items rather than bytes.
pub fn count_items<'b, S, T, E>(
    stream: S,
    progress: &Progress,
) -> impl Stream<Item = Result<T, E>> + Send + 'b
where
    S: Stream<Item = Result<T, E>> + Send + 'b,
    // Deliberately no `T: Send`. The *stream* must be `Send`; its item type need not be, and
    // requiring it would break callers whose `T` carries no such bound. Same reasoning as
    // [`ProgressItem`].
    T: 'b,
    E: 'b,
{
    let progress = progress.clone();
    stream.inspect_ok(move |_| {
        if let Some(state) = &progress.0 {
            state.record_item();
        }
    })
}

/// Reports errors, owns the outcome state machine, and optionally counts items.
///
/// Wrap the **outermost** stream with this: every error passes through there, and its `Drop`
/// is the only way to notice an end that stopped early.
pub fn instrument<'b, S>(
    stream: S,
    progress: Progress,
    counting: Counting,
) -> impl Stream<Item = S::Item> + Send + 'b
where
    S: Stream + Unpin + Send + 'b,
    S::Item: ProgressItem,
{
    ProgressStream {
        inner: stream,
        progress,
        counting,
    }
}

struct ProgressStream<S> {
    inner: S,
    progress: Progress,
    counting: Counting,
}

impl<S> Stream for ProgressStream<S>
where
    S: Stream + Unpin,
    S::Item: ProgressItem,
{
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // Safe without any projection: `Self: Unpin` whenever `S: Unpin`, which is exactly the
        // bound above. That keeps `#![forbid(unsafe_code)]` intact.
        let this = self.get_mut();

        // Borrowed, not cloned: `progress` and `inner` are disjoint fields, so this avoids an
        // atomic refcount bump on every single poll.
        let Some(state) = this.progress.0.as_ref() else {
            return Pin::new(&mut this.inner).poll_next(cx);
        };

        // Polling here drives the whole pipeline synchronously, so entering the span gives
        // everything it touches the stream's context. `poll_next` is synchronous, so this
        // guard is never held across an await.
        #[cfg(feature = "tracing")]
        let _entered = state.span.enter();

        state.mark_polled();

        match Pin::new(&mut this.inner).poll_next(cx) {
            Poll::Ready(Some(item)) => {
                match item.progress_error() {
                    Some(info) => state.record_error(&info),
                    None => {
                        if this.counting == Counting::Items {
                            state.record_item();
                        }
                    }
                }
                Poll::Ready(Some(item))
            }
            Poll::Ready(None) => {
                state.finalize(false);
                Poll::Ready(None)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

impl<S> Drop for ProgressStream<S> {
    fn drop(&mut self) {
        if let Some(state) = &self.progress.0 {
            // A no-op when the stream already ran to completion, or was never polled.
            state.finalize(true);
        }
    }
}