Move metrics::Io to io::SensorIo (#610)

This change decouples the I/O stream wrapper from the metrics
implementation so that it can implement the Io trait (so that it can be
boxed).

Author: olix0r

Cargo.lock

@@ -1122,6 +1122,7 @@ version = "0.1.0"
 dependencies = [
  "bytes 0.5.4",
  "futures 0.3.5",
+ "linkerd2-errno",
  "pin-project",
  "tokio",
  "tokio-rustls",

linkerd/app/core/src/proxy/server.rs

@@ -144,7 +144,7 @@ where
 impl<L, F, H, B> Service<Connection> for Server<L, F, H, B>
 where
     L: TransportLabels<Protocol, Labels = TransportKey>,
-    F: Accept<(tls::accept::Meta, transport::metrics::Io<BoxedIo>)> + Clone + Send + 'static,
+    F: Accept<(tls::accept::Meta, transport::metrics::SensorIo<BoxedIo>)> + Clone + Send + 'static,
     F::Future: Send + 'static,
     F::ConnectionFuture: Send + 'static,
     H: NewService<tls::accept::Meta> + Send + 'static,

linkerd/io/Cargo.toml

@@ -11,6 +11,7 @@ General I/O primitives.
 [dependencies]
 futures = "0.3"
 bytes = "0.5"
+linkerd2-errno = { path = "../errno" }
 tokio = { version = "0.2", features = ["io-util", "net", "macros"] }
 tokio-rustls = "0.13"
 pin-project = "0.4"

linkerd/io/src/lib.rs

@@ -1,10 +1,13 @@
 mod boxed;
 mod peek;
 mod prefixed;
+mod sensor;
+
 pub use self::{
     boxed::BoxedIo,
     peek::{Peek, Peekable},
     prefixed::PrefixedIo,
+    sensor::{Sensor, SensorIo},
 };
 pub use std::io::{Error, Read, Result, Write};
 pub use tokio::io::{AsyncRead, AsyncWrite};

linkerd/io/src/sensor.rs

@@ -0,0 +1,106 @@
+use crate::{internal::Io, Poll};
+use bytes::{Buf, BufMut};
+use futures::ready;
+use linkerd2_errno::Errno;
+use pin_project::pin_project;
+use std::mem::MaybeUninit;
+use std::pin::Pin;
+use std::task::Context;
+use tokio::io::{AsyncRead, AsyncWrite};
+
+pub trait Sensor {
+    fn record_read(&mut self, sz: usize);
+    fn record_write(&mut self, sz: usize);
+    fn record_close(&mut self, eos: Option<Errno>);
+    fn record_error<T>(&mut self, op: Poll<T>) -> Poll<T>;
+}
+
+/// Wraps a transport with telemetry.
+#[pin_project]
+#[derive(Debug)]
+pub struct SensorIo<T, S> {
+    #[pin]
+    io: T,
+
+    sensor: S,
+}
+
+// === impl SensorIo ===
+
+impl<T, S: Sensor> SensorIo<T, S> {
+    pub fn new(io: T, sensor: S) -> Self {
+        Self { io, sensor }
+    }
+}
+
+impl<T: AsyncRead + AsyncWrite, S: Sensor> AsyncRead for SensorIo<T, S> {
+    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8]) -> Poll<usize> {
+        let this = self.project();
+        let bytes = ready!(this.sensor.record_error(this.io.poll_read(cx, buf)))?;
+        this.sensor.record_read(bytes);
+        Poll::Ready(Ok(bytes))
+    }
+
+    fn poll_read_buf<B: BufMut>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut B,
+    ) -> Poll<usize> {
+        let this = self.project();
+        let bytes = ready!(this.sensor.record_error(this.io.poll_read_buf(cx, buf)))?;
+        this.sensor.record_read(bytes);
+        Poll::Ready(Ok(bytes))
+    }
+
+    unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [MaybeUninit<u8>]) -> bool {
+        self.io.prepare_uninitialized_buffer(buf)
+    }
+}
+
+impl<T: AsyncRead + AsyncWrite, S: Sensor> AsyncWrite for SensorIo<T, S> {
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        let this = self.project();
+        this.sensor.record_error(this.io.poll_shutdown(cx))
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
+        let this = self.project();
+        this.sensor.record_error(this.io.poll_flush(cx))
+    }
+
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<usize> {
+        let this = self.project();
+        let bytes = ready!(this.sensor.record_error(this.io.poll_write(cx, buf)))?;
+        this.sensor.record_write(bytes);
+        Poll::Ready(Ok(bytes))
+    }
+
+    fn poll_write_buf<B: Buf>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut B,
+    ) -> Poll<usize> {
+        let this = self.project();
+        let bytes = ready!(this.sensor.record_error(this.io.poll_write_buf(cx, buf)))?;
+        this.sensor.record_write(bytes);
+        Poll::Ready(Ok(bytes))
+    }
+}
+
+impl<T: Io, S: Sensor + Send> Io for SensorIo<T, S> {
+    fn poll_write_buf_erased(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        mut buf: &mut dyn Buf,
+    ) -> Poll<usize> {
+        self.poll_write_buf(cx, &mut buf)
+    }
+
+    fn poll_read_buf_erased(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        mut buf: &mut dyn BufMut,
+    ) -> Poll<usize> {
+        self.poll_read_buf(cx, &mut buf)
+    }
+}

