raptor 0.21.0 → 0.22.1

data/lib/raptor/detached_body.rb

@@ -0,0 +1,336 @@
+# rbs_inline: enabled
+# frozen_string_literal: true
+
+require "atomic-ruby/atom"
+
+module Raptor
+  # A bounded response body that may be written after its application thread
+  # has returned.
+  #
+  class DetachedBody
+    DEFAULT_MAX_BUFFER_SIZE = 64 * 1024
+
+    # Tracks buffered bytes shared by a group of detached bodies.
+    #
+    class Budget
+      # @rbs @max_size: Integer
+      # @rbs @size: Atom
+
+      # Creates a shared budget with the given byte limit.
+      #
+      # @param max_size [Integer] maximum bytes that may be reserved
+      # @return [void]
+      #
+      # @rbs (Integer max_size) -> void
+      def initialize(max_size)
+        @max_size = max_size
+        @size = Atom.new(0)
+      end
+
+      # Reserves bytes when they fit within the shared limit.
+      #
+      # @param bytes [Integer] number of bytes to reserve
+      # @return [Boolean] whether the bytes were reserved
+      #
+      # @rbs (Integer bytes) -> bool
+      def reserve(bytes)
+        reserved = false
+        @size.swap do |size|
+          reserved = size + bytes <= @max_size
+          reserved ? size + bytes : size
+        end
+        reserved
+      end
+
+      # Releases bytes previously charged to the shared limit.
+      #
+      # @param bytes [Integer] number of bytes to release
+      # @return [void]
+      #
+      # @rbs (Integer bytes) -> void
+      def release(bytes)
+        @size.swap { |size| size - bytes }
+      end
+    end
+
+    # @rbs @max_buffer_size: Integer
+    # @rbs @state: Atom
+    # @rbs @budgets: Array[Budget]
+    # @rbs @wake: ^() -> void
+    # @rbs @dispatch: ^(Proc) -> void
+    # @rbs @finished: ^(Symbol) -> void
+    # @rbs @on_open: (^(DetachedBody) -> void)?
+    # @rbs @on_close: (^(Symbol) -> void)?
+
+    # Creates a detached body with a bounded byte buffer.
+    #
+    # @param max_buffer_size [Integer] maximum body bytes waiting to be written
+    # @return [void]
+    #
+    # @rbs (?max_buffer_size: Integer) -> void
+    def initialize(max_buffer_size: DEFAULT_MAX_BUFFER_SIZE)
+      raise ArgumentError, "max_buffer_size must be positive" unless max_buffer_size.positive?
+
+      @max_buffer_size = max_buffer_size
+      @state = Atom.new({ chunks: [], size: 0, closing: false, closed: false, notified: false, trailers: {} })
+      @budgets = []
+      @wake = proc {}
+      @dispatch = proc { |callback| callback.call }
+      @finished = proc {}
+      @on_open = nil
+      @on_close = nil
+    end
+
+    # Registers a callback that runs on the request thread once the server
+    # accepts the response stream.
+    #
+    # @yieldparam stream [DetachedBody] the opened response stream
+    # @return [DetachedBody]
+    #
+    # @rbs () { (DetachedBody stream) -> void } -> DetachedBody
+    def on_open(&block)
+      @on_open = block
+      self
+    end
+
+    # Registers a callback for when the response stream closes.
+    #
+    # @yieldparam reason [Symbol] why the stream closed
+    # @return [DetachedBody]
+    #
+    # @rbs () { (Symbol reason) -> void } -> DetachedBody
+    def on_close(&block)
+      @on_close = block
+      self
+    end
+
+    # Adds body bytes without waiting for socket or flow-control capacity.
+    #
+    # @param chunk [String] response body bytes
+    # @return [Symbol] `:accepted`, `:full`, or `:closed`
+    #
+    # @rbs (String chunk) -> Symbol
+    def try_write(chunk)
+      raise TypeError, "body must yield String values" unless chunk.is_a?(String)
+
+      chunk = chunk.dup.freeze
+      if chunk.empty?
+        state = @state.value
+        return state[:closed] || state[:closing] ? :closed : :accepted
+      end
+      return :full unless reserve(chunk.bytesize)
+
+      result = :closed
+      wake = false
+      @state.swap do |state|
+        if state[:closed] || state[:closing]
+          result = :closed
+          wake = false
+          state
+        elsif state[:size] + chunk.bytesize > @max_buffer_size
+          result = :full
+          wake = false
+          state
+        else
+          result = :accepted
+          wake = !state[:notified]
+          state.merge(
+            chunks: state[:chunks] + [chunk],
+            size: state[:size] + chunk.bytesize,
+            notified: true
+          )
+        end
+      end
+
+      if result == :accepted
+        @wake.call if wake
+      else
+        release(chunk.bytesize)
+      end
+      result
+    end
+
+    # Finishes the body after all accepted bytes have been written.
+    #
+    # @param trailers [Hash] trailing response headers
+    # @return [void]
+    #
+    # @rbs (?trailers: Hash[String, String | Array[String]]) -> void
+    def close(trailers: {})
+      trailers = trailers.to_h do |name, value|
+        value = value.is_a?(Array) ? value.map { _1.dup.freeze }.freeze : value.dup.freeze
+        [name.dup.freeze, value]
+      end.freeze
+      wake = false
+      @state.swap do |state|
+        if state[:closed] || state[:closing]
+          wake = false
+          next state
+        end
+
+        wake = !state[:notified]
+        state.merge(closing: true, notified: true, trailers: trailers)
+      end
+      @wake.call if wake
+    end
+
+    # Returns whether the stream has closed.
+    #
+    # @return [Boolean]
+    #
+    # @rbs () -> bool
+    def closed?
+      @state.value[:closed]
+    end
+
+    # Connects the body to its reactor-owned response stream.
+    #
+    # @param budgets [Array<Budget>] buffer budgets shared with other bodies
+    # @param wake [Proc] called when buffered data or a close is ready to write
+    # @param dispatch [Proc] schedules application callbacks off the reactor thread
+    # @param finished [Proc] called with the close reason after `on_close`
+    # @return [Boolean] false when the bytes already buffered exceed a budget
+    #
+    # @rbs (Array[Budget] budgets, ^() -> void wake, ^(Proc) -> void dispatch, ^(Symbol) -> void finished) -> bool
+    def attach(budgets, wake, dispatch, finished)
+      @wake = wake
+      @dispatch = dispatch
+      @finished = finished
+      size = @state.value[:size]
+      reserved = []
+      budgets.each do |budget|
+        unless budget.reserve(size)
+          reserved.each { |held| held.release(size) }
+          finish(:full)
+          return false
+        end
+        reserved << budget
+      end
+
+      @budgets = budgets
+      true
+    end
+
+    # Notifies the application that its response stream is ready.
+    #
+    # @return [void]
+    #
+    # @rbs () -> void
+    def open
+      @on_open&.call(self)
+      @wake.call if @state.value[:notified]
+    end
+
+    # Returns the size of the next buffered chunk, 0 when closing, or nil
+    # while waiting for more data.
+    #
+    # @return [Integer, nil]
+    #
+    # @rbs () -> Integer?
+    def next_size
+      state = @state.value
+      state[:chunks].first&.bytesize || (0 if state[:closing])
+    end
+
+    # Removes up to `max_bytes` from the next buffered chunk.
+    #
+    # @param max_bytes [Integer] the largest chunk to return
+    # @return [String, nil] the removed bytes, or nil when nothing is buffered
+    #
+    # @rbs (Integer max_bytes) -> String?
+    def shift(max_bytes)
+      chunk = nil
+      @state.swap do |state|
+        chunk = nil
+        first = state[:chunks].first
+        next state unless first
+
+        size = max_bytes < first.bytesize ? max_bytes : first.bytesize
+        chunk = first.byteslice(0, size)
+        chunks = if size == first.bytesize
+          state[:chunks].drop(1)
+        else
+          [first.byteslice(size..-1).freeze] + state[:chunks].drop(1)
+        end
+        state.merge(
+          chunks: chunks,
+          size: state[:size] - size,
+          notified: !chunks.empty? || state[:closing]
+        )
+      end
+      release(chunk.bytesize) if chunk
+      chunk
+    end
+
+    # Returns the response trailers supplied when the body closed.
+    #
+    # @return [Hash] trailing response headers
+    #
+    # @rbs () -> Hash[String, String | Array[String]]
+    def trailers
+      @state.value[:trailers]
+    end
+
+    # Closes the stream and invokes its callback exactly once.
+    #
+    # @param reason [Symbol] why the stream closed
+    # @return [void]
+    #
+    # @rbs (Symbol reason) -> void
+    def finish(reason)
+      callback = false
+      remaining = 0
+      @state.swap do |state|
+        if state[:closed]
+          callback = false
+          remaining = 0
+          state
+        else
+          callback = true
+          remaining = state[:size]
+          state.merge(chunks: [], size: 0, closed: true, notified: false)
+        end
+      end
+      return unless callback
+
+      release(remaining)
+      @dispatch.call(proc do
+        begin
+          @on_close&.call(reason)
+        ensure
+          @finished.call(reason)
+        end
+      end)
+    end
+
+    private
+
+    # Reserves bytes from every budget attached to this body.
+    #
+    # @param bytes [Integer] number of bytes to reserve
+    # @return [Boolean] whether every budget accepted the reservation
+    #
+    # @rbs (Integer bytes) -> bool
+    def reserve(bytes)
+      reserved = []
+      @budgets.each do |budget|
+        unless budget.reserve(bytes)
+          reserved.each { |held| held.release(bytes) }
+          return false
+        end
+        reserved << budget
+      end
+      true
+    end
+
+    # Releases bytes from every budget attached to this body.
+    #
+    # @param bytes [Integer] number of bytes to release
+    # @return [void]
+    #
+    # @rbs (Integer bytes) -> void
+    def release(bytes)
+      @budgets.each { |budget| budget.release(bytes) }
+    end
+  end
+end

