From d8da61a6930123f8f47c9002e9709d71b5d5b161 Mon Sep 17 00:00:00 2001 From: Dmytro Shteflyuk Date: Sun, 19 Jul 2026 21:28:28 -0400 Subject: [PATCH] THRIFT-6107: Preserve NonblockingServer reply order Client: rb Co-Authored-By: OpenAI Codex (GPT-5.4) --- .../lib/thrift/server/nonblocking_server.rb | 305 ++++++++++++++-- lib/rb/spec/nonblocking_server_spec.rb | 334 +++++++++++++++++- 2 files changed, 608 insertions(+), 31 deletions(-) diff --git a/lib/rb/lib/thrift/server/nonblocking_server.rb b/lib/rb/lib/thrift/server/nonblocking_server.rb index 3f248ea0bb8..249bc1d6b45 100644 --- a/lib/rb/lib/thrift/server/nonblocking_server.rb +++ b/lib/rb/lib/thrift/server/nonblocking_server.rb @@ -107,12 +107,17 @@ def initialize(processor, server_transport, transport_factory, protocol_factory, @num_threads = num @logger = logger @connections = [] + @readable_connections = [] @buffers = Hash.new { |h, k| h[k] = Bytes.empty_byte_buffer } + @response_queues = {} @signal_queue = Queue.new @signal_pipes = IO.pipe @signal_pipes[1].sync = true @worker_queue = Queue.new @shutdown_queue = Queue.new + @capacity_callback = method(:response_capacity_available) + @error_callback = method(:response_write_failed) + @full_callback = method(:pause_connection) end def add_connection(socket) @@ -153,7 +158,7 @@ def run spin_worker_threads loop do - rd, = select([@signal_pipes[0], *@connections]) + rd = select_readable if rd.delete @signal_pipes[0] break if read_signals == :shutdown end @@ -164,7 +169,8 @@ def run else read_connection fd end - rescue Errno::ECONNRESET + rescue IOError, SystemCallError, TransportException => e + @logger.debug "#{self} could not read connection: #{e.inspect}" remove_connection fd end end @@ -174,11 +180,41 @@ def run @shutdown_queue.push :shutdown end + def select_readable + select([@signal_pipes[0], *@readable_connections]).first + rescue IOError, SystemCallError, TransportException => e + @logger.debug "#{self} discarded closed connections after select failed: #{e.inspect}" + raise if remove_closed_connections == 0 + + retry + end + + def remove_closed_connections + closed = @connections.select do |fd| + @response_queues[fd]&.closed? || !fd.open? + rescue IOError, SystemCallError, TransportException + true + end + closed.each { |fd| remove_connection(fd) } + closed.length + end + def read_connection(fd) @buffers[fd] << fd.read(DEFAULT_BUFFER) - while(frame = slice_frame!(@buffers[fd])) + dispatch_frames(fd) + end + + def dispatch_frames(fd) + response_queue = @response_queues[fd] + return unless response_queue + + while (frame_size = complete_frame_size(@buffers[fd])) + sequence = response_queue.reserve + break if sequence.nil? + @logger.debug "#{self} is processing a frame" - @worker_queue.push [:frame, fd, frame] + frame = @buffers[fd].slice!(0, frame_size) + @worker_queue.push [:frame, response_queue, sequence, frame] end end @@ -211,6 +247,15 @@ def read_signals case signal when :connection @connections << obj + @readable_connections << obj + @response_queues[obj] = new_response_queue(obj) + when :capacity + if @response_queues.key?(obj) + @readable_connections << obj unless @readable_connections.include?(obj) + dispatch_frames(obj) + end + when :disconnect + remove_connection(obj) when :shutdown @shutdown_timeout = obj return :shutdown @@ -226,9 +271,37 @@ def read_signals end def remove_connection(fd) - # don't explicitly close it, a thread may still be writing to it @connections.delete fd + @readable_connections.delete fd @buffers.delete fd + @response_queues.delete(fd)&.close + fd.close + rescue IOError, SystemCallError, TransportException + end + + def new_response_queue(fd) + ResponseQueue.new( + fd, + @logger, + max_in_flight: [@num_threads, 1].max, + on_full: @full_callback, + on_capacity: @capacity_callback, + on_error: @error_callback + ) + end + + def response_capacity_available(fd) + signal [:capacity, fd] + rescue IOError, SystemCallError + end + + def response_write_failed(fd, _error) + signal [:disconnect, fd] + rescue IOError, SystemCallError + end + + def pause_connection(fd) + @readable_connections.delete(fd) end def join_worker_threads(shutdown_timeout) @@ -253,13 +326,11 @@ def kill_worker_threads end def close_connections - @connections.each do |fd| - begin - fd.close - rescue IOError, SystemCallError, TransportException - end - end + @connections.dup.each { |fd| remove_connection(fd) } + @response_queues.each_value(&:close) + @response_queues.clear @connections.clear + @readable_connections.clear @buffers.clear end @@ -272,16 +343,193 @@ def close_signal_pipes end end - def slice_frame!(buf) - if buf.length >= 4 - size = buf.unpack('N').first - if buf.length >= size + 4 - buf.slice!(0, size + 4) + def complete_frame_size(buf) + return if buf.length < 4 + + frame_size = buf.unpack1('N') + 4 + frame_size if buf.length >= frame_size + end + + class ResponseBufferTransport < BaseTransport # :nodoc: + def initialize + reset + end + + def reset + @data = nil + end + + def write(buf, size = nil) + chunk = if size && size < buf.bytesize + buf.byteslice(0, size) + else + buf + end + if @data.nil? + @data = chunk + elsif @data.is_a?(Array) + @data << chunk + else + @data = [@data, chunk] + end + end + + def flush; end + + def data + case @data + when nil + Bytes.empty_byte_buffer + when Array + @data.join else - nil + @data end - else - nil + end + + def take + response = data + reset + response + end + end + + class ResponseQueue # :nodoc: + # A connection owns one queue. It bounds outstanding responses and publishes + # only contiguous request sequences to the connection's transport. + def initialize(transport, logger, max_in_flight: nil, on_full: nil, on_capacity: nil, on_error: nil) + if max_in_flight && max_in_flight < 1 + raise ArgumentError, 'max_in_flight must be at least 1' + end + + @transport = transport + @logger = logger + @max_in_flight = max_in_flight + @on_full = on_full + @on_capacity = on_capacity + @on_error = on_error + @mutex = Mutex.new + @next_sequence = 0 + @next_to_publish = 0 + @in_flight = 0 + @completed = nil + @closed = false + end + + def reserve + became_full = false + sequence = @mutex.synchronize do + return unless accepting_without_lock? + + sequence = @next_sequence + @next_sequence += 1 + @in_flight += 1 + became_full = full? + sequence + end + @on_full&.call(@transport) if became_full + sequence + end + + def closed? + @mutex.synchronize { @closed } + end + + def complete(sequence, response) + error = nil + capacity_available = false + + @mutex.synchronize do + return if @closed || sequence < @next_to_publish + + was_full = full? + begin + publish_response(sequence, response) + capacity_available = was_full && accepting_without_lock? + rescue IOError, SystemCallError, TransportException => e + @logger.debug "#{self} could not write response: #{e.inspect}" + close_without_lock + error = e + end + end + + if error + @on_error&.call(@transport, error) + elsif capacity_available + @on_capacity&.call(@transport) + end + end + + def close + @mutex.synchronize { close_without_lock } + end + + private + + def accepting_without_lock? + !@closed && !full? + end + + def full? + @max_in_flight && @in_flight >= @max_in_flight + end + + def close_without_lock + @closed = true + @in_flight = 0 + @completed = nil + end + + def publish_response(sequence, response) + unless sequence == @next_to_publish + (@completed ||= {})[sequence] = response + return + end + + write_response(response) + while @completed&.key?(@next_to_publish) + response = @completed.delete(@next_to_publish) + write_response(response) + end + @completed = nil if @completed&.empty? + end + + def write_response(response) + unless response.nil? || response.empty? + @transport.write(response) + @transport.flush + end + @next_to_publish += 1 + @in_flight -= 1 + end + end + + class OnewayAwareProtocol < BaseProtocol # :nodoc: + include ProtocolDecorator + + def initialize + @protocol = nil + @oneway = false + end + + def reset(protocol, response_queue, sequence) + @protocol = protocol + @response_queue = response_queue + @sequence = sequence + @oneway = false + end + + def read_message_begin + message_begin = @protocol.read_message_begin + if !@oneway && message_begin[1] == MessageTypes::ONEWAY + @oneway = true + @response_queue.complete(@sequence, nil) + end + message_begin + end + + def oneway? + @oneway end end @@ -292,6 +540,8 @@ def initialize(processor, transport_factory, protocol_factory, logger, queue) @protocol_factory = protocol_factory @logger = logger @queue = queue + @response = ResponseBufferTransport.new + @request_protocol = OnewayAwareProtocol.new end def spawn @@ -311,16 +561,27 @@ def run @logger.debug "#{self} is shutting down, goodbye" break when :frame - fd, frame = args + response_queue, sequence, frame = args + next if response_queue.closed? + + @response.reset + @request_protocol.reset(nil, response_queue, sequence) begin - otrans = @transport_factory.get_transport(fd) + otrans = @transport_factory.get_transport(@response) oprot = @protocol_factory.get_protocol(otrans) membuf = MemoryBufferTransport.new(frame) itrans = @transport_factory.get_transport(membuf) iprot = @protocol_factory.get_protocol(itrans) - @processor.process(iprot, oprot) + @request_protocol.reset(iprot, response_queue, sequence) + @processor.process(@request_protocol, oprot) rescue => e @logger.error "#{Thread.current.inspect} raised error: #{e.inspect}\n#{e.backtrace.join("\n")}" + ensure + if @request_protocol.oneway? + @response.reset + else + response_queue.complete(sequence, @response.take) + end end end end diff --git a/lib/rb/spec/nonblocking_server_spec.rb b/lib/rb/spec/nonblocking_server_spec.rb index c8fbe8ef866..f092717674d 100644 --- a/lib/rb/spec/nonblocking_server_spec.rb +++ b/lib/rb/spec/nonblocking_server_spec.rb @@ -25,9 +25,11 @@ class Handler def initialize @queue = Queue.new + @block_started = Queue.new end attr_accessor :server + attr_reader :block_started def greeting(english) if english @@ -38,6 +40,7 @@ def greeting(english) end def block + @block_started << true @queue.pop end @@ -54,10 +57,36 @@ def shutdown end end + class SpecProcessor + def initialize(processor, finished) + @processor = processor + @finished = finished + end + + def process(iprot, oprot) + @processor.process(iprot, oprot) + ensure + @finished << true + end + end + + class BlockingProcessor + def initialize(started, release) + @started = started + @release = release + end + + def process(iprot, _oprot) + @started << iprot.read_message_begin + @release.pop + end + end + class SpecTransport < Thrift::BaseTransport - def initialize(transport, queue) + def initialize(transport, queue, write_queue = nil) @transport = transport @queue = queue + @write_queue = write_queue @flushed = false end @@ -78,7 +107,9 @@ def read(sz) end def write(buf, sz = nil) - @transport.write(buf, sz) + data = sz ? buf[0...sz] : buf + @write_queue << data.dup if @write_queue + @transport.write(data) end def flush @@ -86,32 +117,48 @@ def flush @flushed = true @transport.flush end + + def handle + @transport.handle + end + + def to_io + @transport.to_io + end end class SpecServerSocket < Thrift::ServerSocket - def initialize(host, port, queue) + def initialize(host, port, queue, write_queue = nil) super(host, port) @queue = queue + @write_queue = write_queue end def listen super @queue.push :listen end + + def accept + transport = super + SpecTransport.new(transport, nil, @write_queue) + end end describe Thrift::NonblockingServer do before(:each) do @port = available_port - handler = Handler.new - processor = SpecNamespace::NonblockingService::Processor.new(handler) + @handler = Handler.new + @processor_finished = Queue.new + processor = SpecProcessor.new(SpecNamespace::NonblockingService::Processor.new(@handler), @processor_finished) queue = Queue.new - @transport = SpecServerSocket.new('localhost', @port, queue) + @server_writes = Queue.new + @transport = SpecServerSocket.new('localhost', @port, queue, @server_writes) transport_factory = Thrift::FramedTransportFactory.new logger = Logger.new(STDERR) logger.level = Logger::WARN @server = Thrift::NonblockingServer.new(processor, @transport, transport_factory, nil, 5, logger) - handler.server = @server + @handler.server = @server @server_thread = Thread.new(Thread.current) do |master_thread| begin @server.serve @@ -178,6 +225,12 @@ def setup_client_thread(result) queue end + def reply_sequence_id(reply) + transport = Thrift::FramedTransport.new(Thrift::MemoryBufferTransport.new(reply)) + _name, _type, sequence_id = Thrift::BinaryProtocol.new(transport).read_message_begin + sequence_id + end + it "should handle basic message passing" do client = setup_client expect(client.greeting(true)).to eq(SpecNamespace::Hello.new) @@ -203,6 +256,27 @@ def setup_client_thread(result) @server.shutdown end + it "publishes replies for one connection in request order" do + client = setup_client + + client.send_block + expect(@handler.block_started.pop).to be_truthy + + client.send_greeting(true) + expect(@processor_finished.pop).to be_truthy + + expect do + @server_writes.pop(true) + end.to raise_error(ThreadError) + + @handler.unblock(1) + replies = 2.times.map { Timeout.timeout(1) { @server_writes.pop } } + + expect(replies.map { |reply| reply_sequence_id(reply) }).to eq([0, 1]) + expect(client.recv_block).to be_truthy + expect(client.recv_greeting).to eq(SpecNamespace::Hello.new) + end + it "should handle messages from more than 5 long-lived connections" do queues = [] result = Queue.new @@ -266,7 +340,7 @@ def setup_client_thread(result) end describe Thrift::NonblockingServer::IOManager do - def build_io_manager + def build_io_manager(num_threads: 1) logger = Logger.new(IO::NULL) logger.level = Logger::FATAL Thrift::NonblockingServer::IOManager.new( @@ -274,11 +348,69 @@ def build_io_manager double('server_transport'), Thrift::BaseTransportFactory.new, Thrift::BinaryProtocolFactory.new, - 1, + num_threads, logger ) end + def build_response_queue(max_in_flight: nil, on_capacity: nil, on_error: nil) + transport = double('transport', :write => nil, :flush => nil, :close => nil) + logger = Logger.new(IO::NULL) + logger.level = Logger::FATAL + response_queue = Thrift::NonblockingServer::IOManager::ResponseQueue.new( + transport, + logger, + max_in_flight: max_in_flight, + on_capacity: on_capacity, + on_error: on_error + ) + [response_queue, transport] + end + + def framed_message(type) + output = Thrift::MemoryBufferTransport.new + protocol = Thrift::BinaryProtocol.new(Thrift::FramedTransport.new(output)) + protocol.write_message_begin('request', type, 0) + protocol.write_struct_begin('args') + protocol.write_field_stop + protocol.write_struct_end + protocol.write_message_end + protocol.trans.flush + output.read(output.available) + end + + describe Thrift::NonblockingServer::IOManager::ResponseBufferTransport do + it "returns a single complete write without copying it" do + transport = described_class.new + response = +"response" + + transport.write(response) + + expect(transport.data).to be(response) + end + + it "combines partial and repeated writes" do + transport = described_class.new + + transport.write('first', 3) + transport.write('second') + + expect(transport.data).to eq('firsecond') + end + + it "can be reused without retaining its previous response" do + transport = described_class.new + transport.write('first response') + + expect(transport.take).to eq('first response') + expect(transport.data).to be_empty + + transport.write('second response') + + expect(transport.data).to eq('second response') + end + end + it "closes tracked connections and signal pipes during forced cleanup" do io_manager = build_io_manager connection = double('connection', :close => nil) @@ -318,14 +450,198 @@ def build_io_manager it "drops removed connections from bookkeeping" do io_manager = build_io_manager connection = double('connection', :close => nil) + response_queue, response_transport = build_response_queue + first_sequence = response_queue.reserve + second_sequence = response_queue.reserve + response_queue.complete(second_sequence, 'later reply') io_manager.instance_variable_set(:@connections, [connection]) io_manager.instance_variable_set(:@buffers, { connection => 'frame' }) + io_manager.instance_variable_set(:@response_queues, { connection => response_queue }) io_manager.send(:remove_connection, connection) + response_queue.complete(first_sequence, 'first reply') expect(io_manager.instance_variable_get(:@connections)).to be_empty expect(io_manager.instance_variable_get(:@buffers)).to be_empty + expect(io_manager.instance_variable_get(:@response_queues)).to be_empty + expect(response_transport).not_to have_received(:write) + expect(response_transport).not_to have_received(:flush) + end + + it "removes and closes a connection after its response write fails" do + io_manager = build_io_manager + connection = double('connection', :write => nil, :flush => nil, :close => nil) + allow(connection).to receive(:write).and_raise(IOError, 'closed stream') + response_queue = io_manager.send(:new_response_queue, connection) + + io_manager.instance_variable_set(:@connections, [connection]) + io_manager.instance_variable_set(:@buffers, { connection => '' }) + io_manager.instance_variable_set(:@response_queues, { connection => response_queue }) + + response_queue.complete(response_queue.reserve, 'reply') + io_manager.send(:read_signals) + + expect(connection).to have_received(:close).once + expect(io_manager.instance_variable_get(:@connections)).to be_empty + expect(io_manager.instance_variable_get(:@buffers)).to be_empty + expect(io_manager.instance_variable_get(:@response_queues)).to be_empty + end + + it "discards a closed connection when select observes its stale descriptor" do + io_manager = build_io_manager + connection = double('connection', :open? => false, :close => nil) + allow(connection).to receive(:to_io).and_raise(IOError, 'closed stream') + response_queue = io_manager.send(:new_response_queue, connection) + + io_manager.instance_variable_set(:@connections, [connection]) + io_manager.instance_variable_set(:@readable_connections, [connection]) + io_manager.instance_variable_set(:@buffers, { connection => '' }) + io_manager.instance_variable_set(:@response_queues, { connection => response_queue }) + io_manager.send(:signal, [:noop, nil]) + + begin + expect(io_manager.send(:select_readable)).to eq([io_manager.instance_variable_get(:@signal_pipes)[0]]) + io_manager.send(:read_signals) + + expect(connection).to have_received(:close).once + expect(io_manager.instance_variable_get(:@connections)).to be_empty + expect(io_manager.instance_variable_get(:@readable_connections)).to be_empty + ensure + io_manager.ensure_closed + end + end + + it "keeps excess frames buffered until the connection has response capacity" do + io_manager = build_io_manager(num_threads: 2) + connection = double('connection', :write => nil, :flush => nil, :close => nil) + response_queue = io_manager.send(:new_response_queue, connection) + frame = [1].pack('N') << 'x' + + io_manager.instance_variable_set(:@connections, [connection]) + io_manager.instance_variable_set(:@buffers, { connection => frame * 3 }) + io_manager.instance_variable_set(:@response_queues, { connection => response_queue }) + + io_manager.send(:dispatch_frames, connection) + + expect(io_manager.instance_variable_get(:@worker_queue).size).to eq(2) + expect(io_manager.instance_variable_get(:@buffers)[connection]).to eq(frame) + + response_queue.complete(0, '') + io_manager.send(:read_signals) + + expect(io_manager.instance_variable_get(:@worker_queue).size).to eq(3) + expect(io_manager.instance_variable_get(:@buffers)[connection]).to be_empty + end + + it "advances past oneway completions before publishing later replies" do + response_queue, transport = build_response_queue + oneway_sequence = response_queue.reserve + reply_sequence = response_queue.reserve + + response_queue.complete(reply_sequence, 'later reply') + expect(transport).not_to have_received(:write) + + response_queue.complete(oneway_sequence, '') + + expect(transport).to have_received(:write).with('later reply').once + expect(transport).to have_received(:flush).once + end + + it "does not let a running oneway handler block a later reply" do + response_queue, transport = build_response_queue + worker_queue = Queue.new + started = Queue.new + release = Queue.new + worker = Thrift::NonblockingServer::IOManager::Worker.new( + BlockingProcessor.new(started, release), + Thrift::FramedTransportFactory.new, + Thrift::BinaryProtocolFactory.new, + Logger.new(IO::NULL), + worker_queue + ) + worker_thread = worker.spawn + oneway_sequence = response_queue.reserve + reply_sequence = response_queue.reserve + + begin + worker_queue.push [:frame, response_queue, oneway_sequence, framed_message(Thrift::MessageTypes::ONEWAY)] + expect(started.pop[1]).to eq(Thrift::MessageTypes::ONEWAY) + + response_queue.complete(reply_sequence, 'later reply') + + expect(transport).to have_received(:write).with('later reply').once + expect(transport).to have_received(:flush).once + ensure + release << true + worker_queue.push [:shutdown] + worker_thread.join(1) + end + end + + it "notifies its owner and closes after a response write fails" do + failures = [] + response_queue, transport = build_response_queue( + on_error: ->(failed_transport, error) { failures << [failed_transport, error] } + ) + allow(transport).to receive(:write).and_raise(IOError, 'closed stream') + + response_queue.complete(response_queue.reserve, 'reply') + + expect(failures.length).to eq(1) + expect(failures.first[0]).to be(transport) + expect(failures.first[1]).to be_a(IOError) + expect(response_queue.reserve).to be_nil + end + + it "bounds completed responses until the missing reply is published" do + capacity = [] + response_queue, transport = build_response_queue( + max_in_flight: 2, + on_capacity: ->(available_transport) { capacity << available_transport } + ) + first_sequence = response_queue.reserve + second_sequence = response_queue.reserve + + expect(response_queue.reserve).to be_nil + response_queue.complete(second_sequence, 'second reply') + expect(response_queue.reserve).to be_nil + + response_queue.complete(first_sequence, 'first reply') + + expect(transport).to have_received(:write).with('first reply').ordered + expect(transport).to have_received(:write).with('second reply').ordered + expect(response_queue.reserve).to eq(2) + expect(capacity).to eq([transport]) + end + + it "advances past worker exceptions before publishing later replies" do + response_queue, transport = build_response_queue + worker_queue = Queue.new + logger = Logger.new(IO::NULL) + logger.level = Logger::FATAL + processor = double('processor') + allow(processor).to receive(:process).and_raise(StandardError, 'boom') + worker = Thrift::NonblockingServer::IOManager::Worker.new( + processor, + Thrift::FramedTransportFactory.new, + Thrift::BinaryProtocolFactory.new, + logger, + worker_queue + ) + worker_thread = worker.spawn + failing_sequence = response_queue.reserve + reply_sequence = response_queue.reserve + + worker_queue.push [:frame, response_queue, failing_sequence, framed_message(Thrift::MessageTypes::CALL)] + worker_queue.push [:shutdown] + expect(Timeout.timeout(1) { worker_thread.join }).to be_a(Thread) + expect(processor).to have_received(:process).once + + response_queue.complete(reply_sequence, 'later reply') + + expect(transport).to have_received(:write).with('later reply').once + expect(transport).to have_received(:flush).once end end