From b2e73a5e322ceeec825b6ffe3950177c81aaa4aa Mon Sep 17 00:00:00 2001 From: Dmytro Shteflyuk Date: Fri, 23 Oct 2026 16:31:20 -0400 Subject: [PATCH] THRIFT-6451: Reduce temporary allocations when reading Ruby JSON strings Client: rb Use transport byte reads for string bodies and share frozen escape tables. Document read_byte requirements for custom read hooks in 0.26.0. Co-Authored-By: OpenAI Codex (gpt-6.1-sol) --- lib/rb/.rubocop.yml | 3 + lib/rb/README.md | 9 + .../lib/thrift/protocol/compact_protocol.rb | 2 +- lib/rb/lib/thrift/protocol/json_protocol.rb | 88 +++++---- .../thrift/transport/io_stream_transport.rb | 7 +- lib/rb/spec/base_protocol_spec.rb | 2 +- lib/rb/spec/binary_protocol_spec_shared.rb | 10 +- lib/rb/spec/compact_protocol_spec.rb | 15 +- lib/rb/spec/json_protocol_spec.rb | 170 ++++++++++++++++++ lib/rb/spec/socket_spec_shared.rb | 10 +- 10 files changed, 271 insertions(+), 45 deletions(-) diff --git a/lib/rb/.rubocop.yml b/lib/rb/.rubocop.yml index 4c56aeaece8..4d6127ac817 100644 --- a/lib/rb/.rubocop.yml +++ b/lib/rb/.rubocop.yml @@ -351,6 +351,9 @@ Style/RedundantSelf: Style/RedundantStringEscape: Enabled: true +Style/Semicolon: + Enabled: true + Style/StringLiterals: Enabled: true EnforcedStyle: double_quotes diff --git a/lib/rb/README.md b/lib/rb/README.md index fc8417b97d1..d82c8cd1185 100644 --- a/lib/rb/README.md +++ b/lib/rb/README.md @@ -142,6 +142,15 @@ mount Thrift::RackApplication.new(processor) => "/thrift" ## Breaking Changes +### 0.26.0 + +`Thrift::JsonProtocol` can now use any transport's `read_byte` method to read JSON +strings. This cuts the number of temporary Strings. If you override `read`, keep +`read_byte` consistent with it: `read` alone no longer sees every byte in JSON +strings. `read_byte` must return the next byte as an Integer from 0 to 255, or raise +an error at end of input. The same rule applies to custom `LookaheadReader` hooks. +Transports without `read_byte` still use `read(1)`. + ### 0.25.0 Ruby 2.7 support has been removed because the runtime reached end of life on diff --git a/lib/rb/lib/thrift/protocol/compact_protocol.rb b/lib/rb/lib/thrift/protocol/compact_protocol.rb index 5de798469bb..e4005394e19 100644 --- a/lib/rb/lib/thrift/protocol/compact_protocol.rb +++ b/lib/rb/lib/thrift/protocol/compact_protocol.rb @@ -210,7 +210,7 @@ def write_list_begin(etype, size) end def write_set_begin(etype, size) - write_collection_begin(etype, size); + write_collection_begin(etype, size) end def write_bool(bool) diff --git a/lib/rb/lib/thrift/protocol/json_protocol.rb b/lib/rb/lib/thrift/protocol/json_protocol.rb index e5894162828..7f1e1a2c01c 100644 --- a/lib/rb/lib/thrift/protocol/json_protocol.rb +++ b/lib/rb/lib/thrift/protocol/json_protocol.rb @@ -27,6 +27,18 @@ def initialize(trans) @data = nil end + def byte_reads? + @trans.respond_to?(:read_byte) + end + + def read_byte + if @hasData + read.getbyte(0) + else + @trans.read_byte + end + end + def read if @hasData @hasData = false @@ -143,6 +155,12 @@ class JsonProtocol < BaseProtocol ].freeze private_constant :JSON_CHAR_TABLE + # Characters and decoded values share the same order. + # JSON permits the optional solidus escape "\/"; see RFC 8259, section 7. + JSON_ESCAPE_CHARS = "\"\\/bfnrt" + JSON_ESCAPE_VALUES = ["\"", "\\", "/", "\b", "\f", "\n", "\r", "\t"].map { |ch| ch.b.freeze }.freeze + private_constant :JSON_ESCAPE_CHARS, :JSON_ESCAPE_VALUES + attr_reader :max_string_size @@kJSONObjectStart = "{" @@ -266,7 +284,7 @@ def write_json_char(ch) trans.write(ch) end else - outCh = JSON_CHAR_TABLE[ch_value]; + outCh = JSON_CHAR_TABLE[ch_value] # Check if regular character, backslash escaped, or JSON escaped if outCh.kind_of? String trans.write(@@kJSONBackslash) @@ -307,7 +325,7 @@ def write_json_integer(num) if escapeNum trans.write(@@kJSONStringDelimiter) end - trans.write(num.to_s); + trans.write(num.to_s) if escapeNum trans.write(@@kJSONStringDelimiter) end @@ -318,15 +336,15 @@ def write_json_integer(num) def write_json_double(num) @context.write(trans) # Normalize output of thrift::to_string for NaNs and Infinities - special = false; + special = false if num.nan? - special = true; - val = @@kThriftNan; + special = true + val = @@kThriftNan elsif num.infinite? - special = true; - val = @@kThriftInfinity; + special = true + val = @@kThriftInfinity if (num < 0.0) - val = @@kThriftNegativeInfinity; + val = @@kThriftNegativeInfinity end else val = num.to_s @@ -345,7 +363,7 @@ def write_json_double(num) def write_json_object_start @context.write(trans) trans.write(@@kJSONObjectStart) - push_context(JSONPairContext.new); + push_context(JSONPairContext.new) end def write_json_object_end @@ -356,7 +374,7 @@ def write_json_object_end def write_json_array_start @context.write(trans) trans.write(@@kJSONArrayStart) - push_context(JSONListContext.new); + push_context(JSONListContext.new) end def write_json_array_end @@ -498,40 +516,46 @@ def read_json_escape_char # Decodes a JSON string, including unescaping, and returns the string via str def read_json_string(skipContext = false) - # This string's characters must match up with the elements in escape_char_vals. - # JSON permits the optional solidus escape "\/"; see RFC 8259, section 7. - escape_chars = "\"\\/bfnrt" - - # The elements of this array must match up with the sequence of characters in - # escape_chars - escape_char_vals = [ - "\"", "\\", "/", "\b", "\f", "\n", "\r", "\t", - ] - if !skipContext @context.read(@reader) end read_json_syntax_char(@@kJSONStringDelimiter) str = Bytes.empty_byte_buffer + ch = @reader.read + return decode_json_string(str) if ch == @@kJSONStringDelimiter + + # nil: unchecked, false: unavailable, true: enabled. + byte_reads = nil + delimiter = @@kJSONStringDelimiter + backslash = @@kJSONBackslash while true - ch = @reader.read - if (ch == @@kJSONStringDelimiter) - break - end - if (ch == @@kJSONBackslash) + if (ch == backslash) ch = @reader.read if (ch == "u") - ch = read_json_escape_char + ch = Bytes.force_binary_encoding(read_json_escape_char) else - pos = escape_chars.index(ch); + pos = JSON_ESCAPE_CHARS.index(ch) if pos.nil? # not found raise ProtocolException.new(ProtocolException::INVALID_DATA, "Expected control char, got '#{ch}'.") end - ch = escape_char_vals[pos] + ch = JSON_ESCAPE_VALUES[pos] end + elsif !byte_reads + ch = Bytes.force_binary_encoding(ch) end - str << Bytes.force_binary_encoding(ch) - check_string_size(str.bytesize) + str << ch + size = str.bytesize + check_string_size(size) + # Avoid byte-read setup for tiny strings; select once per value. + if byte_reads.nil? && size >= 16 + byte_reads = @reader.is_a?(LookaheadReader) && @reader.byte_reads? + if byte_reads + delimiter = @@kJSONStringDelimiter.getbyte(0) + backslash = @@kJSONBackslash.getbyte(0) + end + end + ch = byte_reads ? @reader.read_byte : @reader.read + break if ch == delimiter end decode_json_string(str) end @@ -560,7 +584,7 @@ def read_json_numeric_chars while true ch = @reader.peek if !is_json_numeric(ch) - break; + break end ch = @reader.read str << ch @@ -579,7 +603,7 @@ def read_json_integer str = read_json_numeric_chars begin - num = Integer(str); + num = Integer(str) rescue raise ProtocolException.new(ProtocolException::INVALID_DATA, "Expected numeric value; got \"#{str}\"") end diff --git a/lib/rb/lib/thrift/transport/io_stream_transport.rb b/lib/rb/lib/thrift/transport/io_stream_transport.rb index d305014ef03..cb5cf3ffec2 100644 --- a/lib/rb/lib/thrift/transport/io_stream_transport.rb +++ b/lib/rb/lib/thrift/transport/io_stream_transport.rb @@ -34,7 +34,12 @@ def initialize(input, output) def open?; not @input.closed? or not @output.closed? end def read(sz); @input.read(sz) end def write(buf); @output.write(Bytes.force_binary_encoding(buf)) end - def close; @input.close; @output.close end + + def close + @input.close + @output.close + end + def to_io; @input end # we're assuming this is used in a IO.select for reading def to_s diff --git a/lib/rb/spec/base_protocol_spec.rb b/lib/rb/spec/base_protocol_spec.rb index 35ff2f03f7a..7a88297d06a 100644 --- a/lib/rb/spec/base_protocol_spec.rb +++ b/lib/rb/spec/base_protocol_spec.rb @@ -220,7 +220,7 @@ expect(@prot).to receive(:read_map_begin).ordered.and_return([Thrift::Types::STRING, Thrift::Types::STRUCT, 1]) expect(@prot).to receive(:read_string).ordered expect(@prot).to receive(:read_struct_begin).ordered.and_return(["some_struct"]) - expect(@prot).to receive(:read_field_begin).ordered.and_return([nil, Thrift::Types::STOP, nil]); + expect(@prot).to receive(:read_field_begin).ordered.and_return([nil, Thrift::Types::STOP, nil]) expect(@prot).to receive(:read_struct_end).ordered expect(@prot).to receive(:read_map_end).ordered real_skip.call(Thrift::Types::MAP) diff --git a/lib/rb/spec/binary_protocol_spec_shared.rb b/lib/rb/spec/binary_protocol_spec_shared.rb index 1d6ebfb63ed..415ab2c8ec1 100644 --- a/lib/rb/spec/binary_protocol_spec_shared.rb +++ b/lib/rb/spec/binary_protocol_spec_shared.rb @@ -74,7 +74,7 @@ it "should write the map header" do @prot.write_map_begin(Thrift::Types::STRING, Thrift::Types::LIST, 17) - expect(@trans.read(@trans.available)).to eq([Thrift::Types::STRING, Thrift::Types::LIST, 17].pack("ccN")); + expect(@trans.read(@trans.available)).to eq([Thrift::Types::STRING, Thrift::Types::LIST, 17].pack("ccN")) end # map footer is a noop @@ -380,7 +380,7 @@ # field footer is a noop it "should read a stop field" do - @trans.write([Thrift::Types::STOP].pack("c")); + @trans.write([Thrift::Types::STOP].pack("c")) expect(@prot.read_field_begin).to eq([nil, Thrift::Types::STOP, 0]) end @@ -427,7 +427,7 @@ # set footer is a noop it "should read a bool" do - @trans.write("\001\000"); + @trans.write("\001\000") expect(@prot.read_bool).to eq(true) expect(@prot.read_bool).to eq(false) end @@ -442,7 +442,7 @@ it "should read an i16" do # try a scattering of values, including min/max [-2**15, -5237, -353, 0, 1527, 2234, (2**15) - 1].each do |i| - @trans.write([i].pack("n")); + @trans.write([i].pack("n")) expect(@prot.read_i16).to eq(i) end end @@ -488,7 +488,7 @@ it "should read a double" do # try a random scattering of values, including min/max [Float::MIN, -231231.12351, -323.233513, 0, 123.2351235, 2351235.12351235, Float::MAX].each do |f| - @trans.write([f].pack("G")); + @trans.write([f].pack("G")) expect(@prot.read_double).to eq(f) end end diff --git a/lib/rb/spec/compact_protocol_spec.rb b/lib/rb/spec/compact_protocol_spec.rb index af705d6029f..3d20d2d0cb3 100644 --- a/lib/rb/spec/compact_protocol_spec.rb +++ b/lib/rb/spec/compact_protocol_spec.rb @@ -754,7 +754,10 @@ input.freeze if frozen writes = [] transport = Thrift::BaseTransport.new - transport.define_singleton_method(:write) { |bytes| writes << bytes; bytes.setbyte(0, 0) } + transport.define_singleton_method(:write) do |bytes| + writes << bytes + bytes.setbyte(0, 0) + end expect(Thrift::CompactProtocol.new(transport).write_binary(input)).to be_nil expect(input).to eq(original) @@ -800,7 +803,10 @@ def initialize_copy(other) it "honors MemoryBuffer write overrides" do transport = Thrift::MemoryBufferTransport.new writes = [] - transport.define_singleton_method(:write) { |bytes| writes << bytes; super(bytes) } + transport.define_singleton_method(:write) do |bytes| + writes << bytes + super(bytes) + end protocol = Thrift::CompactProtocol.new(transport) protocol.write_binary("payload".b.freeze) @@ -838,7 +844,10 @@ def write(bytes) input = ("\u20ac" * 1024).freeze writes = [] transport = Thrift::BaseTransport.new - transport.define_singleton_method(:write) { |bytes| writes << bytes; bytes.clear } + transport.define_singleton_method(:write) do |bytes| + writes << bytes + bytes.clear + end Thrift::CompactProtocol.new(transport).write_string(input) expect(input).to eq("\u20ac" * 1024) expect(input.encoding).to eq(Encoding::UTF_8) diff --git a/lib/rb/spec/json_protocol_spec.rb b/lib/rb/spec/json_protocol_spec.rb index daf77a5c2f5..4ae58cea270 100644 --- a/lib/rb/spec/json_protocol_spec.rb +++ b/lib/rb/spec/json_protocol_spec.rb @@ -21,6 +21,176 @@ require "spec_helper" +describe "JSON input byte ownership" do + class FragmentedJsonInput < Thrift::BaseTransport + attr_reader :requested + + def initialize(parts) + @parts = parts.map(&:b) + @requested = [] + end + + def read(size) + @requested << size + @parts.shift while @parts.first == "" + raise EOFError, "fragmented input exhausted" if @parts.empty? + + @parts.first.slice!(0, size) + end + end + + it "reads Unicode, escapes, base64 and numbers across every fragment boundary" do + wire = <<~'JSON'.chomp.b + ["str",3,"raw € 😀\u20ac\uD83D\uDE00\\\"\/\b\f\n\r\t","QUJDRA",-12.5e+2] + JSON + (1...wire.bytesize).each do |split| + input = FragmentedJsonInput.new([wire.byteslice(0, split), wire.byteslice(split..-1)]) + protocol = Thrift::JsonProtocol.new(Thrift::BufferedTransport.new(input)) + expect(protocol.read_list_begin).to eq([Thrift::Types::STRING, 3]) + expect(protocol.read_string).to eq("raw € 😀€😀\\\"/\b\f\n\r\t") + expect(protocol.read_binary).to eq("ABCD") + expect(protocol.read_double).to eq(-1250.0) + protocol.read_list_end + end + end + + it "leaves the following value in the transport and owns at most one peeked byte" do + trans = Thrift::MemoryBufferTransport.new('"first"123,"second"') + protocol = Thrift::JsonProtocol.new(trans) + expect(protocol.read_string).to eq("first") + expect(trans.available).to eq('123,"second"'.bytesize) + expect(protocol.read_i32).to eq(123) + expect(trans.read(1)).to eq('"') # Only the comma belongs to the reader. + reader = protocol.instance_variable_get(:@reader) + expect(reader.peek).to eq(",") + expect(reader.read).to eq(",") + end + + it "falls back to read when the transport has no byte API" do + klass = Class.new(Thrift::MemoryBufferTransport) { undef_method :read_byte } + value = "v" * 64 + expect(Thrift::JsonProtocol.new(klass.new('"' + value + '"')).read_string).to eq(value) + end + + it "uses the byte API on a duck-typed transport" do + value = "v" * 64 + input = Thrift::MemoryBufferTransport.new('"' + value + '"') + transport = double("Transport") + allow(transport).to receive(:read).with(1) { input.read(1) } + expect(transport).to receive(:read_byte).at_least(:once) { input.read_byte } + + expect(Thrift::JsonProtocol.new(transport).read_string).to eq(value) + expect(input.available).to eq(0) + end + + it "preserves Unicode decoding hooks" do + protocol_class = Class.new(Thrift::JsonProtocol) do + def read_json_escape_char + super.upcase + end + end + trans = Thrift::MemoryBufferTransport.new('"' + ("v" * 64) + 'a\\u0062"') + expect(protocol_class.new(trans).read_string).to eq(("v" * 64) + "aB") + end + + it "keeps tiny strings on the read path without inspecting byte-read eligibility" do + protocol = Thrift::JsonProtocol.new(Thrift::MemoryBufferTransport.new('"""a"')) + reader = protocol.instance_variable_get(:@reader) + expect(reader).not_to receive(:byte_reads?) + expect(protocol.read_string).to eq("") + expect(protocol.read_string).to eq("a") + end + + it "selects byte reads once when Unicode escapes cross the small-string threshold" do + wire = '"' + ('\\u20ac' * 6) + ("a" * 1024) + '"' + protocol = Thrift::JsonProtocol.new(Thrift::MemoryBufferTransport.new(wire)) + reader = protocol.instance_variable_get(:@reader) + expect(reader).to receive(:byte_reads?).once.and_call_original + expect(reader).to receive(:read_byte).at_least(:once).and_call_original + expect(protocol.read_string).to eq(("€" * 6) + ("a" * 1024)) + end + + [Thrift::MemoryBufferTransport, Thrift::BufferedTransport].each do |klass| + it "uses the byte API when #{klass} customizes it" do + reads = 0 + transport_class = Class.new(klass) do + define_method(:read_byte) do + reads += 1 + super() + end + end + value = "a" * 64 + wire = '"' + value + '"' + trans = if klass == Thrift::BufferedTransport + transport_class.new(FragmentedJsonInput.new([wire])) + else + transport_class.new(wire) + end + expect(Thrift::JsonProtocol.new(trans).read_string).to eq(value) + expect(reads).to be > 0 + end + end + + it "supports duck-typed readers" do + value = "v" * 64 + trans = Thrift::MemoryBufferTransport.new('"' + value + '"') + reader = Object.new + reader.define_singleton_method(:read) { trans.read(1) } + protocol = Thrift::JsonProtocol.new(trans) + protocol.instance_variable_set(:@reader, reader) + expect(protocol.read_string).to eq(value) + end + + it "preserves invalid UTF-8 and Unicode errors across fragments" do + ["\"\xC3\x28\"".b, '"\\uD83D\\u0041"', '"\\uDE00"'].each do |wire| + (1...wire.bytesize).each do |split| + input = FragmentedJsonInput.new([wire.byteslice(0, split), wire.byteslice(split..-1)]) + protocol = Thrift::JsonProtocol.new(Thrift::BufferedTransport.new(input)) + expect { protocol.read_string }.to raise_error(Thrift::ProtocolException) do |error| + expect(error.type).to eq(Thrift::ProtocolException::INVALID_DATA) + end + end + end + end + + it "retains mutable independent read and peek results" do + reader = Thrift::LookaheadReader.new(Thrift::MemoryBufferTransport.new("ab")) + first = reader.peek + expect(reader.read).to equal(first) + second = reader.read + second.replace("z") + expect(first).to eq("a") + end + + it "does not consume the next frame when a string ends at a frame boundary" do + value = "a" * 64 + first = '"' + value + '"' + wire = [first.bytesize].pack("N") + first + [3].pack("N") + '"b"' + input = Thrift::MemoryBufferTransport.new(wire) + protocol = Thrift::JsonProtocol.new(Thrift::FramedTransport.new(input)) + expect(protocol.read_string).to eq(value) + expect(input.available).to eq(7) + expect(protocol.read_string).to eq("b") + expect(input.available).to eq(0) + + value = ('"' + ("a" * 64) + '€\\uD83D\\uDE00"').b + wire = value.bytes.map { |byte| [1, byte].pack("NC") }.join + protocol = Thrift::JsonProtocol.new(Thrift::FramedTransport.new(Thrift::MemoryBufferTransport.new(wire))) + expect(protocol.read_string).to eq(("a" * 64) + "€😀") + end + + it "preserves truncated input errors in plain and Unicode strings" do + ['"abc', '"abc\\'].each do |wire| + protocol = Thrift::JsonProtocol.new(Thrift::MemoryBufferTransport.new(wire)) + expect { protocol.read_string }.to raise_error(EOFError, /Not enough bytes remain/) + end + ['"\\u123', '"\\uD83D\\uDE0'].each do |wire| + protocol = Thrift::JsonProtocol.new(Thrift::MemoryBufferTransport.new(wire)) + expect { protocol.read_string }.to raise_error(Thrift::ProtocolException, "Incomplete Unicode escape") + end + end +end + describe "JsonProtocol" do describe Thrift::JsonProtocol do before(:each) do diff --git a/lib/rb/spec/socket_spec_shared.rb b/lib/rb/spec/socket_spec_shared.rb index 13a524bcd8f..c28a863155e 100644 --- a/lib/rb/spec/socket_spec_shared.rb +++ b/lib/rb/spec/socket_spec_shared.rb @@ -121,7 +121,10 @@ @socket.timeout = 0.5 @socket.open expect(@handle).to receive(:read_nonblock).with(17).and_raise(IO::EAGAINWaitReadable) - expect(IO).to(receive(:select).once { sleep(0.6); nil }) + expect(IO).to(receive(:select).once do + sleep(0.6) + nil + end) expect { @socket.read(17) }.to raise_error(Thrift::TransportException) { |e| expect(e.type).to eq(Thrift::TransportException::TIMED_OUT) } end @@ -129,7 +132,10 @@ @socket.timeout = 0.5 @socket.open expect(@handle).to receive(:write_nonblock).with("test data").and_raise(IO::EAGAINWaitWritable) - expect(IO).to(receive(:select).once { sleep(0.6); nil }) + expect(IO).to(receive(:select).once do + sleep(0.6) + nil + end) expect { @socket.write("test data") }.to raise_error(Thrift::TransportException) { |e| expect(e.type).to eq(Thrift::TransportException::TIMED_OUT) } end