data/lib/raptor/http.rb

@@ -27,6 +27,16 @@ module Raptor
       def message = "could not write response"
     end
 
+    # Returns whether an HTTP status forbids an entity body.
+    #
+    # @param status [Integer] the response status code
+    # @return [Boolean]
+    #
+    # @rbs (Integer status) -> bool
+    def self.no_entity_body_status?(status)
+      (status >= 100 && status < 200) || status == 204 || status == 304
+    end
+
     # Writes `string` in full, retrying on partial writes. Bounded by
     # `timeout` so a slow client can't pin the writing thread.
     #

data/lib/raptor/http1.rb

@@ -8,6 +8,7 @@ require "tempfile"
 require "atomic-ruby/atomic_boolean"
 require "rack"
 
+require_relative "detached_body"
 require_relative "http"
 require_relative "raptor_http"
 require_relative "thread_locals"
@@ -97,6 +98,34 @@ module Raptor
       false
     end
 
+    # Encodes one HTTP/1.1 response chunk.
+    #
+    # @param chunk [String] response body bytes
+    # @return [String] the encoded chunk
+    #
+    # @rbs (String chunk) -> String
+    def self.encode_chunk(chunk)
+      HttpParser.chunked_encode(String.new, chunk)
+    end
+
+    # Encodes the final HTTP/1.1 response chunk and its trailers.
+    #
+    # @param trailers [Hash] trailing response headers
+    # @return [String] the final chunk followed by the trailer section
+    #
+    # @rbs (Hash[String, String | Array[String]] trailers) -> String
+    def self.encode_trailers(trailers)
+      normalized = trailers.each_with_object({}) do |(key, value), headers|
+        name = key.match?(/[A-Z]/) ? key.downcase : key
+        next if name.start_with?(RACK_HEADER_PREFIX) || name == "status"
+
+        headers[name] = value
+      end
+      response = +"0\r\n"
+      HttpParser.format_headers(response, normalized)
+      response << "\r\n"
+    end
+
     # Advances an HTTP/1.x request parse from the state hash's buffered
     # bytes. A complete well-formed request returns with `:complete` set
     # plus populated `:env` and `:body`; malformed or oversized input flips
