THRIFT-6431: Amortize Ruby MemoryBuffer compaction and release exhausted storage

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