refactor: reduce mem alloc during framed reading

black-binary · black-binary · commit b45dc506f4ce · 2023-12-26T12:35:45.000-05:00
diff --git a/src/frame.rs b/src/frame.rs
@@ -95,6 +95,8 @@ impl Decoder for MuxCodec {
     type Error = MuxError;
 
     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
+        src.reserve(HEADER_SIZE + MAX_PAYLOAD_SIZE + HEADER_SIZE);
+
         if src.len() < HEADER_SIZE {
             return Ok(None);
         }
diff --git a/src/lib.rs b/src/lib.rs
@@ -317,4 +317,29 @@ mod tests {
 
         assert!(stream2.write_all(b"1234").await.is_err());
     }
+
+    #[tokio::test]
+    async fn test_inner_shutdown() {
+        let (a, b) = get_tcp_pair().await;
+
+        let (connector_a, mut acceptor_a, worker_a) =
+            MuxBuilder::client().with_connection(a).build();
+        let (connector_b, mut acceptor_b, worker_b) =
+            MuxBuilder::server().with_connection(b).build();
+
+        let a_res = tokio::spawn(worker_a);
+        drop(worker_b);
+        tokio::time::sleep(Duration::from_secs(2)).await;
+
+        assert!(connector_b.connect().is_err());
+        assert!(acceptor_b.accept().await.is_none());
+
+        drop(connector_b);
+        drop(acceptor_b);
+
+        tokio::time::sleep(Duration::from_secs(2)).await;
+        assert!(connector_a.connect().is_err());
+        assert!(acceptor_a.accept().await.is_none());
+        a_res.await.unwrap().unwrap_err();
+    }
 }
diff --git a/src/mux.rs b/src/mux.rs
@@ -626,10 +626,7 @@ impl<T: TokioConn> MuxState<T> {
         self.check_closed()?;
 
         if let Some(r) = ready!(self.inner.poll_next_unpin(cx)) {
-            let frame = r.map_err(|e| {
-                self.close();
-                e
-            })?;
+            let frame = r.inspect_err(|_| self.close())?;
             Poll::Ready(Ok(frame))
         } else {
             self.close();

Original file line number	Diff line number	Diff line change
`@@ -95,6 +95,8 @@ impl Decoder for MuxCodec {`
`95`	`95`	`type Error = MuxError;`
`96`	`96`
`97`	`97`	`fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {`
	`98`	`+ src.reserve(HEADER_SIZE + MAX_PAYLOAD_SIZE + HEADER_SIZE);`
	`99`	`+`
`98`	`100`	`if src.len() < HEADER_SIZE {`
`99`	`101`	`return Ok(None);`
`100`	`102`	`}`