linkerd/proxy/transport/src/metrics/mod.rs renamed to linkerd/proxy/transport/src/metrics.rs

@@ -1,6 +1,7 @@
 use futures::{ready, TryFuture};
 use indexmap::IndexMap;
 use linkerd2_errno::Errno;
+use linkerd2_io as io;
 use linkerd2_metrics::{
     latency, metrics, Counter, FmtLabels, FmtMetric, FmtMetrics, Gauge, Histogram, Metric,
 };
@@ -15,10 +16,6 @@ use std::time::Instant;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::debug;
 
-mod io;
-
-pub use self::io::Io;
-
 metrics! {
     tcp_open_total: Counter { "Total count of opened connections" },
     tcp_open_connections: Gauge { "Number of currently-open connections" },
@@ -68,9 +65,6 @@ pub struct Connecting<F> {
 }
 
 /// Stores a class of transport's metrics.
-///
-/// TODO We should probaby use AtomicUsize for most of these counters so that
-/// simple increments don't require a lock. Especially for read|write_bytes_total.
 #[derive(Debug, Default)]
 struct Metrics {
     open_total: Counter,
@@ -98,11 +92,13 @@ struct EosMetrics {
 
 /// Tracks the state of a single instance of `Io` throughout its lifetime.
 #[derive(Debug)]
-struct Sensor {
+pub struct Sensor {
     metrics: Option<Arc<Metrics>>,
     opened_at: Instant,
 }
 
+pub type SensorIo<T> = io::SensorIo<T, Sensor>;
+
 /// Lazily builds instances of `Sensor`.
 #[derive(Clone, Debug)]
 struct NewSensor(Arc<Metrics>);
@@ -182,14 +178,18 @@ impl<K: Eq + Hash + FmtLabels> Registry<K> {
         ConnectLayer::new(label, self.0.clone())
     }
 
-    pub fn wrap_server_transport<T: AsyncRead + AsyncWrite>(&self, labels: K, io: T) -> Io<T> {
+    pub fn wrap_server_transport<T: AsyncRead + AsyncWrite>(
+        &self,
+        labels: K,
+        io: T,
+    ) -> SensorIo<T> {
         let metrics = self
             .0
             .lock()
             .expect("metrics registry poisoned")
             .get_or_default(labels)
             .clone();
-        Io::new(io, Sensor::open(metrics))
+        SensorIo::new(io, Sensor::open(metrics))
     }
 }
 
@@ -239,7 +239,7 @@ where
     L: TransportLabels<T>,
     M: tower::make::MakeConnection<T>,
 {
-    type Response = Io<M::Connection>;
+    type Response = SensorIo<M::Connection>;
     type Error = M::Error;
     type Future = Connecting<M::Future>;
 
@@ -270,7 +270,7 @@ where
     F: TryFuture,
     F::Ok: AsyncRead + AsyncWrite,
 {
-    type Output = Result<Io<F::Ok>, F::Error>;
+    type Output = Result<SensorIo<F::Ok>, F::Error>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = self.project();
@@ -282,7 +282,7 @@ where
             .take()
             .expect("future must not be polled after ready")
             .new_sensor();
-        let t = Io::new(io, sensor);
+        let t = SensorIo::new(io, sensor);
         Poll::Ready(Ok(t))
     }
 }
@@ -321,28 +321,30 @@ impl<K: Eq + Hash + FmtLabels> FmtMetrics for Report<K> {
 // ===== impl Sensor =====
 
 impl Sensor {
-    pub fn open(metrics: Arc<Metrics>) -> Self {
+    fn open(metrics: Arc<Metrics>) -> Self {
         metrics.open_total.incr();
         metrics.open_connections.incr();
         Self {
             metrics: Some(metrics),
             opened_at: Instant::now(),
         }
     }
+}
 
-    pub fn record_read(&mut self, sz: usize) {
+impl io::Sensor for Sensor {
+    fn record_read(&mut self, sz: usize) {
         if let Some(ref m) = self.metrics {
             m.read_bytes_total.add(sz as u64);
         }
     }
 
-    pub fn record_write(&mut self, sz: usize) {
+    fn record_write(&mut self, sz: usize) {
         if let Some(ref m) = self.metrics {
             m.write_bytes_total.add(sz as u64);
         }
     }
 
-    pub fn record_close(&mut self, eos: Option<Errno>) {
+    fn record_close(&mut self, eos: Option<Errno>) {
         let duration = self.opened_at.elapsed();
         // When closed, the metrics structure is dropped so that no further
         // updates can occur (i.e. so that an additional close won't be recorded
@@ -361,7 +363,7 @@ impl Sensor {
     ///
     /// If the transport operation results in a non-recoverable error, record a
     /// transport closure.
-    fn sense_err<T>(&mut self, op: Poll<std::io::Result<T>>) -> Poll<std::io::Result<T>> {
+    fn record_error<T>(&mut self, op: Poll<std::io::Result<T>>) -> Poll<std::io::Result<T>> {
         match op {
             Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)),
             Poll::Ready(Err(e)) => {
@@ -379,7 +381,7 @@ impl Sensor {
 
 impl Drop for Sensor {
     fn drop(&mut self) {
-        self.record_close(None)
+        io::Sensor::record_close(self, None)
     }
 }