puma-plus 0.1.0

data/lib/puma_plus/worker_thread.rb

@@ -0,0 +1,331 @@
+# frozen_string_literal: true
+
+require "socket"
+require "puma_plus/wire"
+require "puma_plus/rack_env"
+require "puma_plus/input"
+require "puma_plus/gvl"
+
+module PumaPlus
+  # One Ruby thread serving one connection to the Go server.
+  #
+  # This replaces the seam puma occupies with Server#process_client plus
+  # Response#handle_request. There is no thread pool and no reactor: the thread
+  # dials in, blocks on read, serves one request, and blocks again. Because the
+  # connection carries exactly one request at a time, nothing here needs a mutex
+  # or a demultiplexer.
+  #
+  # An idle connection *is* the capacity signal the Go dispatcher counts, so
+  # "blocked in read" and "available" are the same state -- there is no separate
+  # bookkeeping to drift out of sync.
+  class WorkerThread
+    MONOTONIC = Process::CLOCK_MONOTONIC
+    THREAD_CPU = Process::CLOCK_THREAD_CPUTIME_ID
+
+    def initialize(socket_path:, app:, worker_id:, thread_index:,
+                   multithread: true, multiprocess: false, logger: $stderr)
+      @socket_path = socket_path
+      @app = app
+      @worker_id = worker_id
+      @thread_index = thread_index
+      @logger = logger
+      @env_builder = RackEnv.new(multithread: multithread, multiprocess: multiprocess)
+    end
+
+    def run
+      conn = UNIXSocket.new(@socket_path)
+      conn.sync = true
+      hello(conn)
+
+      loop do
+        begin
+          type, payload = Wire.read_frame(conn)
+        rescue Wire::Closed
+          break
+        end
+
+        case type
+        when Wire::REQUEST
+          break if serve(conn, payload) == :hijacked
+        when Wire::GOAWAY
+          # Sent only to an idle conn, so there is never in-flight work to lose.
+          break
+        else
+          @logger.puts "[puma-plus] unexpected frame #{Wire.type_name(type)}; closing"
+          break
+        end
+      end
+    rescue Errno::ECONNRESET, Errno::EPIPE => e
+      @logger.puts "[puma-plus] worker #{@worker_id}/#{@thread_index} lost the server: #{e.class}"
+    ensure
+      conn&.close
+    end
+
+    private
+
+    # The readiness signal. Sent only once the Rack app is loaded and callable,
+    # so the Go side never has to poll for readiness.
+    def hello(conn)
+      conn.write Wire.frame(Wire::HELLO, Wire.encode_kv([
+        ["version", Wire::VERSION],
+        ["role", "data"],
+        ["pid", Process.pid],
+        ["worker_id", @worker_id],
+        ["thread_index", @thread_index]
+      ]))
+    end
+
+    def serve(conn, payload)
+      meta, offset = Wire.decode_request_meta(payload)
+      pairs, offset = Wire.decode_kv(payload, offset)
+      body = payload.byteslice(offset, payload.bytesize - offset)
+
+      env = @env_builder.build(pairs, meta)
+      input = Input.for(meta[:body_mode], body, conn)
+      env["rack.input"] = input
+
+      # Rack full hijack. Calling env['rack.hijack'] hands the application the
+      # raw socket; from that moment Go stops framing and splices bytes straight
+      # to the client, so the app owns the whole conversation including writing
+      # its own status line. That is what makes WebSockets work without
+      # puma-plus knowing anything about WebSockets.
+      hijack = Hijack.new(conn, env)
+      env["rack.hijack"] = hijack.to_proc
+
+      status = headers = app_body = nil
+      error = nil
+
+      wall0 = Process.clock_gettime(MONOTONIC, :nanosecond)
+      cpu0 = Process.clock_gettime(THREAD_CPU, :nanosecond)
+
+
+      begin
+        status, headers, app_body = @app.call(env)
+      rescue Exception => e # rubocop:disable Lint/RescueException
+        error = e
+        log_error(e, env)
+        status, headers, app_body = lowlevel_error(e)
+      end
+
+      # A hijacked connection is no longer ours: no drain, no response, no
+      # RESP_END. The socket is the tunnel and the app has finished with it.
+      return :hijacked if hijack.hijacked?
+
+      # Consume anything the app left unread BEFORE writing the response. Go
+      # streams body frames concurrently with the app running, so an early
+      # return leaves frames in flight; sending RESP_END on top of them would
+      # desynchronise the connection and the next request would read a body
+      # chunk where it expected a REQUEST. Protocol invariant, not politeness.
+      input.drain if input.respond_to?(:drain)
+
+      service_ns = Process.clock_gettime(MONOTONIC, :nanosecond) - wall0
+      # The controller derives io_fraction = 1 - cpu_ns/service_ns from this pair.
+      # Two clock reads per request, and it is the cheapest honest signal for
+      # "would another thread actually help, or is the GVL the bottleneck".
+      cpu_ns = Process.clock_gettime(THREAD_CPU, :nanosecond) - cpu0
+
+      # Time blocked waiting on the uploader is not the application being slow.
+      # Reporting it separately lets the controller subtract it, so a slow client
+      # cannot masquerade as a workload that needs more capacity.
+      body_read_ns = input.respond_to?(:read_ns) ? input.read_ns : 0
+
+      # This PROCESS's cumulative GVL wait, not a per-request delta. Go
+      # differences it between observations and divides by available
+      # thread-time.
+      #
+      # A per-request delta would measure the wrong window: a worker thread does
+      # most of its waiting before the handler starts, while waking from the
+      # socket read and queueing for the GVL. Measured under pure CPU load the
+      # delta reported 0.03 where the truth was near 0.80.
+      gvl_wait_ns = GVL.wait_ns
+
+      write_response(conn, status, headers, app_body, service_ns, cpu_ns,
+                     body_read_ns, gvl_wait_ns)
+      :ok
+    ensure
+      # Ordering transcribed from puma/lib/puma/response.rb:109-131: close the
+      # body, then rack.after_reply in order, then rack.response_finished in
+      # reverse with (env, status, headers, error). Each callback is rescued
+      # individually so one bad hook cannot poison the rest or kill the thread.
+      app_body.close if app_body.respond_to?(:close)
+      run_hooks(env&.[]("rack.after_reply")) { |o| o.call }
+      run_hooks(env&.[]("rack.response_finished")&.reverse) do |o|
+        o.call(env, status, headers, error)
+      end
+    end
+
+    def write_response(conn, status, headers, body, service_ns, cpu_ns,
+                       body_read_ns, gvl_wait_ns = 0)
+      pairs = []
+      headers&.each do |key, value|
+        # Rack 3 permits array-valued headers, and a String value may still carry
+        # embedded newlines from Rack 2-era middleware. Both become repeated kv
+        # entries, which Go turns back into separate header lines -- this is what
+        # keeps multiple Set-Cookie headers intact.
+        if value.is_a?(Array)
+          value.each { |v| pairs << [key, v] }
+        elsif value.is_a?(String) && value.include?("\n")
+          value.split("\n").each { |v| pairs << [key, v] }
+        else
+          pairs << [key, value]
+        end
+      end
+
+      conn.write Wire.frame(Wire::RESPONSE, [status.to_i].pack("n") << Wire.encode_kv(pairs))
+      write_body(conn, body)
+
+      conn.write Wire.frame(Wire::RESP_END,
+        Wire.encode_resp_end_meta(
+          service_ns: service_ns,
+          cpu_ns: cpu_ns,
+          body_read_ns: body_read_ns,
+          gvl_wait_ns: gvl_wait_ns
+        ) << Wire.encode_kv([]))
+    end
+
+    # Emit the response body as RESP_CHUNK frames.
+    #
+    # Each yielded chunk becomes one frame, and Go flushes each frame straight to
+    # the client, so an SSE or long-poll body reaches the browser as the app
+    # produces it. This is the puma-dev bug we are deliberately not repeating:
+    # its reverse proxy sets FlushInterval to 1 second (puma-dev/dev/http.go:63),
+    # which batches event streams into one-second hiccups.
+    def write_body(conn, body)
+      return if body.nil?
+
+      # Rack 3 streaming body: an object responding to #call(stream) rather than
+      # #each. This is how Rack 3 apps do SSE without hijacking.
+      if !body.respond_to?(:each) && body.respond_to?(:call)
+        body.call(ChunkWriter.new(conn))
+        return
+      end
+
+      if body.respond_to?(:each)
+        body.each do |chunk|
+          next if chunk.nil? || chunk.empty?
+
+          conn.write Wire.frame(Wire::RESP_CHUNK, chunk)
+        end
+      else
+        conn.write Wire.frame(Wire::RESP_CHUNK, body.to_s)
+      end
+    end
+
+    # The `stream` object handed to a Rack 3 streaming body.
+    #
+    # The Rack SPEC requires read, write, <<, flush, close, close_read,
+    # close_write and closed?, and says their semantics "must be a best effort
+    # match to those of a normal Ruby IO or Socket object, using standard
+    # arguments and raising standard exceptions" (rack SPEC.rdoc:252). So this
+    # behaves like a Ruby IO opened for writing only -- including raising the
+    # same errors that one does, rather than inventing gentler ones.
+    #
+    # Rack::Lint enforces only that the methods exist, which is how an earlier
+    # version passed every other check while omitting three of them and failing
+    # the moment a streaming body was actually used.
+    class ChunkWriter
+      def initialize(conn)
+        @conn = conn
+        @closed = false
+      end
+
+      # IO#write takes any number of arguments and returns the total bytes.
+      def write(*chunks)
+        raise IOError, "closed stream" if @closed
+
+        total = 0
+        chunks.each do |chunk|
+          data = chunk.to_s
+          next if data.empty?
+
+          @conn.write Wire.frame(Wire::RESP_CHUNK, data)
+          total += data.bytesize
+        end
+        total
+      end
+
+      # IO#<< returns self so it can be chained, where #write returns a count.
+      # Aliasing the two would break `stream << a << b`.
+      def <<(chunk)
+        write(chunk)
+        self
+      end
+
+      # Nothing is buffered on this side: every write is already a frame on the
+      # wire, which is the point (puma-dev batches SSE into one-second hiccups
+      # with a 1s FlushInterval). IO#flush returns self.
+      def flush = self
+
+      # Write-only, so reading raises exactly what a Ruby IO opened "w" raises.
+      # Full-duplex streaming would mean interleaving BODY_CHUNK reads with
+      # RESP_CHUNK writes on one connection, which the protocol's drain
+      # invariant (consume the body to BODY_END before RESP_END) forbids.
+      def read(*)
+        raise IOError, "not opened for reading"
+      end
+
+      def close_read
+        raise IOError, "closing non-duplex IO for reading"
+      end
+
+      # Closing the stream ends the body but must NOT close the connection: the
+      # connection outlives the request and goes back to the idle pool. RESP_END
+      # is what terminates the body, and write_response sends it.
+      def close
+        @closed = true
+        nil
+      end
+      alias close_write close
+
+      def closed? = @closed
+    end
+
+    def run_hooks(hooks)
+      return unless hooks
+
+      hooks.each do |hook|
+        yield hook
+      rescue StandardError => e
+        @logger.puts "[puma-plus] hook error: #{e.class}: #{e.message}"
+      end
+    end
+
+    # The object behind env['rack.hijack'].
+    #
+    # Per the Rack SPEC, calling it returns the IO and also sets
+    # env['rack.hijack_io']. Sending the HIJACK frame first is what tells Go to
+    # stop expecting frames -- if the app wrote to the socket before that frame
+    # went out, Go would try to parse the app's bytes as a response frame.
+    class Hijack
+      def initialize(conn, env)
+        @conn = conn
+        @env = env
+        @hijacked = false
+      end
+
+      def hijacked? = @hijacked
+
+      def to_proc
+        method(:call).to_proc
+      end
+
+      def call
+        return @conn if @hijacked
+
+        @conn.write Wire.frame(Wire::HIJACK)
+        @hijacked = true
+        @env["rack.hijack_io"] = @conn
+      end
+    end
+
+    def lowlevel_error(_error)
+      [500, { "content-type" => "text/plain" }, ["Internal Server Error\n"]]
+    end
+
+    def log_error(error, env)
+      @logger.puts "[puma-plus] #{env["REQUEST_METHOD"]} #{env["PATH_INFO"]}: " \
+                   "#{error.class}: #{error.message}"
+      @logger.puts(error.backtrace&.first(10)&.map { |l| "    #{l}" }&.join("\n"))
+    end
+  end
+end

