puma-plus 0.1.0

data/lib/puma_plus/shepherd.rb

@@ -0,0 +1,234 @@
+# frozen_string_literal: true
+
+require "socket"
+require "puma_plus/wire"
+require "puma_plus/worker"
+require "puma_plus/control_channel"
+require "puma_plus/hooks"
+
+module PumaPlus
+  # The shepherd process: loads the Rack app once, forks worker processes, and
+  # holds the control connection to the Go server.
+  #
+  # Why a shepherd rather than letting Go fork Ruby directly: fork must happen
+  # from a process that has already loaded the application, so children inherit a
+  # warm heap through copy-on-write. Go cannot do that. This mirrors puma's
+  # cluster master (puma/lib/puma/cluster.rb) but is far smaller, because Go owns
+  # the listener, the queue, and every scaling decision -- the shepherd only
+  # executes commands.
+  class Shepherd
+    RESPAWN_BACKOFF = 1.0
+
+    def initialize(socket_path:, app_path:, workers:, threads:, config_path: nil, logger: $stderr)
+      @socket_path = socket_path
+      @app_path = app_path
+      @config_path = config_path
+      @initial_workers = workers
+      @threads = threads
+      @logger = logger
+
+      @children = {}       # pid => worker_id
+      # Pids we killed on purpose. Without this, reap_children cannot tell a
+      # crash from a retirement and respawns the worker it was just told to
+      # retire -- which produced a kill/respawn loop that burned 25 forks in a
+      # scenario needing 3, and failed every request in flight on each dying
+      # worker.
+      @retiring = {}
+      @next_worker_id = 0
+      @running = true
+      @completed = 0
+    end
+
+    def run
+      # Preload before the first fork. This is the whole point of the shepherd:
+      # every child inherits an already-parsed application, so forks are cheap
+      # and copy-on-write keeps them cheap in memory too.
+      @app = Worker.new(
+        socket_path: @socket_path, app_path: @app_path,
+        threads: @threads, logger: @logger
+      ).send(:load_app)
+      log "app preloaded from #{@app_path}"
+
+      @hooks = Hooks.load(@config_path, logger: @logger)
+      log "hooks from #{@hooks.path}: #{@hooks.names.join(', ')}" if @hooks.names.any?
+
+      # Once, before ANY worker is forked -- puma's semantics
+      # (puma/lib/puma/cluster.rb:438). The documented use is closing connections
+      # opened during preload, so that each child dials its own rather than
+      # inheriting a socket that several processes then share.
+      @hooks.run(:before_fork)
+
+      @control = ControlChannel.new(socket_path: @socket_path, logger: @logger).connect!
+      install_signal_handlers
+
+      @initial_workers.times { spawn_worker }
+      @respawn_on_death = true
+
+      @control.run(self)
+    ensure
+      shutdown_children
+      @control&.close
+    end
+
+    # --- ControlChannel handler protocol ---
+
+    def running? = @running
+
+    # SET_SLOTS carries the desired *worker process* count.
+    def set_slots(target) = adjust_workers(target)
+
+    def quiesce
+      @respawn_on_death = false
+      log "quiescing: will not respawn"
+    end
+
+    def shutdown(grace_ms)
+      log "shutdown requested, grace #{grace_ms}ms"
+      @running = false
+    end
+
+    def heartbeat_kv
+      [["workers", active_count], ["worker_pids", @children.keys.join(",")]]
+    end
+
+    # Reaping cannot wait on the socket: a child can die at any moment, and the
+    # CHLD handler only sets a flag because a trap context cannot safely take
+    # locks or allocate and Process.wait can block.
+    def on_idle
+      return unless @child_died
+
+      @child_died = false
+      reap_children
+    end
+
+    private
+
+    def log(msg)
+      @logger.puts "[puma-plus shepherd] #{msg}"
+      @logger.flush if @logger.respond_to?(:flush)
+    end
+
+    def install_signal_handlers
+      # Reap in the main loop rather than in the handler: a trap context cannot
+      # safely take locks or allocate, and Process.wait can block.
+      trap("CHLD") { @child_died = true }
+      trap("TERM") { @running = false }
+      trap("INT") { @running = false }
+    end
+
+    # Fork one worker. The child never returns from here.
+    def spawn_worker
+      worker_id = (@next_worker_id += 1)
+
+      pid = fork do
+        # Reset inherited handlers so the child does not try to shepherd.
+        trap("CHLD", "DEFAULT")
+        trap("TERM") { exit 0 }
+        trap("INT") { exit 0 }
+        @control.close
+
+        Worker.new(
+          socket_path: @socket_path, app_path: @app_path,
+          threads: @threads, worker_id: worker_id, logger: @logger, hooks: @hooks,
+          multiprocess: true
+        ).run_preloaded(@app)
+        exit 0
+      end
+
+      @children[pid] = worker_id
+      log "forked worker #{worker_id} pid=#{pid} (#{@children.size} workers)"
+      worker_id
+    end
+
+    # Retire the newest worker. Go chooses when; the shepherd chooses which pid,
+    # because only it knows the mapping. Go's own GOAWAY handles draining the
+    # individual connections, so by the time this runs the worker is idle.
+    def retire_worker
+      # Skip anything already on its way out, or we would TERM the same pid
+      # repeatedly while waiting for it to die and never make progress.
+      candidates = @children.reject { |pid, _| @retiring.key?(pid) }
+      return if candidates.empty?
+
+      pid, worker_id = candidates.max_by { |_, id| id }
+      @retiring[pid] = worker_id
+      log "retiring worker #{worker_id} pid=#{pid}"
+      begin
+        Process.kill("TERM", pid)
+      rescue Errno::ESRCH
+        # Already gone; reaping will notice.
+      end
+    end
+
+    def reap_children
+      loop do
+        pid, status = Process.wait2(-1, Process::WNOHANG)
+        break unless pid
+
+        worker_id = @children.delete(pid)
+        next unless worker_id
+
+        # A worker we retired is supposed to be gone. Only an unexpected death
+        # earns a respawn.
+        if @retiring.delete(pid)
+          log "worker #{worker_id} pid=#{pid} retired"
+        elsif @running && @respawn_on_death
+          log "worker #{worker_id} pid=#{pid} died (#{status.exitstatus || status.termsig}); respawning"
+          sleep RESPAWN_BACKOFF
+          spawn_worker
+        else
+          log "worker #{worker_id} pid=#{pid} exited"
+        end
+      end
+    rescue Errno::ECHILD
+      nil
+    end
+
+    # Workers that exist and are not on their way out.
+    def active_count
+      @children.count { |pid, _| !@retiring.key?(pid) }
+    end
+
+    # SET_SLOTS carries the desired *worker process* count.
+    #
+    # Counts use active_count, which excludes workers already being retired.
+    # Counting them would make the loop TERM a fresh worker on every pass while
+    # waiting for the previous one to actually exit.
+    def adjust_workers(target)
+      target = 1 if target < 1
+
+      spawn_worker while active_count < target
+
+      while active_count > target
+        retire_worker
+        # Let the reaper run so the next iteration sees reality. Retirement is
+        # asynchronous: the worker finishes what it is doing before exiting.
+        sleep 0.05
+        reap_children
+      end
+    end
+
+    def shutdown_children
+      @respawn_on_death = false
+      @children.each_key do |pid|
+        Process.kill("TERM", pid)
+      rescue Errno::ESRCH
+        nil
+      end
+
+      deadline = Time.now + 5
+      until @children.empty? || Time.now > deadline
+        reap_children
+        sleep 0.05
+      end
+
+      @children.each_key do |pid|
+        Process.kill("KILL", pid)
+      rescue Errno::ESRCH
+        nil
+      end
+      log "all workers stopped"
+    end
+  end
+end
+
+require "etc"

