Read large values without 8 KB chunking in TimeoutRecoveryReader; fix TLS EOF replay

CHANGELOG.md

@@ -7,6 +7,15 @@ Unreleased
 - Add experimental per-request `exp_async_timeout_recovery` option: reads wait with explicit deadlines and poll the
   socket once more before timing out, so replies delayed only by a busy fiber scheduler are not discarded. (nherson)
 
+- `exp_async_timeout_recovery` reads values of 8 KB or more straight from the socket in reads of up to 64 KiB
+  instead of through the 8 KB line buffer, cutting a 2 MB get from ~259 to ~37 syscalls. (nherson)
+
+- Fix `exp_async_timeout_recovery` replaying the previous reply instead of reporting EOF after a TLS server
+  closes the connection cleanly. (nherson)
+
+- Fix `exp_async_timeout_recovery` failing a read whose data is ready when the poll after a timed-out wait
+  misses it, which a zero-timeout poll can do while the thread has a pending interrupt. (nherson)
+
 - Add per-request opaque-correlation controls and a client `opaque_correlation_request_only` mode
   for opt-in rollouts, with eager PRNG initialization and unchanged unflagged reads. (mrattle)
 

bin/benchmark

@@ -7,11 +7,14 @@
 # run with:
 # bundle exec bin/benchmark
 # RUBY_YJIT_ENABLE=1 BENCH_TARGET=get bundle exec bin/benchmark
+# BENCH_TARGET=async_get bundle exec bin/benchmark (not part of BENCH_TARGET=all)
 require 'bundler/inline'
 require 'json'
 
 gemfile do
   source 'https://rubygems.org'
+  # Only async_get needs it; other runs keep the same gems as before that target existed.
+  gem 'async' if ENV['BENCH_TARGET'] == 'async_get'
   gem 'benchmark-ips'
   gem 'logger'
   gem 'opentelemetry-sdk'
@@ -234,3 +237,16 @@ if %w[all delete_multi].include?(bench_target)
     x.compare!
   end
 end
+
+# Large-value gets under a fiber scheduler, through the default read path and the
+# exp_async_timeout_recovery reader. Runs only when selected, so BENCH_TARGET=all is unchanged.
+if bench_target == 'async_get'
+  Sync do
+    Benchmark.ips do |x|
+      x.config(warmup: bench_warmup, time: bench_time)
+      x.report('get dalli async') { client.get('key') }
+      x.report('get dalli async timeout recovery') { client.get('key', exp_async_timeout_recovery: true) }
+      x.compare!
+    end
+  end
+end

lib/dalli/protocol/timeout_recovery_reader.rb

@@ -8,6 +8,9 @@ module Protocol
     ##
     class TimeoutRecoveryReader
       CHUNK_SIZE = 8192
+      # Most a single large-value read asks for. Bounds the per-connection scratch
+      # buffer at 64 KiB while cutting a 2 MB value from ~257 reads to ~35.
+      LARGE_READ_SIZE = 65_536
 
       def initialize(sock, timeout)
         @sock = sock
@@ -36,7 +39,13 @@ def read(count)
         data = String.new(capacity: count, encoding: Encoding::BINARY)
         data << @buffer.byteslice(@offset, available) if available.positive?
         @offset = @buffer.bytesize
-        append_from_socket(data, count - available)
+        remaining = count - available
+        # Large values bypass the 8 KB line buffer: each read asks the socket for at
+        # most what is still owed, so a value costs a few syscalls rather than one per
+        # CHUNK_SIZE bytes plus a copy out of @buffer, and never reads past its end.
+        return append_directly(data, remaining) if remaining >= CHUNK_SIZE
+
+        append_from_socket(data, remaining)
       end
 
       private
@@ -59,31 +68,56 @@ def append_from_socket(data, remaining)
         data
       end
 
