From 1e1816676c97d7d7cb904018b8a45a8106010bbb Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki@neon.tech>
Date: Thu, 26 Jan 2023 00:34:49 +0200
Subject: [PATCH] WIP: Revert the MeasuredStream changes

To demonstrate why the MeasuredStream change makes things simpler.

This doesn't compile. Not that I tried very hard, but I couldn't
figure out how to fill in the generic type parameters to make this
compile.
---
 proxy/src/proxy.rs  | 10 ++++++++--
 proxy/src/stream.rs | 20 ++++++++++----------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs
index 696ef07814..e6284a58f2 100644
--- a/proxy/src/proxy.rs
+++ b/proxy/src/proxy.rs
@@ -347,10 +347,16 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> EstablishedConnection<'_, S> {
             .await?;
 
         let m_sent = NUM_BYTES_PROXIED_COUNTER.with_label_values(&node.aux.traffic_labels("tx"));
-        let client_stream = MeasuredStream::new(stream.into_inner(), m_sent);
+        let client_stream = MeasuredStream::new(stream.into_inner(), |cnt| {
+            // Number of bytes we sent to the client (outbound).
+            m_sent.inc_by(cnt as u64);
+        });
 
         let m_recv = NUM_BYTES_PROXIED_COUNTER.with_label_values(&node.aux.traffic_labels("rx"));
-        let db_stream = MeasuredStream::new(db.stream, m_recv);
+        let db_stream = MeasuredStream::new(db.stream, |cnt| {
+            // Number of bytes the client sent to the compute node (inbound).
+            m_recv.inc_by(cnt as u64);
+        });
 
         Ok(EstablishedConnection {
             client_stream,
diff --git a/proxy/src/stream.rs b/proxy/src/stream.rs
index 5729b61987..f6e7bd15f4 100644
--- a/proxy/src/stream.rs
+++ b/proxy/src/stream.rs
@@ -231,27 +231,27 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for Stream<S> {
 }
 
 pin_project! {
-    /// This stream tracks all writes, and whenever the stream is flushed,
-    /// increments the user-provided counter by the number of bytes flushed.
-    pub struct MeasuredStream<S> {
+    /// This stream tracks all writes and calls user provided
+    /// callback when the underlying stream is flushed.
+    pub struct MeasuredStream<S, W> {
         #[pin]
         stream: S,
         write_count: usize,
-        write_counter: prometheus::IntCounter,
+        inc_write_count: W,
     }
 }
 
-impl<S> MeasuredStream<S> {
-    pub fn new(stream: S, write_counter: prometheus::IntCounter) -> Self {
+impl<S, W> MeasuredStream<S, W> {
+    pub fn new(stream: S, inc_write_count: W) -> Self {
         Self {
             stream,
             write_count: 0,
-            write_counter,
+            inc_write_count,
         }
     }
 }
 
-impl<S: AsyncRead + Unpin> AsyncRead for MeasuredStream<S> {
+impl<S: AsyncRead + Unpin, W> AsyncRead for MeasuredStream<S, W> {
     fn poll_read(
         self: Pin<&mut Self>,
         context: &mut task::Context<'_>,
@@ -261,7 +261,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for MeasuredStream<S> {
     }
 }
 
-impl<S: AsyncWrite + Unpin> AsyncWrite for MeasuredStream<S> {
+impl<S: AsyncWrite + Unpin, W: FnMut(usize)> AsyncWrite for MeasuredStream<S, W> {
     fn poll_write(
         self: Pin<&mut Self>,
         context: &mut task::Context<'_>,
@@ -282,7 +282,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for MeasuredStream<S> {
         let this = self.project();
         this.stream.poll_flush(context).map_ok(|()| {
             // Call the user provided callback and reset the write count.
-            this.write_counter.inc_by(*this.write_count as u64);
+            (this.inc_write_count)(*this.write_count);
             *this.write_count = 0;
         })
     }