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 { + Instrumented { + inner: self, + span + } + } +} + +#[derive(Clone, Debug)] +pub struct Instrumented { + inner: T, + span: Span, +} + +impl Instrument for T {} +impl Instrument for T {} +impl Instrument for T {} + +impl Future for Instrument { + type Item = T::Item; + type Error = T::Error; + + fn poll(&mut self) -> Poll { + self.span.enter(|| { + self.inner.poll() + }) + } +} + +impl Stream for Instrument { + type Item = T::Item; + type Error = T::Error; + + fn poll(&mut self) -> Poll, Self::Error> { + self.span.enter(|| { + self.inner.poll() + }) + } +} + +impl Sink for Instrument { + type SinkItem = T::SinkItem; + type SinkError = T::SinkError; + + fn start_send( + &mut self, + item: Self::SinkItem + ) -> StartSend { + self.span.enter(|| self.inner.start_send(item)) + } + + fn poll_complete(&mut self) -> Poll<(), Self::SinkError> { + self.span.enter(|| { self.inner.poll_complete() }) + } +}