Repository: tokio-rs/bytes (Rust)

Issue:
BytesMut memory usage grows without bound
I'm working on a high-throughput network service with Tokio, and I'm running into an issue where the memory used by a `BytesMut` grows without bound as more and more messages are processed. After some debugging, I was able to come up with this minimal example:

    extern crate bytes;
    use bytes::*;

    fn leaks_memory() {
        let mut data = BytesMut::from(Vec::with_capacity(16));

        loop {
            data.reserve(4);
            data.put(&b"leak"[..]);
            data.advance(4);
        }
    }

    fn also_leaks_memory() {
        let mut data = BytesMut::with_capacity(8192);

        loop {
            data.reserve(4);
            data.put(&b"leak"[..]);
            data.advance(4);
        }
    }

    fn constant_memory() {
        let mut data = BytesMut::from(Vec::with_capacity(16));

        loop {
            data.reserve(4);
            data.put(&b"leak"[..]);

            data = data.split_off(4);
        }
    }

    fn also_constant_memory() {
        let mut data = BytesMut::with_capacity(16);

        loop {
            data.reserve(4);
            data.put(&b"leak"[..]);
            data.advance(4);
        }
    }

It looks like what's happening is that a `BytesMut` backed by a `Vec` isn't actually reusing space at the beginning of the buffer when `reserve` is called, but only if it was moved forwards with `advance`. I'm not sure whether any of the other methods (e.g. `truncate`) cause the problem. The documentation says `reserve` will try to recycle existing space (which it's clearly doing in `also_constant_memory`) but it seems to be failing when it shouldn't.


Discussion:
I see what is going on. This is kind of the intended behavior as `advance` though I can see how the docs could be confusing.

That said, the case you are illustrating could be optimized by "resetting" the cursors when the read cursor and the write cursor line up.

Do you want to take a stab at implementing that?
The `advance` implementation actually looks fine to me; it doesn't look like it's explicitly tracking read/write cursors, just advancing the current pointer. I did implement a fix for the issue though. In `reserve`, when `kind == KIND_VEC`, instead of just calling `Vec::reserve` I check whether there's free space at the beginning of the buffer and copy everything back to the start if so.

With the fix implemented, the test cases I showed all run in constant memory. Does it seem like a reasonable approach to you? If so, I can open a PR.



Patch:
diff --git a/src/bytes.rs b/src/bytes.rs
index 60cf5d41d..c85ebf32f 100644
--- a/src/bytes.rs
+++ b/src/bytes.rs
@@ -2250,20 +2250,42 @@ impl Inner {
         }
 
         if kind == KIND_VEC {
-            // Currently backed by a vector, so just use `Vector::reserve`.
+            // If there's enough free space before the start of the buffer, then
+            // just copy the data backwards and reuse the already-allocated
+            // space.
+            //
+            // Otherwise, since backed by a vector, use `Vec::reserve`
             unsafe {
-                let (off, _) = self.uncoordinated_get_vec_pos();
-                let mut v = rebuild_vec(self.ptr, self.len, self.cap, off);
-                v.reserve(additional);
-
-                // Update the info
-                self.ptr = v.as_mut_ptr().offset(off as isize);
-                self.len = v.len() - off;
-                self.cap = v.capacity() - off;
+                let (off, prev) = self.uncoordinated_get_vec_pos();
+
+                // Only reuse space if we stand to gain at least capacity/2
+                // bytes of space back
+                if off >= additional && off >= (self.cap / 2) {
+                    // There's space - reuse it
+                    //
+                    // Just move the pointer back to the start after copying
+                    // data back.
+                    let base_ptr = self.ptr.offset(-(off as isize));
+                    ptr::copy(self.ptr, base_ptr, self.len);
+                    self.ptr = base_ptr;
+                    self.uncoordinated_set_vec_pos(0, prev);
+
+                    // Length stays constant, but since we moved backwards we
+                    // can gain capacity back.
+                    self.cap += off;
+                } else {
+                    // No space - allocate more
+                    let mut v = rebuild_vec(self.ptr, self.len, self.cap, off);
+                    v.reserve(additional);
 
-                // Drop the vec reference
-                mem::forget(v);
+                    // Update the info
+                    self.ptr = v.as_mut_ptr().offset(off as isize);
+                    self.len = v.len() - off;
+                    self.cap = v.capacity() - off;
 
+                    // Drop the vec reference
+                    mem::forget(v);
+                }
                 return;
             }
         }


Test patch:
diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs
index c75de1f20..4ab86699c 100644
--- a/tests/test_bytes.rs
+++ b/tests/test_bytes.rs
@@ -378,6 +378,21 @@ fn reserve_max_original_capacity_value() {
     assert_eq!(bytes.capacity(), 64 * 1024);
 }
 
+// Without either looking at the internals of the BytesMut or doing weird stuff
+// with the memory allocator, there's no good way to automatically verify from
+// within the program that this actually recycles memory. Instead, just exercise
+// the code path to ensure that the results are correct.
+#[test]
+fn reserve_vec_recycling() {
+    let mut bytes = BytesMut::from(Vec::with_capacity(16));
+    assert_eq!(bytes.capacity(), 16);
+    bytes.put("0123456789012345");
+    bytes.advance(10);
+    assert_eq!(bytes.capacity(), 6);
+    bytes.reserve(8);
+    assert_eq!(bytes.capacity(), 16);
+}
+
 #[test]
 fn reserve_in_arc_unique_does_not_overallocate() {
     let mut bytes = BytesMut::with_capacity(1000);