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};
pub const DEFAULT_PROGRESS_INTERVAL: Duration = Duration::from_secs(1);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum Direction {
Request,
#[default]
Response,
}
impl Direction {
pub fn as_str(&self) -> &'static str {
match self {
Direction::Request => "request",
Direction::Response => "response",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum Side {
#[default]
Client,
Server,
}
impl Side {
pub fn as_str(&self) -> &'static str {
match self {
Side::Client => "client",
Side::Server => "server",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum StreamOutcome {
InProgress,
Completed,
Aborted,
Failed,
}
impl StreamOutcome {
pub fn as_str(&self) -> &'static str {
match self {
StreamOutcome::InProgress => "in_progress",
StreamOutcome::Completed => "completed",
StreamOutcome::Aborted => "aborted",
StreamOutcome::Failed => "failed",
}
}
}
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub struct StreamProgress {
pub items: u64,
pub bytes: u64,
pub errors: u64,
pub elapsed: Duration,
pub outcome: StreamOutcome,
}
pub type StreamErrorHandler = Arc<dyn Fn(&StreamError) + Send + Sync + 'static>;
pub type StreamProgressHandler = Arc<dyn Fn(&StreamProgress) + Send + Sync + 'static>;
#[derive(Clone)]
#[non_exhaustive]
pub struct ProgressOptions {
pub on_error: Option<StreamErrorHandler>,
pub on_progress: Option<StreamProgressHandler>,
pub progress_interval: Option<Duration>,
pub progress_items: Option<u64>,
}
impl ProgressOptions {
pub fn new() -> Self {
Self {
on_error: None,
on_progress: None,
progress_interval: Some(DEFAULT_PROGRESS_INTERVAL),
progress_items: None,
}
}
pub fn on_error(mut self, handler: StreamErrorHandler) -> Self {
self.on_error = Some(handler);
self
}
pub fn on_progress(mut self, handler: StreamProgressHandler) -> Self {
self.on_progress = Some(handler);
self
}
pub fn progress_interval(mut self, interval: Duration) -> Self {
self.progress_interval = Some(interval);
self
}
pub fn progress_items(mut self, items: u64) -> Self {
self.progress_items = Some(items);
self
}
}
impl Default for ProgressOptions {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Default)]
#[non_exhaustive]
pub struct StreamContext {
pub format: std::borrow::Cow<'static, str>,
pub direction: Direction,
pub side: Side,
pub status: Option<u16>,
pub content_length: Option<u64>,
pub content_type: Option<String>,
pub max_obj_len: Option<usize>,
pub buf_capacity: Option<usize>,
}
impl StreamContext {
pub fn new(
format: impl Into<std::borrow::Cow<'static, str>>,
direction: Direction,
side: Side,
) -> Self {
Self {
format: format.into(),
direction,
side,
..Default::default()
}
}
pub fn status(mut self, status: u16) -> Self {
self.status = Some(status);
self
}
pub fn content_length(mut self, len: Option<u64>) -> Self {
self.content_length = len;
self
}
pub fn content_type(mut self, ct: impl Into<String>) -> Self {
self.content_type = Some(ct.into());
self
}
pub fn max_obj_len(mut self, len: usize) -> Self {
self.max_obj_len = Some(len);
self
}
pub fn buf_capacity(mut self, cap: usize) -> Self {
self.buf_capacity = Some(cap);
self
}
}
pub struct ErrorInfo<'a> {
display: &'a dyn std::fmt::Display,
stream_error: Option<&'a StreamError>,
}
impl<'a> ErrorInfo<'a> {
pub fn display(&self) -> &dyn std::fmt::Display {
self.display
}
pub fn stream_error(&self) -> Option<&'a StreamError> {
self.stream_error
}
pub fn kind_str(&self) -> &'static str {
self.stream_error.map_or("unknown", |e| e.kind().as_str())
}
}
pub trait ProgressItem {
fn progress_error(&self) -> Option<ErrorInfo<'_>>;
}
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,
stream_error: (err as &dyn std::any::Any).downcast_ref::<StreamError>(),
})
}
}
#[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
}
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 {
#[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;
}
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),
}))
}
#[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(),
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,
);
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"
);
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;
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),
);
}
}
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"
);
if let (Some(handler), Some(err)) = (&self.on_error, info.stream_error()) {
handler(err);
}
}
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
}
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;
}
self.next_item_step
.store(items - (items % step) + step, Ordering::Relaxed);
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);
}
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
};
#[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 {
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"
),
_ => 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);
}
}
}
#[derive(Clone)]
pub struct Progress(Option<Arc<ProgressState>>);
impl Progress {
pub fn new(context: &StreamContext, options: &ProgressOptions) -> Self {
Progress(ProgressState::maybe_new(context, options))
}
pub fn disabled() -> Self {
Progress(None)
}
pub fn is_enabled(&self) -> bool {
self.0.is_some()
}
pub fn record_item(&self) {
if let Some(state) = &self.0 {
state.record_item();
}
}
pub fn record_bytes(&self, len: u64) {
if let Some(state) = &self.0 {
state.record_bytes(len);
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Counting {
Items,
Bytes,
}
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);
}
})
}
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,
T: 'b,
E: 'b,
{
let progress = progress.clone();
stream.inspect_ok(move |_| {
if let Some(state) = &progress.0 {
state.record_item();
}
})
}
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>> {
let this = self.get_mut();
let Some(state) = this.progress.0.as_ref() else {
return Pin::new(&mut this.inner).poll_next(cx);
};
#[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 {
state.finalize(true);
}
}
}