Propagates the tracing `Span` down to the iterators returned by functions

instrument.rs

use tracing::Span;

/// Iterator adapter that enters a span on each iteration.
pub struct InstrumentedIter<I> {
    iter: I,
    span: Span,
}

impl<I> Iterator for InstrumentedIter<I>
where
    I: Iterator,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        let _enter = self.span.enter();
        self.iter.next()
    }
}

/// Extension trait for instrumenting iterators with a tracing span.
pub trait IteratorInstrumentExt: Iterator + Sized {
    fn instrumented(self, span: Span) -> InstrumentedIter<Self> {
        InstrumentedIter { iter: self, span }
    }
}

impl<I: Iterator> IteratorInstrumentExt for I {}