From 6a1bd055fcac0fcc1834f888a81a276fca384332 Mon Sep 17 00:00:00 2001
From: Sean McArthur <sean@seanmonstar.com>
Date: Tue, 3 Mar 2020 14:23:49 -0800
Subject: [PATCH] refactor(http2): store bdp sampler in Body H2 variant

---
 Cargo.toml          |  2 +-
 src/body/body.rs    | 19 +++++--------------
 src/proto/h2/bdp.rs |  6 +-----
 3 files changed, 7 insertions(+), 20 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6879e8a3..e1de6b24 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -27,7 +27,7 @@ futures-util = { version = "0.3", default-features = false }
 http = "0.2"
 http-body = "0.3.1"
 httparse = "1.0"
-h2 = "0.2.1"
+h2 = "0.2.2"
 itoa = "0.4.1"
 log = "0.4"
 pin-project = "0.4"
diff --git a/src/body/body.rs b/src/body/body.rs
index 339b77d6..228a996c 100644
--- a/src/body/body.rs
+++ b/src/body/body.rs
@@ -38,6 +38,7 @@ enum Kind {
         rx: mpsc::Receiver<Result<Bytes, crate::Error>>,
     },
     H2 {
+        bdp: bdp::Sampler,
         content_length: DecodedLength,
         recv: h2::RecvStream,
     },
@@ -62,11 +63,6 @@ struct Extra {
     /// connection yet.
     delayed_eof: Option<DelayEof>,
     on_upgrade: OnUpgrade,
-
-    /// Records bytes read to compute the BDP.
-    ///
-    /// Only used with `H2` variant.
-    h2_bdp: bdp::Sampler,
 }
 
 type DelayEofUntil = oneshot::Receiver<Never>;
@@ -186,15 +182,12 @@ impl Body {
         content_length: DecodedLength,
         bdp: bdp::Sampler,
     ) -> Self {
-        let mut body = Body::new(Kind::H2 {
+        let body = Body::new(Kind::H2 {
+            bdp,
             content_length,
             recv,
         });
 
-        if bdp.is_enabled() {
-            body.extra_mut().h2_bdp = bdp;
-        }
-
         body
     }
 
@@ -220,7 +213,6 @@ impl Body {
             Box::new(Extra {
                 delayed_eof: None,
                 on_upgrade: OnUpgrade::none(),
-                h2_bdp: bdp::disabled(),
             })
         })
     }
@@ -273,15 +265,14 @@ impl Body {
                 }
             }
             Kind::H2 {
+                ref bdp,
                 recv: ref mut h2,
                 content_length: ref mut len,
             } => match ready!(h2.poll_data(cx)) {
                 Some(Ok(bytes)) => {
                     let _ = h2.flow_control().release_capacity(bytes.len());
                     len.sub_if(bytes.len() as u64);
-                    if let Some(ref extra) = self.extra {
-                        extra.h2_bdp.sample(bytes.len());
-                    }
+                    bdp.sample(bytes.len());
                     Poll::Ready(Some(Ok(bytes)))
                 }
                 Some(Err(e)) => Poll::Ready(Some(Err(crate::Error::new_body(e)))),
diff --git a/src/proto/h2/bdp.rs b/src/proto/h2/bdp.rs
index 611eb7ac..2a2c99bc 100644
--- a/src/proto/h2/bdp.rs
+++ b/src/proto/h2/bdp.rs
@@ -30,7 +30,7 @@ type WindowSize = u32;
 /// Any higher than this likely will be hitting the TCP flow control.
 const BDP_LIMIT: usize = 1024 * 1024 * 16;
 
-pub(crate) fn disabled() -> Sampler {
+pub(super) fn disabled() -> Sampler {
     Sampler {
         shared: Weak::new(),
     }
@@ -105,10 +105,6 @@ impl Sampler {
 
         inner.bytes += bytes;
     }
-
-    pub(crate) fn is_enabled(&self) -> bool {
-        self.shared.strong_count() > 0
-    }
 }
 
 impl Estimator {