THRIFT-6431: Amortize Ruby MemoryBuffer compaction and release exhausted storage (#4020)

Client: rb

Compact unread storage only after consuming at least 4 KiB and at least
as many bytes as remain. Detach copied suffixes and replace exhausted
storage, including after partial reads.

Cache the next compaction boundary in pure Ruby to keep ordinary byte
reads inexpensive. Cover rolling reads, input ownership, resets and
exhaustion, and update diagnostic expectations for released storage.

Co-authored-by: OpenAI Codex (gpt-6.1-sol) <codex@openai.com>

Author: kpumuk

lib/rb/ext/memory_buffer.c

@@ -40,6 +40,21 @@ VALUE rb_thrift_memory_buffer_read_all(VALUE self, VALUE length_value);
 VALUE rb_thrift_memory_buffer_read_byte(VALUE self);
 VALUE rb_thrift_memory_buffer_read_into_buffer(VALUE self, VALUE buffer_value, VALUE size_value);
 
+static void compact_buffer(VALUE self, VALUE buf, long index) {
+  long remaining = RSTRING_LEN(buf) - index;
+  if (remaining == 0) {
+    /* Drop exhausted storage, including shared suffixes of larger strings. */
+    rb_ivar_set(self, buf_ivar_id, rb_str_new(NULL, 0));
+    index = 0;
+  } else if (index >= GARBAGE_BUFFER_SIZE && index >= remaining) {
+    /* Bound copying by bytes consumed and detach the unread suffix now. */
+    rb_ivar_set(self, buf_ivar_id, rb_str_new(RSTRING_PTR(buf) + index, remaining));
+    index = 0;
+  }
+  rb_ivar_set(self, index_ivar_id, LONG2NUM(index));
+  RB_GC_GUARD(buf);
+}
+
 VALUE rb_thrift_memory_buffer_write(VALUE self, VALUE str) {
   VALUE buf = GET_BUF(self);
   str = force_binary_encoding(str);
@@ -65,11 +80,7 @@ VALUE rb_thrift_memory_buffer_read(VALUE self, VALUE length_value) {
   } else {
     index += (long)length;
   }
-  if (index >= GARBAGE_BUFFER_SIZE) {
-    rb_ivar_set(self, buf_ivar_id, rb_funcall(buf, slice_method_id, 2, LONG2NUM(index), LONG2NUM(RSTRING_LEN(buf) - 1)));
-    index = 0;
-  }
-  rb_ivar_set(self, index_ivar_id, LONG2NUM(index));
+  compact_buffer(self, buf, index);
 
   if (RSTRING_LEN(data) < length) {
     rb_raise(rb_eEOFError, "Not enough bytes remain in memory buffer");
@@ -96,73 +107,61 @@ VALUE rb_thrift_memory_buffer_read_all(VALUE self, VALUE length_value) {
   VALUE data = rb_str_subseq(buf, index, (long)length);
 
   index += (long)length;
-  if (index >= GARBAGE_BUFFER_SIZE) {
-    rb_ivar_set(self, buf_ivar_id, rb_str_subseq(buf, index, RSTRING_LEN(buf) - index));
-    index = 0;
-  }
-  rb_ivar_set(self, index_ivar_id, LONG2NUM(index));
+  compact_buffer(self, buf, index);
 
   return data;
 }
 
 VALUE rb_thrift_memory_buffer_read_byte(VALUE self) {
   VALUE index_value = rb_ivar_get(self, index_ivar_id);
-  int index = FIX2INT(index_value);
+  long index = NUM2LONG(index_value);
 
   VALUE buf = GET_BUF(self);
   if (index >= RSTRING_LEN(buf)) {
     rb_raise(rb_eEOFError, "Not enough bytes remain in memory buffer");
   }
   unsigned char byte = (unsigned char)RSTRING_PTR(buf)[index++];
 
-  if (index >= GARBAGE_BUFFER_SIZE) {
-    rb_ivar_set(self, buf_ivar_id, rb_funcall(buf, slice_method_id, 2, INT2FIX(index), INT2FIX(RSTRING_LEN(buf) - 1)));
-    index = 0;
-  }
-  rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+  compact_buffer(self, buf, index);
 
   return INT2FIX(byte);
 }
 
 VALUE rb_thrift_memory_buffer_read_into_buffer(VALUE self, VALUE buffer_value, VALUE size_value) {
   int i = 0;
   int size = FIX2INT(size_value);
-  int index;
+  long index;
   VALUE buf = GET_BUF(self);
 
-  index = FIX2INT(rb_ivar_get(self, index_ivar_id));
+  index = NUM2LONG(rb_ivar_get(self, index_ivar_id));
   if (size > 0) {
     if (index >= RSTRING_LEN(buf)) {
-      rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+      compact_buffer(self, buf, index);
       rb_raise(rb_eEOFError, "Not enough bytes remain in memory buffer");
     }
     StringValue(buffer_value);
     if (RSTRING_LEN(buffer_value) == 0) {
-      rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+      rb_ivar_set(self, index_ivar_id, LONG2NUM(index));
       rb_raise(rb_eIndexError, "index 0 out of string");
     }
     rb_str_modify(buffer_value);
   }
 
   while (i < size) {
     if (index >= RSTRING_LEN(buf)) {
-      rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+      compact_buffer(self, buf, index);
       rb_raise(rb_eEOFError, "Not enough bytes remain in memory buffer");
     }
     if (i >= RSTRING_LEN(buffer_value)) {
-      rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+      compact_buffer(self, buf, index);
       rb_raise(rb_eIndexError, "index %d out of string", i);
     }
     char byte = RSTRING_PTR(buf)[index++];
     ((char*)RSTRING_PTR(buffer_value))[i] = byte;
     i++;
   }
 
-  if (index >= GARBAGE_BUFFER_SIZE) {
-    rb_ivar_set(self, buf_ivar_id, rb_funcall(buf, slice_method_id, 2, INT2FIX(index), INT2FIX(RSTRING_LEN(buf) - 1)));
-    index = 0;
-  }
-  rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+  compact_buffer(self, buf, index);
 
   return INT2FIX(i);
 }

lib/rb/lib/thrift/transport/memory_buffer_transport.rb

@@ -29,6 +29,7 @@ class MemoryBufferTransport < BaseTransport
     def initialize(buffer = nil)
       @buf = buffer ? Bytes.force_binary_encoding(buffer.dup) : Bytes.empty_byte_buffer
       @index = 0
+      @compact_at = 0
     end
 
     def open?
@@ -49,6 +50,7 @@ def peek
     def reset_buffer(new_buf = "")
       @buf.replace Bytes.force_binary_encoding(new_buf)
       @index = 0
+      @compact_at = 0
     end
 
     def available
@@ -61,10 +63,7 @@ def read(len)
       data = @buf.slice(@index, len)
       @index += len
       @index = @buf.size if @index > @buf.size
-      if @index >= GARBAGE_BUFFER_SIZE
-        @buf = @buf.slice(@index..-1)
-        @index = 0
-      end
+      compact_buffer if @index >= @compact_at
       if data.size < len
         raise EOFError, "Not enough bytes remain in buffer"
       end
@@ -81,10 +80,7 @@ def read_byte
       raise EOFError.new("Not enough bytes remain in buffer") if @index >= @buf.size
       val = Bytes.get_string_byte(@buf, @index)
       @index += 1
-      if @index >= GARBAGE_BUFFER_SIZE
-        @buf = @buf.slice(@index..-1)
-        @index = 0
-      end
+      compact_buffer if @index >= @compact_at
       val
     end
 
@@ -99,11 +95,9 @@ def read_into_buffer(buffer, size)
         @index += 1
         i += 1
       end
-      if @index >= GARBAGE_BUFFER_SIZE
-        @buf = @buf.slice(@index..-1)
-        @index = 0
-      end
       i
+    ensure
+      compact_buffer if @index >= @compact_at
     end
 
     def write(wbuf)
@@ -131,5 +125,26 @@ def inspect_buffer
     def to_s
       "memory"
     end
+
+    private
+
+    def compact_buffer
+      remaining = @buf.bytesize - @index
+      if remaining == 0
+        # Replace rather than clear: even a small suffix can share a large input.
+        @buf = Bytes.empty_byte_buffer
+        @index = 0
+      elsif @index >= GARBAGE_BUFFER_SIZE && @index >= remaining
+        # Copy at most as many bytes as have been consumed since the last copy.
+        # Appending to fresh storage also detaches a shared suffix immediately.
+        @buf = Bytes.empty_byte_buffer << @buf.slice(@index, remaining)
+        @index = 0
+      end
+      # Cache the next boundary for pure Ruby reads. Appends can only move it
+      # forward, so an earlier cached boundary stays safe until rechecked.
+      size = @buf.bytesize
+      half = (size + 1) / 2
+      @compact_at = size < GARBAGE_BUFFER_SIZE ? size : (half < GARBAGE_BUFFER_SIZE ? GARBAGE_BUFFER_SIZE : half)
+    end
   end
 end

lib/rb/spec/base_transport_spec.rb

@@ -445,7 +445,7 @@
         (0..bytes.length).each do |position|
           @buffer.reset_buffer(input)
           @buffer.read(position)
-          expected = spellings.dup
+          expected = position == bytes.length ? [] : spellings.dup
           expected.insert(position, ">") if position < bytes.length
 
           output = @buffer.inspect_buffer
@@ -474,9 +474,9 @@
         @buffer.read_byte
         expect(@buffer.inspect_buffer).to eq("0 > f 10 ff")
         @buffer.read(3)
-        expect(@buffer.inspect_buffer).to eq("0 f 10 ff")
+        expect(@buffer.inspect_buffer).to eq("")
         @buffer.write([128].pack("C"))
-        expect(@buffer.inspect_buffer).to eq("0 f 10 ff > 80")
+        expect(@buffer.inspect_buffer).to eq("> 80")
         @buffer.reset_buffer([1].pack("C"))
         expect(@buffer.inspect_buffer).to eq("> 1")
       end
@@ -487,7 +487,7 @@
 
         expect(@buffer.inspect_buffer).to eq("> 0 ff")
         expect(@buffer.read(2)).to eq([0, 255].pack("C*"))
-        expect(@buffer.inspect_buffer).to eq("0 ff")
+        expect(@buffer.inspect_buffer).to eq("")
       end
 
       it "keeps dumps independent across calls and subclasses without consuming input" do
@@ -680,6 +680,106 @@
 
       expect { @buffer.read_all(3_397_380_576) }.to raise_error(EOFError)
     end
+
+    {
+      read: ->(transport, size) { transport.read(size) },
+      read_all: ->(transport, size) { transport.read_all(size) },
+      read_byte: ->(transport, size) { Array.new(size) { transport.read_byte }.pack("C*") },
+      read_into_buffer: ->(transport, size) do
+        destination = "\0".b * size
+        transport.read_into_buffer(destination, size)
+        destination
+      end,
+    }.each do |operation, consume|
+      it "releases exhausted storage through #{operation}, including a shared suffix" do
+        @buffer.reset_buffer(("a" * 65_536) + ("z" * 1024))
+        @buffer.read(65_536)
+        consume.call(@buffer, 1024)
+
+        expect(@buffer.available).to eq(0)
+        expect(@buffer.instance_variable_get(:@buf)).to be_empty
+        expect { @buffer.read_byte }.to raise_error(EOFError)
+        expect(@buffer.read(0)).to eq("")
+        @buffer.write("new")
+        expect(@buffer.read_all(3)).to eq("new")
+      end
+
+      it "preserves a rolling backlog through #{operation}" do
+        backlog = 65_536
+        chunk = "b" * 4096
+        @buffer.reset_buffer("a" * backlog)
+        expected = ("a" * backlog) + (chunk * 40)
+        40.times do |iteration|
+          @buffer.write(chunk)
+          expect(consume.call(@buffer, chunk.bytesize)).to eq(expected.slice!(0, chunk.bytesize))
+          expect(@buffer.available).to eq(backlog)
+          expected_index = ((iteration + 1) * chunk.bytesize) % backlog
+          expect(@buffer.instance_variable_get(:@index)).to eq(expected_index)
+        end
+
+        expect(@buffer.read_all(backlog)).to eq("b" * backlog)
+      end
+
+      it "releases exhausted storage after resetting a partially read buffer through #{operation}" do
+        @buffer.reset_buffer("x" * 65_536)
+        expect(consume.call(@buffer, 1)).to eq("x")
+        @buffer.reset_buffer("abc")
+
+        expect(consume.call(@buffer, 3)).to eq("abc")
+        expect(@buffer.available).to eq(0)
+        expect(@buffer.instance_variable_get(:@buf)).to be_empty
+        @buffer.write("new")
+        expect(@buffer.read_all(3)).to eq("new")
+      end
+
+      it "preserves a caller-owned shared suffix through #{operation}" do
+        source = ("x" * 65_536).byteslice(-1024, 1024).freeze
+        @buffer = Thrift::MemoryBufferTransport.new(source)
+        consume.call(@buffer, 1024)
+
+        expect(source).to eq("x" * 1024)
+        expect(@buffer.available).to eq(0)
+        @buffer.write("new")
+        expect(@buffer.read_all(3)).to eq("new")
+      end
+    end
+
+    it "preserves a small unread suffix across compaction and subsequent writes" do
+      @buffer.reset_buffer(("a" * 65_536) + ("z" * 1024))
+      expect(@buffer.read(65_536)).to eq("a" * 65_536)
+      @buffer.write("tail")
+
+      expect(@buffer.available).to eq(1028)
+      expect(@buffer.read_all(1028)).to eq(("z" * 1024) + "tail")
+    end
+
+    it "preserves returned strings and caller input across compaction, exhaustion and reset" do
+      source = (("\xff" * 32_768) + ("z" * 8192)).b.freeze
+      @buffer.reset_buffer(source)
+      first = @buffer.read(32_768)
+      @buffer.write("tail")
+      last = @buffer.read_all(@buffer.available)
+      @buffer.reset_buffer("replacement")
+      @buffer.write("!")
+
+      expect(first).to eq("\xff".b * 32_768)
+      expect(last).to eq(("z" * 8192) + "tail")
+      first.setbyte(0, 0)
+      expect(source.getbyte(0)).to eq(255)
+      expect(@buffer.read_all(@buffer.available)).to eq("replacement!")
+    end
+
+    it "releases storage after a partial read exhausts input" do
+      @buffer.reset_buffer("abc")
+      expect { @buffer.read(4) }.to raise_error(EOFError)
+      expect(@buffer.instance_variable_get(:@buf)).to be_empty
+
+      @buffer.reset_buffer("abc")
+      destination = +"xxxx"
+      expect { @buffer.read_into_buffer(destination, 4) }.to raise_error(EOFError)
+      expect(destination).to eq("abcx")
+      expect(@buffer.instance_variable_get(:@buf)).to be_empty
+    end
   end
 
   describe Thrift::IOStreamTransport do