From 8b777797646b2163844aa2425ee68fcee067bc5e Mon Sep 17 00:00:00 2001
From: Matt Schulte <mschulte@cloudflare.com>
Date: Fri, 29 Jul 2022 16:39:59 -0700
Subject: [PATCH] Fix reserve over allocating underlying buffer

Fixes calls to `reserve` when the underlying shared buffer was already
big enough to fit the requested capacity. Previously a new even larger
buffer was created anyways. This could eventually lead to an OOM
condition.
---
 src/bytes_mut.rs    |  5 ++++-
 tests/test_bytes.rs | 28 ++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs
index 99e849716..2282bdc48 100644
--- a/src/bytes_mut.rs
+++ b/src/bytes_mut.rs
@@ -670,7 +670,10 @@ impl BytesMut {
 
                 // Compare the condition in the `kind == KIND_VEC` case above
                 // for more details.
-                if v_capacity >= new_cap && offset >= len {
+                if v_capacity >= new_cap + offset {
+                    self.cap = new_cap;
+                    // no copy is necessary
+                } else if v_capacity >= new_cap && offset >= len {
                     // The capacity is sufficient, and copying is not too much
                     // overhead: reclaim the buffer!
 
diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs
index 6bf01d67d..4ddb24de5 100644
--- a/tests/test_bytes.rs
+++ b/tests/test_bytes.rs
@@ -515,6 +515,34 @@ fn reserve_in_arc_unique_doubles() {
     assert_eq!(2000, bytes.capacity());
 }
 
+#[test]
+fn reserve_in_arc_unique_does_not_overallocate_after_split() {
+    let mut bytes = BytesMut::from(LONG);
+    let orig_capacity = bytes.capacity();
+    drop(bytes.split_off(LONG.len() / 2));
+
+    // now bytes is Arc and refcount == 1
+
+    let new_capacity = bytes.capacity();
+    bytes.reserve(orig_capacity - new_capacity);
+    assert_eq!(bytes.capacity(), orig_capacity);
+}
+
+#[test]
+fn reserve_in_arc_unique_does_not_overallocate_after_multiple_splits() {
+    let mut bytes = BytesMut::from(LONG);
+    let orig_capacity = bytes.capacity();
+    for _ in 0..10 {
+        drop(bytes.split_off(LONG.len() / 2));
+
+        // now bytes is Arc and refcount == 1
+
+        let new_capacity = bytes.capacity();
+        bytes.reserve(orig_capacity - new_capacity);
+    }
+    assert_eq!(bytes.capacity(), orig_capacity);
+}
+
 #[test]
 fn reserve_in_arc_nonunique_does_not_overallocate() {
     let mut bytes = BytesMut::with_capacity(1000);