From 13e32604398deb7731d90f20f6c255d856176d02 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Fri, 21 Sep 2018 14:41:29 -0700
Subject: [PATCH] Futures, Streams, and Sinks can now be instrumented with
 spans

---
 tokio-trace/src/instrument.rs | 60 +++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 tokio-trace/src/instrument.rs

diff --git a/tokio-trace/src/instrument.rs b/tokio-trace/src/instrument.rs
new file mode 100644
index 00000000..e4254f48
--- /dev/null
+++ b/tokio-trace/src/instrument.rs
@@ -0,0 +1,60 @@
+use super::Span;
+use futures::{Future, Sink, Stream, Poll};
+
+
+pub trait Instrument: Sized {
+    fn instrument(self, span: Span) -> Instrumented<Self> {
+        Instrumented {
+            inner: self,
+            span
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Instrumented<T> {
+    inner: T,
+    span: Span,
+}
+
+impl<T: Future + Sized> Instrument for T {}
+impl<T: Stream + Sized> Instrument for T {}
+impl<T: Sink + Sized> Instrument for T {}
+
+impl<T: Future> Future for Instrument<T> {
+    type Item = T::Item;
+    type Error = T::Error;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        self.span.enter(|| {
+            self.inner.poll()
+        })
+    }
+}
+
+impl<T: Stream> Stream for Instrument<T> {
+    type Item = T::Item;
+    type Error = T::Error;
+
+    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        self.span.enter(|| {
+            self.inner.poll()
+        })
+    }
+}
+
+impl<T: Sink> Sink for Instrument<T> {
+    type SinkItem = T::SinkItem;
+    type SinkError = T::SinkError;
+
+    fn start_send(
+        &mut self,
+        item: Self::SinkItem
+    ) -> StartSend<Self::SinkItem, Self::SinkError> {
+        self.span.enter(|| self.inner.start_send(item))
+    }
+
+    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
+        self.span.enter(|| { self.inner.poll_complete() })
+    }
+}