data/lib/puma_plus/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module PumaPlus
+  VERSION = "0.1.0"
+end

data/lib/puma_plus/wire.rb

@@ -0,0 +1,236 @@
+# frozen_string_literal: true
+
+module PumaPlus
+  # The puma-plus frame protocol, Ruby side.
+  #
+  # Mirror of internal/wire/wire.go. See docs/PROTOCOL.md for the normative spec;
+  # internal/wire/golden_test.go enforces that the two agree byte for byte.
+  #
+  # Deliberately pure Ruby with no dependencies: String#pack/unpack and byteslice
+  # are enough for a length-prefixed framing over a flat string->string map, and
+  # every dependency avoided is one the gem does not have to install.
+  module Wire
+    # Frame types, data connections.
+    HELLO      = 0x01
+    REQUEST    = 0x02
+    BODY_CHUNK = 0x03
+    BODY_END   = 0x04
+    RESPONSE   = 0x05
+    RESP_CHUNK = 0x06
+    RESP_END   = 0x07
+    HIJACK     = 0x08
+    GOAWAY     = 0x09
+
+    # Frame types, WebSocket control connections (role "ws"). Outbound commands
+    # from Ruby, issuable from any thread at any time rather than only while a
+    # message is being handled.
+    WS_SEND        = 0x50
+    WS_PUBLISH     = 0x51
+    WS_SUBSCRIBE   = 0x52
+    WS_UNSUBSCRIBE = 0x53
+    WS_CLOSE       = 0x54
+    WS_QUERY       = 0x55
+    WS_QUERY_REPLY = 0x56
+
+    # WebTransport-only. Separate frames rather than a flag: best-effort
+    # delivery is a different contract from reliable delivery.
+    WS_DATAGRAM         = 0x57
+    WS_PUBLISH_DATAGRAM = 0x58
+    WS_OPEN_STREAM      = 0x59
+
+    # Frame types, control connections.
+    WORKER_STATUS = 0x40
+    SET_SLOTS     = 0x41
+    QUIESCE       = 0x42
+    SHUTDOWN      = 0x43
+    PING          = 0x44
+    PONG          = 0x45
+
+    # body_mode values in the REQUEST fixed meta.
+    BODY_NONE   = 0
+    BODY_INLINE = 1
+    BODY_STREAM = 2
+
+    # http_version values in the REQUEST fixed meta.
+    HTTP_10 = 10
+    HTTP_11 = 11
+    HTTP_2  = 20
+    HTTP_3  = 30
+
+    VERSION           = "1"
+    HEADER_SIZE       = 8
+    MAX_PAYLOAD       = 16 << 20
+    CHUNK_TARGET      = 64 << 10
+    INLINE_BODY_MAX   = 64 << 10
+    REQUEST_META_SIZE = 26
+    RESP_END_META_SIZE = 32
+
+    TYPE_NAMES = {
+      HELLO => "HELLO", REQUEST => "REQUEST", BODY_CHUNK => "BODY_CHUNK",
+      BODY_END => "BODY_END", RESPONSE => "RESPONSE", RESP_CHUNK => "RESP_CHUNK",
+      RESP_END => "RESP_END", HIJACK => "HIJACK", GOAWAY => "GOAWAY",
+      WORKER_STATUS => "WORKER_STATUS", SET_SLOTS => "SET_SLOTS",
+      QUIESCE => "QUIESCE", SHUTDOWN => "SHUTDOWN", PING => "PING", PONG => "PONG",
+      WS_SEND => "WS_SEND", WS_PUBLISH => "WS_PUBLISH", WS_SUBSCRIBE => "WS_SUBSCRIBE",
+      WS_UNSUBSCRIBE => "WS_UNSUBSCRIBE", WS_CLOSE => "WS_CLOSE",
+      WS_QUERY => "WS_QUERY", WS_QUERY_REPLY => "WS_QUERY_REPLY",
+      WS_DATAGRAM => "WS_DATAGRAM", WS_PUBLISH_DATAGRAM => "WS_PUBLISH_DATAGRAM",
+      WS_OPEN_STREAM => "WS_OPEN_STREAM"
+    }.freeze
+
+    BINARY = Encoding::BINARY
+
+    # Raised when a peer announces a frame larger than MAX_PAYLOAD. The only
+    # defense against a corrupt or hostile peer causing an OOM.
+    class PayloadTooLarge < StandardError; end
+
+    # Raised when a payload ends in the middle of a field.
+    class Truncated < StandardError; end
+
+    # Raised when the peer closes cleanly between frames. Not an error at the
+    # call site that owns the connection lifecycle -- it means "we're done".
+    class Closed < StandardError; end
+
+    module_function
+
+    def type_name(type)
+      TYPE_NAMES[type] || format("UNKNOWN(0x%02x)", type)
+    end
+
+    # Encode a frame: u8 type, u24 reserved (zero), u32 length, payload.
+    # Returns a single string so the caller emits one write per frame; frames
+    # must never appear interleaved on the wire.
+    def frame(type, payload = "")
+      [type, 0, 0, 0, payload.bytesize].pack("C4N") << payload.b
+    end
+
+    # Encode a kv blob: u32 count, then count x (u16 klen, k, u32 vlen, v).
+    #
+    # +pairs+ is an array of [key, value] rather than a Hash because duplicate
+    # keys are meaningful in header blobs, where they encode Rack 3 array-valued
+    # headers in wire order.
+    def encode_kv(pairs)
+      out = [pairs.size].pack("N")
+      pairs.each do |k, v|
+        k = k.to_s.b
+        v = v.to_s.b
+        out << [k.bytesize].pack("n") << k << [v.bytesize].pack("N") << v
+      end
+      out
+    end
+
+    # Decode a kv blob from the front of +buf+ starting at +offset+.
+    # Returns [pairs, next_offset]. Values are byteslices of buf, not copies --
+    # on the request hot path that aliasing is the point, since the whole env
+    # becomes byteslices of a single payload string.
+    def decode_kv(buf, offset = 0)
+      raise Truncated, "kv count" if buf.bytesize - offset < 4
+
+      count = buf.unpack1("N", offset: offset)
+      offset += 4
+
+      # A count larger than the remaining bytes can never be satisfied (the
+      # smallest possible pair is 6 bytes). Checking up front stops a corrupt
+      # length from driving a huge allocation.
+      raise Truncated, "kv count #{count} exceeds remaining bytes" if count > buf.bytesize - offset
+
+      pairs = Array.new(count)
+      count.times do |i|
+        raise Truncated, "kv key length" if buf.bytesize - offset < 2
+        kl = buf.unpack1("n", offset: offset)
+        offset += 2
+        raise Truncated, "kv key" if buf.bytesize - offset < kl
+        key = buf.byteslice(offset, kl)
+        offset += kl
+
+        raise Truncated, "kv value length" if buf.bytesize - offset < 4
+        vl = buf.unpack1("N", offset: offset)
+        offset += 4
+        raise Truncated, "kv value" if buf.bytesize - offset < vl
+        val = buf.byteslice(offset, vl)
+        offset += vl
+
+        pairs[i] = [key, val]
+      end
+      [pairs, offset]
+    end
+
+    # Encode the fixed 26-byte REQUEST meta.
+    def encode_request_meta(queue_ns:, runnable_wall_us:, body_wait_ns:,
+                            http_version:, body_mode:)
+      # Ruby has no native u64 pack directive that is endian-explicit in older
+      # syntax, so split each into two u32s. Q> exists and is used here.
+      [queue_ns, runnable_wall_us, body_wait_ns].pack("Q>3") <<
+        [http_version, body_mode].pack("C2")
+    end
+
+    # Decode the fixed 26-byte REQUEST meta from +buf+ at +offset+.
+    # Returns [meta_hash, next_offset].
+    def decode_request_meta(buf, offset = 0)
+      raise Truncated, "request meta" if buf.bytesize - offset < REQUEST_META_SIZE
+
+      queue_ns, runnable_wall_us, body_wait_ns = buf.unpack("Q>3", offset: offset)
+      http_version = buf.getbyte(offset + 24)
+      body_mode    = buf.getbyte(offset + 25)
+
+      [{
+        queue_ns: queue_ns,
+        runnable_wall_us: runnable_wall_us,
+        body_wait_ns: body_wait_ns,
+        http_version: http_version,
+        body_mode: body_mode
+      }, offset + REQUEST_META_SIZE]
+    end
+
+    # Encode the fixed 24-byte RESP_END meta. Piggybacking per-request timing on
+    # the response costs no extra frame and no extra syscall.
+    def encode_resp_end_meta(service_ns:, cpu_ns:, body_read_ns:, gvl_wait_ns: 0)
+      [service_ns, cpu_ns, body_read_ns, gvl_wait_ns].pack("Q>4")
+    end
+
+    # Decode the fixed 24-byte RESP_END meta. Returns [meta_hash, next_offset].
+    def decode_resp_end_meta(buf, offset = 0)
+      raise Truncated, "resp_end meta" if buf.bytesize - offset < RESP_END_META_SIZE
+
+      service_ns, cpu_ns, body_read_ns, gvl_wait_ns = buf.unpack("Q>4", offset: offset)
+      [{ service_ns: service_ns, cpu_ns: cpu_ns, body_read_ns: body_read_ns,
+         gvl_wait_ns: gvl_wait_ns },
+       offset + RESP_END_META_SIZE]
+    end
+
+    # Read one frame from +io+. Returns [type, payload].
+    #
+    # Raises Closed on a clean EOF between frames, Truncated on EOF mid-frame --
+    # the distinction matters because the first is normal shutdown and the second
+    # means the peer died holding our request.
+    def read_frame(io)
+      hdr = read_exactly(io, HEADER_SIZE)
+      raise Closed if hdr.nil?
+
+      type = hdr.getbyte(0)
+      # hdr[1..3] is the reserved u24. Per spec we ignore its value rather than
+      # rejecting non-zero, so it can later become a stream id.
+      len = hdr.unpack1("N", offset: 4)
+
+      if len > MAX_PAYLOAD
+        raise PayloadTooLarge, "#{type_name(type)} announced #{len} bytes"
+      end
+      return [type, ""] if len.zero?
+
+      payload = read_exactly(io, len)
+      raise Truncated, "#{type_name(type)} body" if payload.nil?
+
+      [type, payload]
+    end
+
+    # Read exactly +n+ bytes. Returns nil on clean EOF before any byte was read,
+    # which read_frame maps to Closed; a partial read raises Truncated.
+    def read_exactly(io, n)
+      buf = io.read(n)
+      return nil if buf.nil?
+      raise Truncated, "wanted #{n} bytes, got #{buf.bytesize}" if buf.bytesize < n
+
+      buf
+    end
+  end
+end

