diff --git a/lib/rb/ext/memory_buffer.c b/lib/rb/ext/memory_buffer.c index b2d36c57a0..e8986d5386 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 014dbe1f38..7654a9c322 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 0c58b8136f..b763fc43d0 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