@@ -492,7 +521,7 @@ module Raptor
     #
     # @rbs (TCPSocket socket, Integer id, Hash[String, untyped] env, Hash[Symbol, untyped] parse_data, String? body, Reactor reactor, AtomicThreadPool thread_pool, Integer request_count, String remote_addr, String url_scheme) -> void
     def process_client(socket, id, env, parse_data, body, reactor, thread_pool, request_count, remote_addr, url_scheme)
-      keep_alive = process_request(socket, env, parse_data, body, request_count, remote_addr, url_scheme)
+      keep_alive = process_request(socket, id, env, parse_data, body, reactor, request_count, remote_addr, url_scheme)
       eager_keepalive(socket, id, reactor, thread_pool, request_count, remote_addr, url_scheme) if keep_alive
     end
 
@@ -501,20 +530,22 @@ module Raptor
     # for another request.
     #
     # @param socket [TCPSocket] the client socket
+    # @param id [Integer] unique client identifier
     # @param env [Hash] partial env hash from the HTTP parser
     # @param parse_data [Hash] metadata from the parsing pass
     # @param body [String, nil] decoded request body
+    # @param reactor [Reactor] the reactor managing the client connection
     # @param request_count [Integer] number of requests handled on this connection
     # @param remote_addr [String] client IP address
     # @param url_scheme [String] "http" or "https"
     # @return [Boolean] true if the connection should be kept alive
     #
