puma-plus 0.1.0

data/lib/puma_plus/control_channel.rb

@@ -0,0 +1,138 @@
+# frozen_string_literal: true
+
+require "socket"
+require "etc"
+require "puma_plus/wire"
+
+module PumaPlus
+  # The control connection to the Go server, shared by every kind of worker
+  # supervisor.
+  #
+  # Extracted because Shepherd (which forks processes) and RactorWorker (which
+  # spawns Ractors) had byte-identical copies of connect, command loop,
+  # heartbeat and RSS reading. The only genuinely different parts are what
+  # SET_SLOTS means and what goes in the heartbeat -- which is exactly what the
+  # handler protocol below covers. Duplicating the rest meant a protocol change
+  # had to be made in two files, and nothing would have failed if it were made
+  # in one.
+  #
+  # A supervisor is the handler. It must respond to:
+  #
+  #   running?       -- keep looping?
+  #   set_slots(n)   -- SET_SLOTS arrived; n is the desired capacity
+  #   quiesce        -- stop replacing capacity that dies
+  #   shutdown(ms)   -- stop, with this much grace
+  #   heartbeat_kv   -- extra [key, value] pairs for WORKER_STATUS
+  #
+  # and may respond to:
+  #
+  #   on_idle        -- called each pass, whether or not a frame arrived, for
+  #                     work that cannot wait on the socket (the shepherd reaps
+  #                     dead children here)
+  class ControlChannel
+    # Long enough to be cheap, short enough that signal flags and child deaths
+    # are noticed promptly rather than only when Go happens to send something.
+    SELECT_TIMEOUT = 0.25
+
+    def initialize(socket_path:, logger: $stderr)
+      @socket_path = socket_path
+      @logger = logger
+      # Guards the socket. The heartbeat thread writes WORKER_STATUS while the
+      # command loop writes PONG, and without this the two can interleave
+      # mid-frame and hand Go a corrupt stream -- a race that existed in both
+      # copies of this code and had simply never been hit, since PING is only
+      # sent when a worker looks suspect.
+      @write_mu = Mutex.new
+    end
+
+    def connect!
+      @conn = UNIXSocket.new(@socket_path)
+      @conn.sync = true
+      write(Wire::HELLO, Wire.encode_kv([
+                                          ["version", Wire::VERSION],
+                                          ["role", "control"],
+                                          ["pid", Process.pid]
+                                        ]))
+      self
+    end
+
+    # Serve commands until the connection closes or the handler stops running.
+    def run(handler)
+      heartbeat = Thread.new { heartbeat_loop(handler) }
+      command_loop(handler)
+    ensure
+      heartbeat&.kill
+    end
+
+    def close
+      @conn&.close
+    rescue IOError
+      nil
+    end
+
+    # Advisory only. Go reads /proc itself for the memory guard rather than
+    # trusting a number reported by a process it may need to kill.
+    def self.rss_kb
+      File.read("/proc/self/statm").split[1].to_i * (Etc.sysconf(Etc::SC_PAGESIZE) / 1024)
+    rescue StandardError
+      0
+    end
+
+    private
+
+    def log(msg)
+      @logger.puts "[puma-plus control] #{msg}"
+      @logger.flush if @logger.respond_to?(:flush)
+    end
+
+    def write(type, payload = nil)
+      @write_mu.synchronize { @conn.write Wire.frame(type, payload) }
+    end
+
+    def command_loop(handler)
+      while handler.running?
+        handler.on_idle if handler.respond_to?(:on_idle)
+
+        ready = IO.select([@conn], nil, nil, SELECT_TIMEOUT)
+        next unless ready
+
+        begin
+          type, payload = Wire.read_frame(@conn)
+        rescue Wire::Closed
+          log "connection closed by server"
+          break
+        end
+
+        case type
+        when Wire::SET_SLOTS
+          handler.set_slots(payload.unpack1("N"))
+        when Wire::QUIESCE
+          handler.quiesce
+        when Wire::SHUTDOWN
+          handler.shutdown(payload.unpack1("N") || 5000)
+        when Wire::PING
+          write(Wire::PONG)
+        else
+          log "unexpected frame #{Wire.type_name(type)}"
+        end
+      end
+    end
+
+    def heartbeat_loop(handler)
+      seq = 0
+      while handler.running?
+        sleep 1
+        seq += 1
+        begin
+          write(Wire::WORKER_STATUS, Wire.encode_kv(
+                                       [["seq", seq], ["pid", Process.pid]] +
+                                       handler.heartbeat_kv +
+                                       [["rss_hint", self.class.rss_kb * 1024]]
+                                     ))
+        rescue IOError, Errno::EPIPE
+          break
+        end
+      end
+    end
+  end
+end

