From fb50fbfe0a27489ba8d78095228f379e8daf3ffa Mon Sep 17 00:00:00 2001 From: Dmytro Shteflyuk Date: Fri, 16 Oct 2026 17:31:14 -0400 Subject: [PATCH] THRIFT-6431: Amortize Ruby MemoryBuffer compaction and release exhausted storage 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) --- lib/rb/ext/memory_buffer.c | 53 +++++---- .../transport/memory_buffer_transport.rb | 39 +++++-- lib/rb/spec/base_transport_spec.rb | 108 +++++++++++++++++- 3 files changed, 157 insertions(+), 43 deletions(-) diff --git a/lib/rb/ext/memory_buffer.c b/lib/rb/ext/memory_buffer.c index b2d36c57a00..e8986d5386a 100644 --- a/lib/rb/ext/memory_buffer.c +++ b/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,18 +107,14 @@ 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)) { @@ -115,11 +122,7 @@ VALUE rb_thrift_memory_buffer_read_byte(VALUE self) { } 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); } @@ -127,18 +130,18 @@ 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) { 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); @@ -146,11 +149,11 @@ VALUE rb_thrift_memory_buffer_read_into_buffer(VALUE self, VALUE buffer_value, V 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++]; @@ -158,11 +161,7 @@ VALUE rb_thrift_memory_buffer_read_into_buffer(VALUE self, VALUE buffer_value, V 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); } diff --git a/lib/rb/lib/thrift/transport/memory_buffer_transport.rb b/lib/rb/lib/thrift/transport/memory_buffer_transport.rb index 014dbe1f38b..7654a9c322d 100644 --- a/lib/rb/lib/thrift/transport/memory_buffer_transport.rb +++ b/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 diff --git a/lib/rb/spec/base_transport_spec.rb b/lib/rb/spec/base_transport_spec.rb index 0c58b8136f6..b763fc43d05 100644 --- a/lib/rb/spec/base_transport_spec.rb +++ b/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