-    # @rbs (TCPSocket socket, Hash[String, untyped] env, Hash[Symbol, untyped] parse_data, String? body, Integer request_count, String remote_addr, String url_scheme) -> bool
-    def process_request(socket, env, parse_data, body, request_count, remote_addr, url_scheme)
+    # @rbs (TCPSocket socket, Integer id, Hash[String, untyped] env, Hash[Symbol, untyped] parse_data, String? body, Reactor reactor, Integer request_count, String remote_addr, String url_scheme) -> bool
+    def process_request(socket, id, env, parse_data, body, reactor, request_count, remote_addr, url_scheme)
       if @clean_fiber_locals
-        Fiber.new { perform_request(socket, env, parse_data, body, request_count, remote_addr, url_scheme) }.resume
+        Fiber.new { perform_request(socket, id, env, parse_data, body, reactor, request_count, remote_addr, url_scheme) }.resume
       else
-        perform_request(socket, env, parse_data, body, request_count, remote_addr, url_scheme)
+        perform_request(socket, id, env, parse_data, body, reactor, request_count, remote_addr, url_scheme)
       end
     ensure
       ThreadLocals.clear if @clean_thread_locals
@@ -522,13 +553,25 @@ module Raptor
 
     # Calls the Rack app and writes its response for one request.
     #
-    # @rbs (TCPSocket socket, Hash[String, untyped] env, Hash[Symbol, untyped] parse_data, String? body, Integer request_count, String remote_addr, String url_scheme) -> bool
-    def perform_request(socket, env, parse_data, body, request_count, remote_addr, url_scheme)
+    # @param socket [TCPSocket] the client socket
+    # @param id [Integer] unique client identifier
+    # @param env [Hash] partial env hash from the HTTP parser
+    # @param parse_data [Hash] metadata from the parsing pass
+    # @param body [String, nil] decoded request body
+    # @param reactor [Reactor] the reactor managing the client connection
+    # @param request_count [Integer] number of requests handled on this connection
+    # @param remote_addr [String] client IP address
+    # @param url_scheme [String] "http" or "https"
+    # @return [Boolean] true if the connection should be kept alive
+    #
+    # @rbs (TCPSocket socket, Integer id, Hash[String, untyped] env, Hash[Symbol, untyped] parse_data, String? body, Reactor reactor, Integer request_count, String remote_addr, String url_scheme) -> bool
+    def perform_request(socket, id, env, parse_data, body, reactor, request_count, remote_addr, url_scheme)
       rack_env = nil
       status = nil
       headers = nil
       hijacked = false
       keep_alive = false
+      detached = false
       response_started = false
 
       begin
@@ -540,16 +583,40 @@ module Raptor
           body.close if body.respond_to?(:close)
         else
           hijacked = headers.is_a?(Hash) && !!headers[Rack::RACK_HIJACK]