data/lib/puma_plus/gvl.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module PumaPlus
+  # GVL wait instrumentation.
+  #
+  # Measures how long threads spend runnable but unable to run because another
+  # thread in the same process holds the GVL. That number decides the hardest
+  # question the scaling controller asks: add a thread, or fork a process?
+  #
+  # Without it, the cheap heuristic is io_fraction = 1 - cpu/service, and it is
+  # not merely imprecise but backwards in the case that matters. A thread queued
+  # behind the GVL accrues no CPU time, so a purely CPU-bound request whose wall
+  # time is mostly queueing reads as "60% IO-bound" -- and the controller adds
+  # threads to a process where more threads make latency worse. Measured
+  # directly, the same request reports a 0.60 GVL fraction and the controller
+  # forks instead.
+  #
+  # The native extension is optional. Without it every method here returns 0 and
+  # the controller falls back to io_fraction plus gain probing, which is what it
+  # did before this existed.
+  module GVL
+    @native = begin
+      require "puma_plus_gvl"
+      true
+    rescue LoadError
+      false
+    end
+
+    # Fallbacks, defined only when the extension is absent so the real ones win.
+    unless @native
+      NATIVE = false
+
+      class << self
+        def native_start! = false
+        def native_stop! = false
+        def running? = false
+        def wait_ns = 0
+        def waits = 0
+        def thread_wait_ns = 0
+        def thread_waits = 0
+      end
+    end
+
+    class << self
+      # Is the native extension available?
+      def native? = const_defined?(:NATIVE) && const_get(:NATIVE)
+
+      # Begin measuring. Called once per worker process.
+      #
+      # Returns false when unavailable, which is not an error: the controller
+      # simply uses its cheaper signals.
+      def start!
+        return false unless native?
+
+        native_start!
+      end
+
+      def stop!
+        return false unless native?
+
+        native_stop!
+      end
+
+      # What fraction of an interval was spent waiting for the GVL.
+      #
+      # Above ~0.3 means this process is GVL-bound: its threads are queueing on
+      # each other rather than doing work, and adding more of them will deepen
+      # the queue rather than raise throughput.
+      def fraction(wait_ns, service_ns)
+        return 0.0 if service_ns.nil? || service_ns <= 0
+
+        wait_ns.to_f / service_ns
+      end
+    end
+  end
+end

