feat: add minitrace-futures to support instrumentation for futures

Cargo.toml

@@ -6,6 +6,7 @@ members = [
  "minitrace-jaeger",
  "minitrace-datadog",
  "minitrace-opentelemetry",
+ "minitrace-futures",
  "test-statically-disable",
 ]
 

minitrace-futures/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "minitrace-futures"
+version = "0.6.3"
+edition = "2021"
+license = "Apache-2.0"
+description = "Utilities for instrumenting `futures` with minitrace-rust"
+homepage = "https://github.com/tikv/minitrace-rust"
+repository = "https://github.com/tikv/minitrace-rust"
+documentation = "https://docs.rs/minitrace-futures"
+readme = "README.md"
+categories = ["development-tools::debugging"]
+keywords = ["tracing", "span", "datadog", "jaeger", "opentelemetry"]
+
+[dependencies]
+futures = "0.3"
+minitrace = { version = "0.6.3", path = "../minitrace" }
+pin-project-lite = "0.2.13"
+
+[dev-dependencies]
+async-stream = "0.3.5"
+tokio = { version = "1", features = ["rt", "time", "macros"] }

minitrace-futures/README.md

@@ -0,0 +1,7 @@
+# minitrace-futures
+
+[![Documentation](https://docs.rs/minitrace-futures/badge.svg)](https://docs.rs/minitrace-futures/)
+[![Crates.io](https://img.shields.io/crates/v/minitrace-futures.svg)](https://crates.io/crates/minitrace-futures)
+[![LICENSE](https://img.shields.io/github/license/tikv/minitrace-rust.svg)](https://github.com/tikv/minitrace-rust/blob/master/LICENSE)
+
+Some utilities for instrumenting `futures` with [`minitrace`](https://crates.io/crates/minitrace).

minitrace-futures/src/lib.rs

@@ -0,0 +1,160 @@
+#![doc = include_str!("../README.md")]
+
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+use futures::Sink;
+use futures::Stream;
+use minitrace::Span;
+use pin_project_lite::pin_project;
+
+// There is no boundary in order to support types that
+// implement both stream and sink.
+impl<T> Instrumented for T {}
+
+/// Instrument [`futures::Stream`]s and [`futures::Sink`]s.
+pub trait Instrumented: Sized {
+ /// For [`Stream`]s :
+ ///
+ /// Binds a [`Span`] to the [`Stream`] that continues to record until the stream is **finished**.
+ ///
+ /// In addition, it sets the span as the local parent at every poll so that [`minitrace::local::LocalSpan`]
+ /// becomes available within the future. Internally, it calls [`Span::set_local_parent`] when
+ /// the executor polls it.
+ ///
+ /// # Examples:
+ ///
+ /// ```
+ /// # #[tokio::main]
+ /// # async fn main() {
+ /// use async_stream::stream;
+ /// use futures::StreamExt;
+ /// use minitrace::prelude::*;
+ /// use minitrace_futures::*;
+ ///
+ /// let root = Span::root("Root", SpanContext::random());
+ /// let s = stream! {
+ /// for i in 0..2 {
+ /// yield i;
+ /// }
+ /// }
+ /// .in_span(Span::enter_with_parent("Task", &root));
+ ///
+ /// tokio::pin!(s);
+ ///
+ /// assert_eq!(s.next().await.unwrap(), 0);
+ /// assert_eq!(s.next().await.unwrap(), 1);
+ /// assert_eq!(s.next().await, None);
+ /// // span ends here.
+ ///
+ /// # }
+ /// ```
+ ///
+ /// For [`Sink`]s :
+ ///
+ /// Binds a [`Span`] to the [`Sink`] that continues to record until the sink is **closed**.
+ ///
+ /// In addition, it sets the span as the local parent at every poll so that [`minitrace::local::LocalSpan`]
+ /// becomes available within the future. Internally, it calls [`Span::set_local_parent`] when
+ /// the executor polls it.
+ ///
+ /// # Examples:
+ ///
+ /// ```
+ /// # #[tokio::main]
+ /// # async fn main() {
+ /// use futures::sink;
+ /// use futures::sink::SinkExt;
+ /// use minitrace::prelude::*;
+ /// use minitrace_futures::*;
+ ///
+ /// let root = Span::root("Root", SpanContext::random());
+ ///
+ /// let mut drain = sink::drain().in_span(Span::enter_with_parent("Task", &root));
+ ///
+ /// drain.send(1).await.unwrap();
+ /// drain.send(2).await.unwrap();
+ /// drain.close().await.unwrap();
+ /// // span ends here.
+ ///
+ /// # }
+ /// ```
+ fn in_span(self, span: Span) -> InSpan<Self> {
+ InSpan {
+ inner: self,
+ span: Some(span),
+ }
+ }
+}
+
+pin_project! {
+ pub struct InSpan<T> {
+ #[pin]
+ inner: T,
+ span: Option<Span>,
+ }
+}
+
+impl<T> Stream for InSpan<T>
+where T: Stream
+{
+ type Item = T::Item;
+
+ fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ let this = self.project();
+
+ let _guard = this.span.as_ref().map(|s| s.set_local_parent());
+ let res = this.inner.poll_next(cx);
+
+ match res {
+ r @ Poll::Pending => r,
+ r @ Poll::Ready(None) => {
+ // finished
+ this.span.take();
+ r
+ }
+ other => other,
+ }
+ }
+}
+
+impl<T, I> Sink<I> for InSpan<T>
+where T: Sink<I>
+{
+ type Error = T::Error;
+
+ fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+ let this = self.project();
+ let _guard = this.span.as_ref().map(|s| s.set_local_parent());
+ this.inner.poll_ready(cx)
+ }
+
+ fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
+ let this = self.project();
+ let _guard = this.span.as_ref().map(|s| s.set_local_parent());
+ this.inner.start_send(item)
+ }
+
+ fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+ let this = self.project();
+ let _guard = this.span.as_ref().map(|s| s.set_local_parent());
+ this.inner.poll_flush(cx)
+ }
+
+ fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+ let this = self.project();
+
+ let _guard = this.span.as_ref().map(|s| s.set_local_parent());
+ let res = this.inner.poll_close(cx);
+
+ match res {
+ r @ Poll::Pending => r,
+ other => {
+ // closed
+ this.span.take();
+ other
+ }
+ }
+ }
+}