-          streaming = body.respond_to?(:call) && !body.respond_to?(:each)
-          keep_alive = (hijacked || streaming) ? false : keep_alive?(rack_env, request_count)
-          response_size = response_size(headers, body) if @access_log_io && !hijacked
-          response_started = true
-          write_response(socket, rack_env, status, headers, body, keep_alive: keep_alive)
+          no_body = rack_env[Rack::REQUEST_METHOD] == "HEAD" || Http.no_entity_body_status?(status)
+          detached = body.is_a?(DetachedBody) && !no_body
+          if detached
+            chunked = rack_env[Rack::SERVER_PROTOCOL] == HTTP_11
+            keep_alive = chunked && keep_alive?(rack_env, request_count)
+            response_started = true
+            write_detached_response(socket, status, headers, chunked: chunked, keep_alive: keep_alive)
+            finished = proc do |reason|
+              error = IOError.new("detached response closed") unless reason == :closed
+              write_access_log(rack_env, status, "-", remote_addr) if @access_log_io
+              Http.call_response_finished(rack_env, status, headers, error)
+            end
+            reactor.attach_http1_body(
+              socket,
+              id,
+              body,
+              { chunked: chunked, keep_alive: keep_alive, request_count: request_count, remote_addr: remote_addr, url_scheme: url_scheme },
+              finished
+            )
+          else
+            streaming = body.respond_to?(:call) && !body.respond_to?(:each)
+            keep_alive = (hijacked || streaming) ? false : keep_alive?(rack_env, request_count)
+            response_size = response_size(headers, body) if @access_log_io && !hijacked
+            response_started = true
+            write_response(socket, rack_env, status, headers, body, keep_alive: keep_alive)
+            body.finish(:closed) if body.is_a?(DetachedBody)
+          end
         end
 
-        write_access_log(rack_env, status, response_size, remote_addr) if @access_log_io && !hijacked
-        Http.call_response_finished(rack_env, status, headers, nil)
-        keep_alive && !hijacked
+        unless detached
+          write_access_log(rack_env, status, response_size, remote_addr) if @access_log_io && !hijacked
+          Http.call_response_finished(rack_env, status, headers, nil)
+        end
+        keep_alive && !hijacked && !detached
       rescue => error
         keep_alive = false
         handle_app_error(socket, rack_env, status, headers, error, response_started: response_started, hijacked: hijacked)
@@ -557,7 +624,7 @@ module Raptor
         rack_input = rack_env && rack_env[Rack::RACK_INPUT]
         rack_input.close! rescue nil if rack_input.respond_to?(:close!)
 
-        unless hijacked || keep_alive
+        unless hijacked || keep_alive || detached
           socket.close rescue nil
         end
       end
@@ -666,9 +733,11 @@ module Raptor
 
         keep_alive = process_request(
           socket,
+          id,
           env,
           parse_data,
           body,
+          reactor,
           request_count,
           remote_addr,
           url_scheme
@@ -895,6 +964,31 @@ module Raptor
       socket_write(socket, response)
     end
 
+    # Starts a detached response, chunked on HTTP/1.1 and ended by closing
+    # the connection on HTTP/1.0.
+    #
+    # @param socket [TCPSocket] the client socket to write to
+    # @param status [Integer] HTTP status code
+    # @param headers [Hash] response headers from the Rack application
+    # @param chunked [Boolean] whether to use chunked transfer encoding
+    # @param keep_alive [Boolean] whether to send a keep-alive connection header
+    # @return [void]
+    #
+    # @rbs (TCPSocket socket, Integer status, Hash[String, String | Array[String]] headers, chunked: bool, keep_alive: bool) -> void
+    def write_detached_response(socket, status, headers, chunked:, keep_alive:)
+      validate_status(status)
+      headers = normalize_headers(headers)
+      validate_headers(headers, status, false)
+      headers.delete(Rack::CONTENT_LENGTH)
+      headers[Rack::TRANSFER_ENCODING] = TRANSFER_ENCODING_CHUNKED if chunked
+      headers["connection"] = keep_alive ? CONNECTION_KEEPALIVE : CONNECTION_CLOSE
+
+      response = build_status_line(chunked ? HTTP_11 : HTTP_10, status)
+      HttpParser.format_headers(response, headers)
+      response << "\r\n"
+      socket_write(socket, response)
+    end
+
     # Writes a complete HTTP response to the socket.
     #
     # @param socket [TCPSocket] the client socket to write to
@@ -910,7 +1004,7 @@ module Raptor
       validate_status(status)
       response_hijack = headers.is_a?(Hash) ? headers.delete(Rack::RACK_HIJACK) : nil
       headers = normalize_headers(headers)
-      no_entity_body = (status >= 100 && status < 200) || status == 204 || status == 304
+      no_entity_body = Http.no_entity_body_status?(status)
       validate_headers(headers, status, no_entity_body)
 
       headers["connection"] = keep_alive ? CONNECTION_KEEPALIVE : CONNECTION_CLOSE