data/lib/puma_plus/hooks.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+require "puma_plus/config_file"
+
+module PumaPlus
+  # Lifecycle hooks from the config file.
+  #
+  # These are the one part of the configuration that cannot be translated into
+  # flags for the Go server, because they are blocks. So the config file is read
+  # twice: once by the `puma-plus` launcher, which turns settings into flags and
+  # then execs Go, and again here by the Ruby process that actually runs the
+  # application, which is the only process where a hook can meaningfully run.
+  #
+  # Reading it twice means any top-level side effect in the config file happens
+  # twice. That is worth knowing and is why config files should declare rather
+  # than do -- the same advice puma's own docs give, for the same reason.
+  #
+  # Semantics follow puma exactly (verified against puma/lib/puma/cluster.rb:438
+  # and cluster/worker.rb:58,151):
+  #
+  #   before_fork          once in the shepherd, before ANY worker is forked
+  #   on_worker_boot(idx)  in each worker after fork, before it serves
+  #   on_worker_shutdown   in each worker, before it exits
+  #
+  # The timing of on_worker_boot is unusually clean here. A worker becomes
+  # visible to Go by dialing in, so running boot hooks before that dial means Go
+  # cannot dispatch a request to a worker whose hooks have not finished. In puma
+  # the equivalent guarantee needs the master to track a booted state; here it
+  # falls out of the connection being the readiness signal.
+  class Hooks
+    EMPTY = {}.freeze
+
+    def self.load(path, logger: $stderr)
+      return new(EMPTY, logger: logger) unless path && File.exist?(path)
+
+      # A throwaway Options: every non-hook directive in the file still runs and
+      # still sets values, and all of them are discarded. The launcher already
+      # turned those into flags, and honouring them here would let a worker
+      # disagree with the server about its own configuration.
+      dsl = ConfigFile.load(path, Options.new({}))
+      new(dsl.hooks, logger: logger, path: path)
+    rescue StandardError => e
+      # A config file that raises when read for hooks would otherwise take down
+      # a worker with a stack trace and no indication that hooks were the cause.
+      raise ConfigError, "loading hooks from #{path}: #{e.class}: #{e.message}"
+    end
+
+    def initialize(hooks, logger: $stderr, path: nil)
+      @hooks = hooks || EMPTY
+      @logger = logger
+      @path = path
+    end
+
+    attr_reader :path
+
+    def any?(name) = !(@hooks[name].nil? || @hooks[name].empty?)
+
+    def names = @hooks.keys.select { |k| any?(k) }
+
+    # Run every block registered for +name+, in declaration order.
+    #
+    # fatal: is the difference between "this process cannot do its job" and
+    # "this process is leaving anyway". A failed on_worker_boot means the worker
+    # would serve requests without whatever the hook was supposed to establish
+    # -- a database connection, most often -- so it dies loudly instead of
+    # serving wrongly. A failed shutdown hook is logged and the shutdown
+    # continues, because refusing to exit helps nobody.
+    def run(name, *args, fatal: true)
+      blocks = @hooks[name]
+      return if blocks.nil? || blocks.empty?
+
+      blocks.each_with_index do |blk, i|
+        blk.call(*args)
+      rescue StandardError, ScriptError => e
+        where = "#{name} hook #{i + 1}/#{blocks.size}#{@path ? " from #{@path}" : ''}"
+        if fatal
+          @logger.puts "[puma-plus] #{where} failed: #{e.class}: #{e.message}"
+          Array(e.backtrace).first(8).each { |l| @logger.puts "[puma-plus]   #{l}" }
+          raise HookError, "#{where} failed: #{e.class}: #{e.message}"
+        end
+
+        @logger.puts "[puma-plus] #{where} failed (continuing): #{e.class}: #{e.message}"
+      end
+    end
+  end
+
+  class HookError < StandardError; end
+end

data/lib/puma_plus/input.rb