data/lib/puma_plus/ws.rb

@@ -0,0 +1,233 @@
+# frozen_string_literal: true
+
+require "socket"
+require "monitor"
+require "puma_plus/wire"
+
+module PumaPlus
+  # Outbound WebSocket API.
+  #
+  # The application never touches a socket. Go owns every WebSocket connection;
+  # this module sends it commands over a per-process control channel, which means
+  # the app can address a connection at any time from any thread -- not only
+  # while handling a message from it.
+  #
+  #   PumaPlus::WS.send(conn_id, "pong")
+  #   PumaPlus::WS.subscribe(conn_id, "room:42")
+  #   PumaPlus::WS.publish("room:42", payload)   # one call -> N sockets
+  #   PumaPlus::WS.close(conn_id, code: 1000)
+  #
+  # #publish is the reason this exists. Go does the fan-out on goroutines, so a
+  # broadcast to ten thousand subscribers costs the application exactly one frame
+  # instead of ten thousand. Doing that in Ruby is what currently pushes
+  # ActionCable into a separate process with Redis behind it.
+  module WS
+    class NotConnected < StandardError; end
+
+    TEXT = "text"
+    BINARY = "binary"
+
+    class << self
+      # Called once per worker process, before any threads start.
+      def connect!(socket_path, logger: $stderr)
+        @logger = logger
+        @socket_path = socket_path
+        @monitor = Monitor.new
+        @pending = {}
+        @seq = 0
+        @conn = UNIXSocket.new(socket_path)
+        @conn.sync = true
+        @conn.write Wire.frame(Wire::HELLO, Wire.encode_kv([
+          ["version", Wire::VERSION],
+          ["role", "ws"],
+          ["pid", Process.pid]
+        ]))
+        start_reader
+        true
+      rescue SystemCallError => e
+        @conn = nil
+        logger.puts "[puma-plus] ws control channel unavailable: #{e.class}: #{e.message}"
+        false
+      end
+
+      # Is the control channel usable?
+      def connected? = !@conn.nil?
+
+      # Send a reliable, ordered message to one connection.
+      #
+      # +stream:+ targets one stream of a WebTransport session. Without it the
+      # message goes to the connection's default channel, which is all a
+      # WebSocket has anyway -- so code that ignores streams works on both.
+      def send(conn_id, payload, opcode: TEXT, stream: nil)
+        args = { "conn_id" => conn_id, "opcode" => opcode }
+        args["stream_id"] = stream if stream
+        command(Wire::WS_SEND, args, payload)
+      end
+
+      # Send a best-effort datagram to one connection. WebTransport only.
+      #
+      # Deliberately a separate method from #send rather than a keyword on it.
+      # A datagram may be dropped, duplicated or delivered out of order; that is
+      # a fundamentally different contract, and a one-word argument is far too
+      # easy to copy between call sites without noticing the guarantee changed.
+      #
+      # Silently does nothing for a WebSocket client, which has no unreliable
+      # channel. Check PumaPlus::WSEvent#transport if you need to know.
+      def datagram(conn_id, payload)
+        command(Wire::WS_DATAGRAM, { "conn_id" => conn_id }, payload)
+      end
+
+      # Open a server-initiated stream on a WebTransport session.
+      # Returns the new stream id, or nil if the transport has no streams.
+      def open_stream(conn_id)
+        id = query_raw(Wire::WS_OPEN_STREAM, "conn_id" => conn_id)
+        id.nil? || id.empty? ? nil : id
+      end
+
+      # Broadcast to every subscriber of a topic.
+      #
+      # One frame regardless of how many subscribers there are; Go performs the
+      # writes. Returns nil rather than a delivery count because waiting for one
+      # would serialise the caller behind every recipient's socket -- use
+      # #subscribers if you need the number.
+      def publish(topic, payload, opcode: TEXT)
+        command(Wire::WS_PUBLISH, { "topic" => topic, "opcode" => opcode }, payload)
+      end
+
+      # Broadcast a best-effort datagram to a topic. WebTransport subscribers
+      # receive it; WebSocket subscribers are skipped, since they have no
+      # unreliable channel.
+      #
+      #   WS.publish("room:42", chat)            # reliable -- everyone
+      #   WS.publish_datagram("game:7", pos)     # lossy -- WebTransport only
+      #
+      # The asymmetry is the point: one of those is safe for chat, the other is
+      # not.
+      def publish_datagram(topic, payload)
+        command(Wire::WS_PUBLISH_DATAGRAM, { "topic" => topic }, payload)
+      end
+
+      # Which transport a connection is using: "websocket" or "webtransport".
+      def transport(conn_id)
+        query("transport", "conn_id" => conn_id)
+      end
+
+      def subscribe(conn_id, topic)
+        command(Wire::WS_SUBSCRIBE, { "conn_id" => conn_id, "topic" => topic })
+      end
+
+      def unsubscribe(conn_id, topic)
+        command(Wire::WS_UNSUBSCRIBE, { "conn_id" => conn_id, "topic" => topic })
+      end
+
+      def close(conn_id, code: 1000, reason: "")
+        command(Wire::WS_CLOSE,
+                { "conn_id" => conn_id, "code" => code.to_s, "reason" => reason })
+      end
+
+      # How many connections are subscribed to a topic. Synchronous.
+      def subscribers(topic)
+        query("subscribers", "topic" => topic).to_i
+      end
+
+      # Total live WebSocket connections across the whole server, not just this
+      # worker -- Go holds them all.
+      def connections
+        query("connections").to_i
+      end
+
+      # Is this connection still open?
+      def open?(conn_id)
+        query("exists", "conn_id" => conn_id) == "true"
+      end
+
+      private
+
+      def command(type, args, payload = nil)
+        raise NotConnected, "ws control channel is not connected" unless @conn
+
+        frame = Wire.frame(type, Wire.encode_kv(args.to_a) + (payload ? payload.to_s.b : ""))
+        # One writer at a time: frames must not interleave on the wire, and any
+        # thread may be publishing.
+        @monitor.synchronize { @conn.write(frame) }
+        nil
+      rescue IOError, SystemCallError => e
+        @conn = nil
+        raise NotConnected, "ws control channel lost: #{e.class}"
+      end
+
+      # Synchronous request/reply on a command frame that answers directly,
+      # rather than through the generic query path.
+      def query_raw(type, extra = {})
+        raise NotConnected, "ws control channel is not connected" unless @conn
+
+        seq = nil
+        q = Queue.new
+        @monitor.synchronize do
+          @seq += 1
+          seq = @seq.to_s
+          @pending[seq] = q
+          @conn.write Wire.frame(type, Wire.encode_kv([["seq", seq]] + extra.to_a))
+        end
+        begin
+          Timeout.timeout(5) { q.pop }
+        rescue Timeout::Error
+          @monitor.synchronize { @pending.delete(seq) }
+          nil
+        end
+      end
+
+      # Synchronous request/reply, correlated by sequence number so many threads
+      # can have questions outstanding on the one channel.
+      def query(kind, extra = {})
+        raise NotConnected, "ws control channel is not connected" unless @conn
+
+        seq = nil
+        q = Queue.new
+        @monitor.synchronize do
+          @seq += 1
+          seq = @seq.to_s
+          @pending[seq] = q
+          @conn.write Wire.frame(Wire::WS_QUERY,
+            Wire.encode_kv([["seq", seq], ["kind", kind]] + extra.to_a))
+        end
+
+        # Bounded wait: a lost reply must not park an application thread forever.
+        result = nil
+        begin
+          Timeout.timeout(5) { result = q.pop }
+        rescue Timeout::Error
+          @monitor.synchronize { @pending.delete(seq) }
+          raise NotConnected, "ws query #{kind} timed out"
+        end
+        result
+      end
+
+      def start_reader
+        Thread.new do
+          Thread.current.name = "puma-plus ws-control"
+          loop do
+            type, payload = Wire.read_frame(@conn)
+            next unless type == Wire::WS_QUERY_REPLY
+
+            pairs, = Wire.decode_kv(payload)
+            fields = pairs.to_h
+            seq = fields["seq"]
+            waiter = @monitor.synchronize { @pending.delete(seq) }
+            waiter&.push(fields["result"])
+          rescue Wire::Closed, IOError, SystemCallError
+            break
+          end
+          # Wake anything still waiting rather than leaving it parked forever.
+          @monitor.synchronize do
+            @pending.each_value { |w| w.push(nil) }
+            @pending.clear
+            @conn = nil
+          end
+        end
+      end
+    end
+  end
+end
+
+require "timeout"

