chore: Add an explicit +Unpin to our buffers

lib/vector-core/buffers/src/disk/mod.rs

@@ -80,7 +80,7 @@ pub fn open<'a, T>(
 ) -> Result<
     (
         Writer<T>,
-        Box<dyn Stream<Item = T> + 'a + Send>,
+        Box<dyn Stream<Item = T> + 'a + Unpin + Send>,
         super::Acker,
     ),
     DataDirError,

lib/vector-core/buffers/src/lib.rs

@@ -66,7 +66,7 @@ pub fn build<'a, T>(
 ) -> Result<
     (
         BufferInputCloner<T>,
-        Box<dyn Stream<Item = T> + 'a + Send>,
+        Box<dyn Stream<Item = T> + 'a + Unpin + Send>,
         Acker,
     ),
     String,

src/buffers.rs

@@ -32,6 +32,8 @@ impl Default for BufferConfig {
     }
 }
 
+pub(crate) type EventStream = Box<dyn Stream<Item = Event> + Unpin + Send>;
+
 impl BufferConfig {
     #[inline]
     const fn memory_max_events() -> usize {
@@ -43,14 +45,7 @@ impl BufferConfig {
         &self,
         data_dir: &Option<PathBuf>,
         sink_name: &str,
-    ) -> Result<
-        (
-            BufferInputCloner<Event>,
-            Box<dyn Stream<Item = Event> + Send>,
-            Acker,
-        ),
-        String,
-    > {
+    ) -> Result<(BufferInputCloner<Event>, EventStream, Acker), String> {
         let variant = match &self {
             BufferConfig::Memory {
                 max_events,

src/topology/builder.rs

@@ -13,6 +13,7 @@ use crate::{
     Pipeline,
 };
 use futures::{future, stream, FutureExt, SinkExt, StreamExt, TryFutureExt};
+use std::pin::Pin;
 use std::{
     collections::HashMap,
     future::ready,
@@ -119,9 +120,13 @@ pub async fn build_pieces(
             Ok(transform) => transform,
         };
 
-        let (input_tx, input_rx) = futures::channel::mpsc::channel(100);
-        let input_tx = buffers::BufferInputCloner::Memory(input_tx, buffers::WhenFull::Block);
-        let input_rx = crate::utilization::wrap(input_rx);
+        let (input_tx, input_rx, _) =
+            vector_core::buffers::build(vector_core::buffers::Variant::Memory {
+                max_events: 100,
+                when_full: vector_core::buffers::WhenFull::Block,
+            })
+            .unwrap();
+        let input_rx = crate::utilization::wrap(Pin::new(input_rx));
 
         let (output, control) = Fanout::new();
 

src/topology/mod.rs

@@ -11,7 +11,7 @@ pub mod fanout;
 mod task;
 
 use crate::{
-    buffers,
+    buffers::{self, EventStream},
     config::{Config, ConfigDiff, HealthcheckOptions, Resource},
     event::Event,
     shutdown::SourceShutdownCoordinator,
@@ -21,7 +21,7 @@ use crate::{
     },
     trigger::DisabledTrigger,
 };
-use futures::{future, Future, FutureExt, SinkExt, Stream};
+use futures::{future, Future, FutureExt, SinkExt};
 use std::{
     collections::{HashMap, HashSet},
     panic::AssertUnwindSafe,
@@ -38,7 +38,7 @@ type TaskHandle = tokio::task::JoinHandle<Result<TaskOutput, ()>>;
 
 type BuiltBuffer = (
     buffers::BufferInputCloner<Event>,
-    Arc<Mutex<Option<Pin<Box<dyn Stream<Item = Event> + Send>>>>>,
+    Arc<Mutex<Option<Pin<EventStream>>>>,
     buffers::Acker,
 );
 

src/topology/task.rs

@@ -1,5 +1,5 @@
-use crate::{buffers::Acker, event::Event};
-use futures::{future::BoxFuture, FutureExt, Stream};
+use crate::buffers::{Acker, EventStream};
+use futures::{future::BoxFuture, FutureExt};
 use pin_project::pin_project;
 use std::{
     fmt,
@@ -12,7 +12,7 @@ pub enum TaskOutput {
     Source,
     Transform,
     /// Buffer of sink
-    Sink(Pin<Box<dyn Stream<Item = Event> + Send>>, Acker),
+    Sink(Pin<EventStream>, Acker),
     Healthcheck,
 }
 

src/utilization.rs

@@ -1,40 +1,72 @@
-use crate::{event::Event, stats};
-use async_stream::stream;
+use crate::{buffers::EventStream, event::Event, stats};
 use futures::{Stream, StreamExt};
-use std::time::{Duration, Instant};
+use pin_project::pin_project;
+use std::{pin::Pin, task::Context};
+use std::{
+    task::Poll,
+    time::{Duration, Instant},
+};
+use tokio::time::interval;
+use tokio_stream::wrappers::IntervalStream;
 
-/// Wrap a stream to emit stats about utilization. This is designed for use with the input channels
-/// of transform and sinks components, and measures the amount of time that the stream is waiting
-/// for input from upstream. We make the simplifying assumption that this wait time is when the
-/// component is idle and the rest of the time it is doing useful work. This is more true for sinks
-/// than transforms, which can be blocked by downstream components, but with knowledge of the
-/// config the data is still useful.
-pub fn wrap(inner: impl Stream<Item = Event>) -> impl Stream<Item = Event> {
-    let mut timer = Timer::new();
-    let mut interval = tokio::time::interval(Duration::from_secs(5));
-    stream! {
-        tokio::pin!(inner);
+#[pin_project]
+struct Utilization {
+    timer: Timer,
+    intervals: IntervalStream,
+    inner: Pin<EventStream>,
+}
+
+impl Stream for Utilization {
+    type Item = Event;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        // The goal of this function is to measure the time between when the
+        // caller requests the next Event from the stream and before one is
+        // ready, with the side-effect of reporting every so often about how
+        // long the wait gap is.
+        //
+        // To achieve this we poll the `intervals` stream and if a new interval
+        // is ready we hit `Timer::report` and loop back around again to poll
+        // for a new `Event`. Calls to `Timer::start_wait` will only have an
+        // effect if `stop_wait` has been called, so the structure of this loop
+        // avoids double-measures.
+        let this = self.project();
         loop {
-            timer.start_wait();
-            let value = tokio::select! {
-                value = inner.next() => {
-                    timer.stop_wait();
-                    value
-                },
-                _ = interval.tick() => {
-                    timer.report();
-                    continue
+            this.timer.start_wait();
+            match this.intervals.poll_next_unpin(cx) {
+                Poll::Ready(_) => {
+                    this.timer.report();
+                    continue;
                 }
-            };
-            if let Some(value) = value {
-                yield value
-            } else {
-                break
+                Poll::Pending => match this.inner.poll_next_unpin(cx) {
+                    pe @ Poll::Ready(_) => {
+                        this.timer.stop_wait();
+                        return pe;
+                    }
+                    Poll::Pending => return Poll::Pending,
+                },
             }
         }
     }
 }
 
+/// Wrap a stream to emit stats about utilization. This is designed for use with
+/// the input channels of transform and sinks components, and measures the
+/// amount of time that the stream is waiting for input from upstream. We make
+/// the simplifying assumption that this wait time is when the component is idle
+/// and the rest of the time it is doing useful work. This is more true for
+/// sinks than transforms, which can be blocked by downstream components, but
+/// with knowledge of the config the data is still useful.
+pub fn wrap(inner: Pin<EventStream>) -> Pin<EventStream> {
+    let utilization = Utilization {
+        timer: Timer::new(),
+        intervals: IntervalStream::new(interval(Duration::from_secs(5))),
+        inner,
+    };
+
+    Box::pin(utilization)
+}
+
 struct Timer {
     overall_start: Instant,
     span_start: Instant,
@@ -43,13 +75,14 @@ struct Timer {
     ewma: stats::Ewma,
 }
 
-/// A simple, specialized timer for tracking spans of waiting vs not-waiting time and reporting
-/// a smoothed estimate of utilization.
+/// A simple, specialized timer for tracking spans of waiting vs not-waiting
+/// time and reporting a smoothed estimate of utilization.
 ///
-/// This implementation uses the idea of spans and reporting periods. Spans are a period of time
-/// spent entirely in one state, aligning with state transitions but potentially more granular.
-/// Reporting periods are expected to be of uniform length and used to aggregate span data into
-/// time-weighted averages.
+/// This implementation uses the idea of spans and reporting periods. Spans are
+/// a period of time spent entirely in one state, aligning with state
+/// transitions but potentially more granular.  Reporting periods are expected
+/// to be of uniform length and used to aggregate span data into time-weighted
+/// averages.
 impl Timer {
     fn new() -> Self {
         Self {
@@ -63,8 +96,10 @@ impl Timer {
 
     /// Begin a new span representing time spent waiting
     fn start_wait(&mut self) {
-        self.end_span();
-        self.waiting = true;
+        if !self.waiting {
+            self.end_span();
+            self.waiting = true;
+        }
     }
 
     /// Complete the current waiting span and begin a non-waiting span
@@ -75,11 +110,13 @@ impl Timer {
         self.waiting = false;
     }
 
-    /// Meant to be called on a regular interval, this method calculates wait ratio  since the last
-    /// time it was called and reports the resulting utilization average.
+    /// Meant to be called on a regular interval, this method calculates wait
+    /// ratio since the last time it was called and reports the resulting
+    /// utilization average.
     fn report(&mut self) {
-        // End the current span so it can be accounted for, but do not change whether or not we're
-        // in the waiting state. This way the next span inherits the correct status.
+        // End the current span so it can be accounted for, but do not change
+        // whether or not we're in the waiting state. This way the next span
+        // inherits the correct status.
         self.end_span();
 
         let total_duration = self.overall_start.elapsed();