@@ -0,0 +1,216 @@
+# frozen_string_literal: true
+
+require "stringio"
+require "puma_plus/wire"
+
+module PumaPlus
+  module Input
+    # An empty rack.input.
+    #
+    # Transcribed from puma/lib/puma/null_io.rb rather than depended upon: the
+    # class is a handful of methods, and the gem takes no dependencies.
+    class NullIO
+      def gets = nil
+      def each; end
+      def close; end
+      def size = 0
+      def eof? = true
+      def rewind = 0
+      def sync = true
+      def string = ""
+      def closed? = false
+      def read_ns = 0
+
+      def sync=(value)
+        value
+      end
+
+      def read(count = nil, buffer = nil)
+        result = count.nil? || count.zero? ? "" : nil
+        if buffer
+          buffer.clear
+          buffer << result if result
+          return count.nil? || count.zero? ? buffer : nil
+        end
+        result
+      end
+
+      def readpartial(_count, buffer = nil)
+        raise EOFError if buffer.nil?
+
+        buffer.clear
+        raise EOFError
+      end
+    end
+
+    # Frozen, which NullIO can afford because it holds no state at all. Two
+    # reasons, and the first stands on its own: a rack.input shared across every
+    # bodyless request must not be mutable, or one request could scribble on
+    # another's. The second is that a frozen stateless object is Ractor-shareable,
+    # and this constant is on the request path -- unfrozen, it raises
+    # IsolationError the first time a Ractor worker serves a GET.
+    EMPTY = NullIO.new.freeze
+
+    # A rack.input backed by BODY_CHUNK frames pulled from the worker connection
+    # on demand.
+    #
+    # Deliberately NOT rewindable. Rack 3 dropped the rewind requirement, and the
+    # payoff is that an arbitrarily large or chunked upload never touches a
+    # tempfile in either process. Compare puma, which decodes chunked bodies in
+    # Ruby and spools every one of them to a Tempfile
+    # (puma/lib/puma/client.rb:600), then rewrites CONTENT_LENGTH to the decoded
+    # size.
+    #
+    # Time spent blocked here is accumulated into #read_ns and reported in
+    # RESP_END, so the controller can subtract it from service time. Without
+    # that, a slow uploader would look exactly like an application that got
+    # slower, and the scaler would add capacity to fix a problem capacity cannot
+    # fix.
+    class FrameStream
+      # Sentinel returned by #read at EOF when a length was requested.
+      def initialize(conn, prefix = nil)
+        @conn = conn
+        @buf = +""
+        @buf << prefix if prefix && !prefix.empty?
+        @eof = false
+        @closed = false
+        @read_ns = 0
+      end
+
+      # Nanoseconds blocked reading body frames.
+      attr_reader :read_ns
+
+      def read(length = nil, outbuf = nil)
+        if length.nil?
+          fill_all
+          data = @buf
+          @buf = +""
+          outbuf ? outbuf.replace(data) : data
+        else
+          return outbuf ? outbuf.clear : +"" if length.zero?
+
+          fill_until(length)
+          if @buf.empty?
+            outbuf&.clear
+            return nil # EOF, per IO#read semantics
+          end
+          data = @buf.byteslice(0, length)
+          @buf = @buf.byteslice(length, @buf.bytesize - length) || +""
+          outbuf ? outbuf.replace(data) : data
+        end
+      end
+
+      def gets(sep = $INPUT_RECORD_SEPARATOR || "\n")
+        loop do
+          if (i = @buf.index(sep))
+            line = @buf.byteslice(0, i + sep.bytesize)
+            @buf = @buf.byteslice(i + sep.bytesize, @buf.bytesize - i - sep.bytesize) || +""
+            return line
+          end
+          break if @eof
+
+          pull or break
+        end
+
+        return nil if @buf.empty?
+
+        line = @buf
+        @buf = +""
+        line
+      end
+
+      def each
+        while (line = gets)
+          yield line
+        end
+      end
+
+      def eof?
+        @eof && @buf.empty?
+      end
+
+      def close
+        @closed = true
+      end
+
+      def closed? = @closed
+
+      # rewind is part of the Rack 2 SPEC but not Rack 3's. Raise rather than
+      # silently returning 0, so an app that depends on it fails loudly instead
+      # of quietly reading an empty body.
+      def rewind
+        raise Errno::ESPIPE, "puma-plus streams large request bodies; rack.input is not rewindable"
+      end
+
+      # Consume and discard anything the application did not read.
+      #
+      # Required before RESP_END: Go writes body frames concurrently with the
+      # application running, so if the app returns early (a 413, an auth
+      # rejection, a HEAD-like handler) the unread frames are still in flight.
+      # Leaving them would desynchronise the connection and the next request on
+      # it would read a body chunk where it expected a REQUEST.
+      def drain
+        pull until @eof
+        @buf = +""
+        nil
+      end
+
+      private
+
+      def fill_all
+        pull until @eof
+      end
+
+      def fill_until(length)
+        pull while @buf.bytesize < length && !@eof
+      end
+
+      # Read one more frame from the connection. Returns false at end of body.
+      def pull
+        return false if @eof
+
+        t0 = Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond)
+        type, payload = Wire.read_frame(@conn)
+        @read_ns += Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond) - t0
+
+        case type
+        when Wire::BODY_CHUNK
+          @buf << payload
+          true
+        when Wire::BODY_END
+          # Trailers ride on BODY_END. Nothing consumes them yet; decoding here
+          # keeps the frame accounted for and gives a place to surface them.
+          @trailers, = Wire.decode_kv(payload) unless payload.empty?
+          @eof = true
+          false
+        else
+          @eof = true
+          raise Wire::Truncated,
+                "expected BODY_CHUNK or BODY_END during body read, got #{Wire.type_name(type)}"
+        end
+      rescue Wire::Closed
+        @eof = true
+        false
+      end
+    end
+
+    # Build a rack.input for a decoded REQUEST.
+    #
+    # NONE   -> the shared frozen NullIO.
+    # INLINE -> a StringIO over a byteslice of the frame payload. Rewindable for
+    #           free, and covers essentially every form post and JSON API.
+    # STREAM -> FrameStream, pulling BODY_CHUNK frames on demand.
+    def self.for(mode, body, conn = nil)
+      case mode
+      when Wire::BODY_NONE   then EMPTY
+      when Wire::BODY_INLINE then StringIO.new(body)
+      when Wire::BODY_STREAM
+        raise ArgumentError, "STREAM body requires a connection" unless conn
+
+        FrameStream.new(conn, body)
+      else
+        raise ArgumentError, "unknown body_mode #{mode}"
+      end
+    end
+  end
+end