data/lib/puma_plus/ws_event.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module PumaPlus
+  # Convenience wrapper over the WebSocket keys in a Rack env.
+  #
+  # A WebSocket event arrives as an ordinary Rack call with REQUEST_METHOD set to
+  # WEBSOCKET, so existing routing and middleware work unchanged. This just saves
+  # the app from string-keyed env lookups.
+  class WSEvent
+    def self.websocket?(env) = env["REQUEST_METHOD"] == "WEBSOCKET"
+
+    def initialize(env)
+      @env = env
+    end
+
+    # open | message | datagram | stream_open | stream_close | close
+    def kind      = @env["puma_plus.ws.event"]
+    def open?     = kind == "open"
+    def message?  = kind == "message"
+    def close?    = kind == "close"
+
+    # An unreliable datagram. A distinct event kind from #message? on purpose:
+    # this payload may have been dropped, duplicated or reordered, and handling
+    # it identically to a guaranteed message is usually a bug.
+    def datagram? = kind == "datagram"
+
+    def stream_open?  = kind == "stream_open"
+    def stream_close? = kind == "stream_close"
+
+    # "websocket" or "webtransport". Usually irrelevant -- that is the point of
+    # a shared hub -- but needed before reaching for a datagram.
+    def transport   = @env["puma_plus.ws.transport"]
+    def webtransport? = transport == "webtransport"
+
+    # Which stream a WebTransport message arrived on. nil for WebSocket, whose
+    # single channel needs no name.
+    def stream_id = @env["puma_plus.ws.stream_id"]
+
+    def conn_id = @env["puma_plus.ws.conn_id"]
+    def mode    = @env["puma_plus.ws.mode"]
+    def sticky? = mode == "sticky"
+    def binary? = @env["puma_plus.ws.opcode"] == "binary"
+
+    def close_code   = @env["puma_plus.ws.close_code"]&.to_i
+    def close_reason = @env["puma_plus.ws.close_reason"]
+
+    # The message payload.
+    def data
+      @data ||= (@env["rack.input"]&.read || "")
+    end
+
+    # Metadata captured at upgrade time and echoed on every event, so a flexible
+    # handler knows who it is talking to without a lookup.
+    def meta
+      @meta ||= @env.each_with_object({}) do |(k, v), h|
+        h[k.delete_prefix("puma_plus.ws.meta.")] = v if k.start_with?("puma_plus.ws.meta.")
+      end
+    end
+
+    # Queue time for this event, in seconds. WebSocket messages are queued and
+    # measured exactly like HTTP requests, which is what lets them drive
+    # autoscaling.
+    def queue_time = @env["puma.request_queue_time"]
+
+    # Build the 101 response that asks Go to take over a connection.
+    #
+    #   PumaPlus::WSEvent.upgrade(mode: :sticky, topics: ["room:42"])
+    def self.upgrade(mode: :flexible, topics: [], meta: {}, transport: nil)
+      headers = { "puma-plus-websocket" => mode.to_s }
+      headers["puma-plus-transport"] = transport.to_s if transport
+      headers["puma-plus-websocket-topic"] = Array(topics).map(&:to_s) unless Array(topics).empty?
+      unless meta.empty?
+        headers["puma-plus-websocket-meta"] = meta.map { |k, v| "#{k}=#{v}" }
+      end
+      [101, headers, []]
+    end
+  end
+end