Fixes #270 - Add `Session::digest_mut()` method

Includes-commit: 5dd90aadd29b63564433ad8caccff7f4ab10eb5b
Replicated-from: https://github.com/cloudflare/pingora/pull/271

8 files changed, 46 insertions, 1 deletions

diff --git a/.bleep b/.bleep
index 700b347..10944e6 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-e8c45299beef05e26f59b620d998b98d31139aac
\ No newline at end of file
+04ff13f360fd33f2a508a47df6fbd2f6ca95ac89
\ No newline at end of file
diff --git a/pingora-core/src/connectors/http/v2.rs b/pingora-core/src/connectors/http/v2.rs
index 2c11ede..a082943 100644
--- a/pingora-core/src/connectors/http/v2.rs
+++ b/pingora-core/src/connectors/http/v2.rs
@@ -102,6 +102,10 @@ impl ConnectionRef {
         &self.0.digest
     }
 
+    pub fn digest_mut(&mut self) -> Option<&mut Digest> {
+        Arc::get_mut(&mut self.0).map(|inner| &mut inner.digest)
+    }
+
     pub fn ping_timedout(&self) -> bool {
         self.0.ping_timeout_occurred.load(Ordering::Relaxed)
     }
diff --git a/pingora-core/src/protocols/http/client.rs b/pingora-core/src/protocols/http/client.rs
index ed73152..af7a614 100644
--- a/pingora-core/src/protocols/http/client.rs
+++ b/pingora-core/src/protocols/http/client.rs
@@ -161,6 +161,17 @@ impl HttpSession {
         }
     }
 
+    /// Return a mutable [Digest] reference for the connection
+    ///
+    /// For reused connection, the timing in the digest will reflect its initial handshakes
+    /// The caller should check if the connection is reused to avoid misuse of the timing field
+    pub fn digest_mut(&mut self) -> Option<&mut Digest> {
+        match self {
+            Self::H1(s) => Some(s.digest_mut()),
+            Self::H2(s) => s.digest_mut(),
+        }
+    }
+
     /// Return the server (peer) address of the connection.
     pub fn server_addr(&self) -> Option<&SocketAddr> {
         match self {
diff --git a/pingora-core/src/protocols/http/server.rs b/pingora-core/src/protocols/http/server.rs
index 78c1856..6e702b1 100644
--- a/pingora-core/src/protocols/http/server.rs
+++ b/pingora-core/src/protocols/http/server.rs
@@ -395,6 +395,14 @@ impl Session {
         }
     }
 
+    /// Return a mutable digest reference for the session.
+    pub fn digest_mut(&mut self) -> Option<&mut Digest> {
+        match self {
+            Self::H1(s) => Some(s.digest_mut()),
+            Self::H2(s) => s.digest_mut(),
+        }
+    }
+
     /// Return the client (peer) address of the connection.
     pub fn client_addr(&self) -> Option<&SocketAddr> {
         match self {
diff --git a/pingora-core/src/protocols/http/v1/client.rs b/pingora-core/src/protocols/http/v1/client.rs
index 7a72ed4..c567d78 100644
--- a/pingora-core/src/protocols/http/v1/client.rs
+++ b/pingora-core/src/protocols/http/v1/client.rs
@@ -635,6 +635,10 @@ impl HttpSession {
         &self.digest
     }
 
+    pub fn digest_mut(&mut self) -> &mut Digest {
+        &mut self.digest
+    }
+
     /// Return the server (peer) address recorded in the connection digest.
     pub fn server_addr(&self) -> Option<&SocketAddr> {
         self.digest()
diff --git a/pingora-core/src/protocols/http/v1/server.rs b/pingora-core/src/protocols/http/v1/server.rs
index f545b6b..82a93a2 100644
--- a/pingora-core/src/protocols/http/v1/server.rs
+++ b/pingora-core/src/protocols/http/v1/server.rs
@@ -829,6 +829,11 @@ impl HttpSession {
         &self.digest
     }
 
+    /// Return a mutable [Digest] reference for the connection.
+    pub fn digest_mut(&mut self) -> &mut Digest {
+        &mut self.digest
+    }
+
     /// Return the client (peer) address of the underlying connection.
     pub fn client_addr(&self) -> Option<&SocketAddr> {
         self.digest()
diff --git a/pingora-core/src/protocols/http/v2/client.rs b/pingora-core/src/protocols/http/v2/client.rs
index 15f86d0..c639165 100644
--- a/pingora-core/src/protocols/http/v2/client.rs
+++ b/pingora-core/src/protocols/http/v2/client.rs
@@ -310,6 +310,14 @@ impl Http2Session {
         Some(self.conn.digest())
     }
 
+    /// Return a mutable [Digest] reference for the connection
+    ///
+    /// For reused connection, the timing in the digest will reflect its initial handshakes
+    /// The caller should check if the connection is reused to avoid misuse the timing field
+    pub fn digest_mut(&mut self) -> Option<&mut Digest> {
+        self.conn.digest_mut()
+    }
+
     /// Return the server (peer) address recorded in the connection digest.
     pub fn server_addr(&self) -> Option<&SocketAddr> {
         self.conn
diff --git a/pingora-core/src/protocols/http/v2/server.rs b/pingora-core/src/protocols/http/v2/server.rs
index 57238d1..f9072d6 100644
--- a/pingora-core/src/protocols/http/v2/server.rs
+++ b/pingora-core/src/protocols/http/v2/server.rs
@@ -454,6 +454,11 @@ impl HttpSession {
         Some(&self.digest)
     }
 
+    /// Return a mutable [Digest] reference for the connection.
+    pub fn digest_mut(&mut self) -> Option<&mut Digest> {
+        Arc::get_mut(&mut self.digest)
+    }
+
     /// Return the server (local) address recorded in the connection digest.
     pub fn server_addr(&self) -> Option<&SocketAddr> {
         self.digest.socket_digest.as_ref().map(|d| d.local_addr())?