+      def append_directly(data, remaining)
+        while remaining.positive?
+          chunk = if data.empty?
+                    read_partial(remaining, data)
+                  else
+                    read_partial([remaining, LARGE_READ_SIZE].min, scratch)
+                  end
+          return data.empty? ? nil : data unless chunk
+
+          data << chunk unless chunk.equal?(data)
+          remaining -= chunk.bytesize
+        end
+        data
+      end
+
+      def scratch
+        @scratch ||= String.new(encoding: Encoding::BINARY)
+      end
+
+      # Reads at most maxlen bytes into outbuf, waiting while the socket has none; nil at EOF.
+      def read_partial(maxlen, outbuf)
+        last_try = false
+        while (result = @sock.read_nonblock(maxlen, outbuf, exception: false)).is_a?(Symbol)
+          raise IO::TimeoutError, "IO timeout after #{@timeout}s" if last_try
+
+          last_try = !wait_for(result)
+        end
+        result
+      end
+
       # read_nonblock overwrites the buffer, so unconsumed bytes are set aside and restored
-      # after the read; @offset only moves once the read returns.
+      # after the read; @offset only moves once the read returns. At EOF a socket empties the
+      # buffer but OpenSSL leaves it untouched, so it is cleared here to avoid replaying old bytes.
       def fill
         partial = @buffer.byteslice(@offset, @buffer.bytesize - @offset) if pending?
-        filled = read_chunk
+        filled = !read_partial(CHUNK_SIZE, @buffer).nil?
+        @buffer.clear unless filled
         @offset = 0
-        if partial
-          filled ? @buffer.prepend(partial) : @buffer.replace(partial)
-        end
+        @buffer.prepend(partial) if partial
         filled
       end
 
-      # Waits on :wait_readable or :wait_writable, which name the socket method to call.
-      def read_chunk
-        while (result = @sock.read_nonblock(CHUNK_SIZE, @buffer, exception: false)).is_a?(Symbol)
-          wait_for(result)
-        end
-        !result.nil?
-      end
-
+      # Waits on :wait_readable or :wait_writable, which name the socket method to call. After a
+      # wait that times out, polls once outside the fiber scheduler. Returns false when that poll
+      # finds nothing too, and the caller then reads once more before failing: a zero-timeout poll
+      # can report nothing without checking the socket while the thread has a pending interrupt.
       def wait_for(event)
-        return if @sock.public_send(event, @timeout)
-        return if Fiber.blocking { @sock.public_send(event, 0) }
+        return true if @sock.public_send(event, @timeout)
 
-        raise IO::TimeoutError, "IO timeout after #{@timeout}s"
+        Fiber.blocking { @sock.public_send(event, 0) } ? true : false
       end
     end
   end

test/integration/test_async_timeout_recovery.rb

@@ -7,8 +7,22 @@
 
 # Replies from a separate process, so they arrive while a fiber in this process holds the thread.
 # A key named "delay<ms>-..." delays its reply; every get writes one byte to #signal on receipt.
+# A key containing "size<n>" gets an n-byte value; with "parts<n>" and "gap<ms>" the value is
+# written in n parts with a pause before each later part. Halfway through each pause, once the
+# client has drained the previous part and is waiting, a byte is written to #signal.
 class DelayedReplyServer
-  SCRIPT = <<~'RUBY'
+  VALUE_OF = <<~'RUBY'
+    def value_of(key)
+      size = key[/size(\d+)/, 1]&.to_i
+      return "value-of-#{key}" unless size
+
+      pattern = "EN\r\nVA 3 k\r\n\0\xFF".b
+      (pattern * ((size / pattern.bytesize) + 1)).byteslice(0, size)
+    end
+  RUBY
+  singleton_class.class_eval(VALUE_OF)
+
+  SCRIPT = VALUE_OF + <<~'RUBY'
     require 'socket'
     listener = UNIXServer.new(ARGV.fetch(0))
     signal = IO.for_fd(3, autoclose: false)