data/lib/puma_plus/worker.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+require "puma_plus/worker_thread"
+require "puma_plus/app_loader"
+require "puma_plus/hooks"
+require "puma_plus/ws"
+require "puma_plus/gvl"
+
+module PumaPlus
+  # A Ruby worker process: loads the Rack app once, then runs N threads that each
+  # dial the Go server and serve requests.
+  #
+  # Phase 1 runs a single worker process directly. Phase 4 puts a shepherd above
+  # this that forks several of these and handles SET_SLOTS/QUIESCE.
+  class Worker
+    def initialize(socket_path:, app_path:, threads:, worker_id: 0, config_path: nil,
+                   hooks: nil, multiprocess: false, logger: $stderr)
+      @socket_path = socket_path
+      @app_path = app_path
+      @threads = threads
+      @worker_id = worker_id
+      @logger = logger
+      # The shepherd passes hooks it already loaded, so a forked child does not
+      # re-read and re-evaluate the config file once per worker.
+      @hooks = hooks || Hooks.load(config_path, logger: logger)
+      # Matches puma: multithread when more than one thread can call the app in
+      # this process, multiprocess when a shepherd is managing worker processes
+      # (puma/lib/puma/binder.rb:33-34).
+      @multithread = threads > 1
+      @multiprocess = multiprocess
+    end
+
+    def run
+      # Load before any thread dials in. A connection appearing on the Go side is
+      # the readiness signal, so it must not appear until the app can actually
+      # serve -- that is the whole reason puma-plus needs no dial-polling.
+      run_preloaded(load_app)
+    end
+
+    # Serve using an already-loaded app. Used by the shepherd, which loads once
+    # before forking so children inherit a warm heap through copy-on-write.
+    def run_preloaded(app)
+      # Start GVL instrumentation before any thread serves, so the very first
+      # request is measured. A no-op without the native extension.
+      PumaPlus::GVL.start!
+
+      # One outbound WebSocket control channel per worker process, opened before
+      # any thread serves. The app can then address any connection at any time,
+      # from any thread, rather than only while handling a message.
+      PumaPlus::WS.connect!(@socket_path, logger: @logger)
+
+      # Before any thread dials in, which is what makes this ordering meaningful:
+      # a connection appearing on the Go side is the readiness signal, so hooks
+      # finishing first means Go cannot dispatch a request to a worker whose
+      # setup has not completed.
+      @hooks.run(:on_worker_boot, @worker_id)
+
+      threads = @threads.times.map do |i|
+        Thread.new do
+          Thread.current.name = "puma-plus #{@worker_id}/#{i}"
+          WorkerThread.new(
+            socket_path: @socket_path,
+            app: app,
+            worker_id: @worker_id,
+            thread_index: i,
+            multithread: @multithread,
+            multiprocess: @multiprocess,
+            logger: @logger
+          ).run
+        end
+      end
+
+      @logger.puts "[puma-plus] worker #{@worker_id} pid=#{Process.pid} " \
+                   "serving with #{@threads} threads"
+
+      trap("TERM") { threads.each(&:kill) }
+      trap("INT") { threads.each(&:kill) }
+
+      threads.each(&:join)
+    ensure
+      # Not fatal: the process is leaving either way, and refusing to exit
+      # because a cleanup hook raised helps nobody.
+      @hooks.run(:on_worker_shutdown, @worker_id, fatal: false)
+    end
+
+    private
+
+    def load_app = AppLoader.load(@app_path)
+  end
+end