From 94442bd8a589f507716737e1da8b1984959cc71f Mon Sep 17 00:00:00 2001 From: Dmytro Shteflyuk Date: Thu, 15 Oct 2026 20:40:24 -0400 Subject: [PATCH] THRIFT-6427: Avoid copying payload bytes in native Ruby Compact writes Client: rb Use a separate mutable String whose storage Ruby can share with the input until either value changes. Preserve transport write overrides, callback mutation isolation, frozen inputs, and retained payload snapshots. Keep the existing copying path when a header callback changes the source length. Normalize the callback payload after header writes so a callback that re-encodes the source cannot change the payload's binary encoding. Add behavioral coverage for callback ownership, encodings, and transport overrides. Allocation measurements remain outside the regression suite. Co-Authored-By: OpenAI Codex (gpt-6.1-sol) --- lib/rb/ext/compact_protocol.c | 10 +- lib/rb/spec/compact_protocol_spec.rb | 195 +++++++++++++++++++++++++++ 2 files changed, 204 insertions(+), 1 deletion(-) diff --git a/lib/rb/ext/compact_protocol.c b/lib/rb/ext/compact_protocol.c index 9cdf39d8327..ef78096f2f8 100644 --- a/lib/rb/ext/compact_protocol.c +++ b/lib/rb/ext/compact_protocol.c @@ -370,7 +370,15 @@ VALUE rb_thrift_compact_proto_write_binary(VALUE self, VALUE buf) { rb_raise(rb_eRangeError, "integer out of bounds"); } write_varint32(transport, (uint32_t)size); - WRITE(transport, StringValuePtr(buf), size); + if (RSTRING_LEN(buf) == size) { + // Keep a separate mutable String for transport callbacks while allowing + // Ruby to share the payload storage until either String is modified. + rb_funcall(transport, write_method_id, 1, force_binary_encoding(rb_str_subseq(buf, 0, size))); + } else { + // A header callback may have changed the source length. Preserve the + // original byte-range write in that case. + WRITE(transport, StringValuePtr(buf), size); + } return Qnil; } diff --git a/lib/rb/spec/compact_protocol_spec.rb b/lib/rb/spec/compact_protocol_spec.rb index bdba8aa5b1f..ba46bd7a03d 100644 --- a/lib/rb/spec/compact_protocol_spec.rb +++ b/lib/rb/spec/compact_protocol_spec.rb @@ -565,6 +565,201 @@ def write(data) end end + describe "binary payload storage" do + [false, true].each do |frozen| + [Encoding::BINARY, Encoding::UTF_8, Encoding::ISO_8859_1].each do |encoding| + [:memory, :framed].each do |kind| + it "round trips #{frozen ? "frozen" : "mutable"} #{encoding} input through #{kind}" do + input = ("a\x00\xff".b * 1024).force_encoding(encoding) + expected = input.b + input.freeze if frozen + memory = Thrift::MemoryBufferTransport.new + transport = kind == :memory ? memory : Thrift::FramedTransport.new(memory) + protocol = Thrift::CompactProtocol.new(transport) + + protocol.write_binary(input) + transport.flush + input.replace("changed") unless frozen + expect(protocol.read_binary).to eq(expected) + expect(input.encoding).to eq(encoding) if frozen + end + end + end + end + + context "with native callback ownership" do + before do + skip "native Compact ownership" unless Thrift::CompactProtocol.instance_method(:write_binary).source_location.nil? + end + + [false, true].each do |frozen| + [Encoding::BINARY, Encoding::UTF_8, Encoding::ISO_8859_1].each do |encoding| + it "normalizes a separate #{frozen ? "frozen" : "mutable"} #{encoding} callback argument" do + input = ("a\x00\xff".b * 1024).force_encoding(encoding) + expected = input.b + input.freeze if frozen + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) { |bytes| writes << bytes } + Thrift::CompactProtocol.new(transport).write_binary(input) + + expect(writes.last).to eq(expected) + expect(writes.last.encoding).to eq(Encoding::BINARY) + expect(writes.last).not_to equal(input) + expect(writes.last).not_to be_frozen + expect(input.encoding).to eq(frozen ? encoding : Encoding::BINARY) + end + end + end + + [false, true].each do |frozen| + it "isolates a #{frozen ? "frozen" : "mutable"} input from a mutating callback" do + input = "payload".b * 1024 + original = input.dup + input.freeze if frozen + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) { |bytes| writes << bytes; bytes.setbyte(0, 0) } + + expect(Thrift::CompactProtocol.new(transport).write_binary(input)).to be_nil + expect(input).to eq(original) + expect(writes.last).not_to equal(input) + expect(writes.last).not_to be_frozen + expect(writes.last.encoding).to eq(Encoding::BINARY) + end + end + + it "keeps retained payloads independent of later source and callback mutations" do + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) { |bytes| writes << bytes } + protocol = Thrift::CompactProtocol.new(transport) + input = "a".b * 4096 + 2.times { protocol.write_binary(input) } + input.replace("changed") + GC.start + GC.verify_compaction_references(double_heap: true, toward: :empty) if GC.respond_to?(:verify_compaction_references) + + expect(writes[1]).to eq("a".b * 4096) + writes[1].setbyte(0, 98) + expect(writes[3]).to eq("a".b * 4096) + expect(input).to eq("changed") + end + + it "passes a plain String without invoking input subclass copy hooks" do + klass = Class.new(String) do + def initialize_copy(other) + raise "unexpected copy hook" + end + end + input = klass.new("a".b * 4096).freeze + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) { |bytes| writes << bytes } + Thrift::CompactProtocol.new(transport).write_binary(input) + + expect(writes.last).to be_instance_of(String) + expect(writes.last).to eq(input) + end + + it "honors MemoryBuffer write overrides" do + transport = Thrift::MemoryBufferTransport.new + writes = [] + transport.define_singleton_method(:write) { |bytes| writes << bytes; super(bytes) } + protocol = Thrift::CompactProtocol.new(transport) + protocol.write_binary("payload".b.freeze) + + expect(writes.map(&:bytesize)).to eq([1, 7]) + expect(protocol.read_binary).to eq("payload") + end + + it "honors MemoryBuffer subclass write overrides" do + klass = Class.new(Thrift::MemoryBufferTransport) do + attr_reader :writes + + def write(bytes) + (@writes ||= []) << bytes + super + end + end + transport = klass.new + protocol = Thrift::CompactProtocol.new(transport) + protocol.write_binary("a".b * 4096) + + expect(transport.writes.map(&:bytesize)).to eq([2, 4096]) + expect(protocol.read_binary).to eq("a".b * 4096) + end + + it "passes a mutable empty payload after its length prefix" do + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) { |bytes| writes << bytes } + Thrift::CompactProtocol.new(transport).write_binary("".b.freeze) + expect(writes).to eq(["\x00".b, "".b]) + expect(writes.last).not_to be_frozen + end + + it "keeps write_string input separate from a mutating callback" do + input = ("\u20ac" * 1024).freeze + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) { |bytes| writes << bytes; bytes.clear } + Thrift::CompactProtocol.new(transport).write_string(input) + expect(input).to eq("\u20ac" * 1024) + expect(input.encoding).to eq(Encoding::UTF_8) + expect(writes.last.encoding).to eq(Encoding::BINARY) + end + + it "preserves the declared length when the header callback appends to the source" do + input = "abc".b + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) do |bytes| + input << "d" if writes.empty? + writes << bytes + end + Thrift::CompactProtocol.new(transport).write_binary(input) + expect(writes).to eq(["\x03".b, "abc".b]) + end + + [false, true].each do |frozen| + it "passes binary bytes when the header callback re-encodes #{frozen ? "and freezes" : "a mutable"} input" do + input = "a".b * 4096 + expected = input.dup + output = "".b + writes = [] + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) do |bytes| + if writes.empty? + input.force_encoding(Encoding::UTF_16LE) + input.freeze if frozen + end + writes << bytes + output << bytes + end + + Thrift::CompactProtocol.new(transport).write_binary(input) + + expect(writes.last.encoding).to eq(Encoding::BINARY) + expect(writes.last).not_to equal(input) + expect(writes.last).not_to be_frozen + expect(input.encoding).to eq(Encoding::UTF_16LE) + expect(Thrift::CompactProtocol.new(Thrift::MemoryBufferTransport.new(output)).read_binary).to eq(expected) + end + end + + it "propagates payload callback exceptions without changing the input" do + input = "abc".b.freeze + transport = Thrift::BaseTransport.new + transport.define_singleton_method(:write) do |bytes| + raise IOError, "payload rejected" if bytes.bytesize == 3 + end + expect { Thrift::CompactProtocol.new(transport).write_binary(input) }.to raise_error(IOError, "payload rejected") + expect(input).to eq("abc") + end + end + end + it "should accept a valid 10-byte varint" do trans = Thrift::MemoryBufferTransport.new((([0x80] * 9) + [0x01]).pack("C*")) proto = Thrift::CompactProtocol.new(trans)