Merge pull request #1114 from byroot/response-buffer-clear

Author: petergoldstein

.rubocop_todo.yml

@@ -23,6 +23,7 @@ Metrics/MethodLength:
   Exclude:
     - 'lib/dalli/pipelined_getter.rb'
     - 'lib/dalli/protocol/base.rb'
+    - 'lib/dalli/socket.rb'
 
 # Offense count: 1
 # Configuration parameters: CountComments, CountAsOne.

lib/dalli/protocol/connection_manager.rb

@@ -193,8 +193,8 @@ def flushed_write(bytes)
 
       # Non-blocking read.  Here to support the operation
       # of the get_multi operation
-      def read_available
-        @sock.read_available
+      def read_available(...)
+        @sock.read_available(...)
       end
 
       def max_allowed_failures

lib/dalli/protocol/response_buffer.rb

@@ -23,22 +23,34 @@ def initialize(io_source, response_processor)
       end
 
       def read
-        @buffer << @io_source.read_available
+        remaining_bytes = @buffer.bytesize - @offset
+        if remaining_bytes.zero?
+          @offset = 0
+          @buffer = @io_source.read_available(@buffer)
+        else
+          if remaining_bytes > COMPACT_THRESHOLD && remaining_bytes < @buffer.bytesize
+            # NB: we're freezing the old buffer before slicing so Ruby doesn't
+            # have to allocate a third hidden string to be the Copy-on-Write onwer.
+            @buffer = @buffer.freeze.byteslice(@offset..)
+            @offset = 0
+          end
+          @buffer << @io_source.read_available
+        end
       end
 
       # Attempts to process a single response from the buffer,
       # advancing the offset past the consumed bytes.
       def process_single_getk_response
         bytes, status, cas, key, value = @response_processor.getk_response_from_buffer(@buffer, @offset)
         @offset += bytes
-        compact_if_needed
         [status, cas, key, value]
       end
 
       # Resets the internal buffer to an empty state,
       # so that we're ready to read pipelined responses
       def reset
-        @buffer = ''.b
+        @buffer&.clear
+        @buffer ||= ''.b
         @offset = 0
       end
 
@@ -54,25 +66,14 @@ def ensure_ready
 
       # Clear the internal response buffer
       def clear
+        @buffer&.clear
         @buffer = nil
         @offset = 0
       end
 
       def in_progress?
         !@buffer.nil?
       end
-
-      private
-
-      # Only compact when we've consumed a significant portion of the buffer.
-      # This avoids per-response string allocation while preventing unbounded
-      # memory growth for large pipelines.
-      def compact_if_needed
-        return unless @offset > COMPACT_THRESHOLD && @offset > @buffer.bytesize / 2
-
-        @buffer = @buffer.byteslice(@offset..)
-        @offset = 0
-      end
     end
   end
 end

lib/dalli/socket.rb

@@ -12,48 +12,64 @@ module Socket
     # Common methods for all socket implementations.
     ##
     module InstanceMethods
-      def readfull(count)
-        value = String.new(capacity: count + 1)
-        loop do
-          result = read_nonblock(count - value.bytesize, exception: false)
-          value << result if append_to_buffer?(result)
-          break if value.bytesize == count
+      def read_available(reusable_buffer = nil)
+        if reusable_buffer
+          value = read_nonblock(8196, reusable_buffer, exception: false)
+          case value
+          when :wait_writable, :wait_readable
+            return reusable_buffer.clear
+          when nil
+            raise Errno::ECONNRESET, "Connection reset: #{logged_options.inspect}"
+          end
+        else
+          value = ''.b
         end
-        value
-      end
 
-      def read_available
-        value = +''
+        buffer = ''.b
         loop do
-          result = read_nonblock(8196, exception: false)
-          break if WAIT_RCS.include?(result)
-          raise Errno::ECONNRESET, "Connection reset: #{logged_options.inspect}" unless result
-
-          value << result
+          result = read_nonblock(8196, buffer, exception: false)
+          case result
+          when :wait_writable, :wait_readable
+            buffer.clear
+            return value
+          when nil
+            raise Errno::ECONNRESET, "Connection reset: #{logged_options.inspect}"
+          else
+            value << result
+          end
         end
-        value
-      end
-
-      WAIT_RCS = %i[wait_writable wait_readable].freeze
-
-      def append_to_buffer?(result)
-        raise Timeout::Error, "IO timeout: #{logged_options.inspect}" if nonblock_timed_out?(result)
-        raise Errno::ECONNRESET, "Connection reset: #{logged_options.inspect}" unless result
-
-        !WAIT_RCS.include?(result)
-      end
-
-      def nonblock_timed_out?(result)
-        return true if result == :wait_readable && !wait_readable(options[:socket_timeout])
-
-        # TODO: Do we actually need this?  Looks to be only used in read_nonblock
-        result == :wait_writable && !wait_writable(options[:socket_timeout])
       end
 
       FILTERED_OUT_OPTIONS = %i[username password].freeze
       def logged_options
         options.except(*FILTERED_OUT_OPTIONS)
       end
+
+      # JRuby doesn't support IO#timeout=, so use custom readfull implementation
+      # CRuby 3.3+ has IO#timeout= which makes IO#read work with timeouts
+      if RUBY_ENGINE == 'jruby'
+        # rubocop:disable Metrics/AbcSize
+        def readfull(count)
+          value = String.new(capacity: count + 1)
+
+          until value.bytesize == count
+            result = read_nonblock(count - value.bytesize, exception: false)
+            case result
+            when :wait_readable
+              wait_readable(options[:socket_timeout]) or raise Timeout::Error, "IO timeout: #{logged_options.inspect}"
+            when :wait_writable
+              wait_writable(options[:socket_timeout]) or raise Timeout::Error, "IO timeout: #{logged_options.inspect}"
+            when nil
+              raise Errno::ECONNRESET, "Connection reset: #{logged_options.inspect}"
+            else
+              value << result
+            end
+          end
+
+          value
+        end
+        # rubocop:enable Metrics/AbcSize
+      end
     end
 
     ##