@@ -24,8 +38,19 @@ class DelayedReplyServer
           when 'mg'
             signal.write('.')
             sleep(key[/\Adelay(\d+)ms/, 1].to_i / 1000.0)
-            value = "value-of-#{key}"
-            connection.write("VA #{value.bytesize}#{" k#{key}" if flags.include?('k')}\r\n#{value}\r\n")
+            value = value_of(key)
+            parts = key[/parts(\d+)/, 1]&.to_i || 1
+            part_size = (value.bytesize + parts - 1) / parts
+            first, *rest = Array.new(parts) { |i| value.byteslice(i * part_size, part_size) }
+            connection.write("VA #{value.bytesize}#{" k#{key}" if flags.include?('k')}\r\n", first,
+                             rest.empty? ? "\r\n" : '')
+            pause = key[/gap(\d+)ms/, 1].to_i / 2000.0
+            rest.each_with_index do |part, index|
+              sleep(pause)
+              signal.write('.')
+              sleep(pause)
+              connection.write(part, index == rest.size - 1 ? "\r\n" : '')
+            end
           else raise "Unexpected command: #{command}"
           end
         end
@@ -144,6 +169,106 @@ def elapsed
   end
 end
 
+module LargeValueHelpers
+  # More than a UNIX socket buffer holds, so the server's write of a first half returns only once the
+  # client is reading inside the value.
+  LARGE_SIZE = 900_000
+
+  ##
+  # Records reads and waits on a connection's socket.
+  ##
+  class SocketLog
+    def self.now
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+
+    def initialize(sock)
+      events = @events = []
+      sock.singleton_class.prepend(Module.new do
+        define_method(:read_nonblock) do |*args, **kwargs|
+          result = super(*args, **kwargs)
+          events << [:read, result.is_a?(String) ? result.bytesize : 0, SocketLog.now]
+          result
+        end
+
+        define_method(:wait_readable) do |timeout = nil|
+          result = super(timeout)
+          events << [:wait, timeout, result ? true : false]
+          result
+        end
+      end)
+    end
+
+    def last_data_at
+      @events.reverse_each.find { |kind, bytes| kind == :read && bytes.positive? }&.last
+    end
+
+    # Bytes received before each wait that reported a timeout after which the reader still found
+    # data: through the poll that follows or, when that poll misses it, the read after the poll.
+    def recoveries(timeout)
+      received = 0
+      @events.each_cons(3).filter_map do |event, poll, after|
+        received += event[1] if event.first == :read
+        received if event == [:wait, timeout, false] && found_data?(poll, after)
+      end
+    end
+
+    private
+
+    def found_data?(poll, after)
+      poll == [:wait, 0, true] || (poll == [:wait, 0, false] && after.first == :read && after[1].positive?)
+    end
+  end
+
+  def in_scheduler(selector, &)
+    return yield unless selector
+
+    with_selector(selector, &)
+  end
+
+  def assert_reads_through_pauses(client, selector)
+    key = "size#{LARGE_SIZE}-parts6-gap30ms"
+    value = nil
+    duration = in_scheduler(selector) do
+      client.alive!
+      elapsed { value = client.get(key, recover) }
+    end
+
+    assert_equal DelayedReplyServer.value_of(key), value
+    # Five 30ms pauses: none outlasts the timeout, together they do.
+    assert_operator duration, :>, 2 * socket_timeout
+  end
+
+  def assert_fails_within_timeout_of_last_byte(client, selector)
+    error = hang = nil
+    in_scheduler(selector) do
+      client.alive!
+      log = SocketLog.new(connection_manager(client).sock)
+      client.get("size#{LARGE_SIZE}-parts2-gap500ms", recover)
+    rescue Dalli::NetworkError => e
+      error = e
+      hang = SocketLog.now - log.last_data_at
+    end
+
+    assert_kind_of Dalli::NetworkError, error
+    assert_match(/IO::TimeoutError/, error.message)
+    assert_operator hang, :>=, socket_timeout
+    assert_operator hang, :<, socket_timeout + 0.1
+    key = "size#{LARGE_SIZE}-after"
+
+    assert_equal DelayedReplyServer.value_of(key), in_scheduler(selector) { client.get(key, recover) }
+  end
+
+  # A fiber that wakes on the server's second signal, sent while the reader waits inside the value
+  # for its second half, and holds the thread until after that half arrives.
+  def start_busy_fiber_inside_value(task, signal)
+    task.async do
+      2.times { await_signal(signal) }
+      busy
+    end
+  end
+end
+
 module RecoveryParityHelpers
   MULTI_KEYS = (Array.new(40) { |i| "multi-#{i}" } << 'missing').freeze
 
@@ -188,6 +313,7 @@ def assert_same_results(client)
 describe 'exp_async_timeout_recovery' do
   include AsyncTimeoutRecoveryHelpers
   include RecoveryParityHelpers
+  include LargeValueHelpers
 
   describe 'without a fiber scheduler' do
     it 'returns the same results as default reads over TCP' do
@@ -247,6 +373,14 @@ def assert_same_results(client)
         assert_equal 'value-of-b', client.get('b', recover)
       end
     end
+
+    it 'reads a large value delivered in pauses that are each shorter than the timeout' do
+      with_delayed_reply_client { |client| assert_reads_through_pauses(client, nil) }
+    end
+
+    it 'fails a large value within one timeout once it stops arriving' do
+      with_delayed_reply_client { |client| assert_fails_within_timeout_of_last_byte(client, nil) }
+    end
   end
 
   AsyncTimeoutRecoveryHelpers::SELECTORS.each do |selector|
@@ -297,6 +431,51 @@ def assert_same_results(client)
         end
       end
 
+      it 'keeps reading a large value whose rest arrives while another fiber holds the thread' do
+        with_delayed_reply_client do |client, server|
+          key = "size#{LargeValueHelpers::LARGE_SIZE}-parts2-gap40ms"
+          value = sock = log = nil
+          duration = with_selector(selector) do |task|
+            client.alive!
+            sock = connection_manager(client).sock
+            log = LargeValueHelpers::SocketLog.new(sock)
+            start_busy_fiber_inside_value(task, server.signal)
+            elapsed { value = client.get(key, recover) }
+          end
+          recoveries = log.recoveries(socket_timeout)
+
+          assert_equal DelayedReplyServer.value_of(key), value
+          assert_operator duration, :>, socket_timeout
+          assert_same sock, connection_manager(client).sock
+          assert recoveries.any? { |received| received.between?(100, LargeValueHelpers::LARGE_SIZE) },
+                 "expected a wait inside the value to time out and find data; recovered after #{recoveries} bytes"
+        end
+      end
+
+      it 'fails a large value without the option when another fiber holds the thread inside it' do
+        with_delayed_reply_client do |client, server|
+          error = with_selector(selector) do |task|
+            client.alive!
+            start_busy_fiber_inside_value(task, server.signal)
+            client.get("size#{LargeValueHelpers::LARGE_SIZE}-parts2-gap40ms")
+            nil
+          rescue Dalli::NetworkError => e
+            e
+          end
+
+          assert_kind_of Dalli::NetworkError, error
+          assert_match(/IO::TimeoutError/, error.message)
+        end
+      end
+
+      it 'reads a large value delivered in pauses that are each shorter than the timeout' do
+        with_delayed_reply_client { |client| assert_reads_through_pauses(client, selector) }
+      end
+
+      it 'fails a large value within one timeout once it stops arriving' do
+        with_delayed_reply_client { |client| assert_fails_within_timeout_of_last_byte(client, selector) }
+      end
+
       it 'still fails a read when no reply arrives before the timeout' do
         with_delayed_reply_client do |client|
           error = nil