pgbus 0.5.1 → 0.6.1

data/lib/pgbus/web/streamer/instance.rb

@@ -0,0 +1,176 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    module Streamer
+      # Composes the Streamer's three background threads (Listener, Dispatcher,
+      # Heartbeat) with the shared Registry and dispatch_queue. One Instance
+      # per Puma worker. Owns the lifecycle of all three threads and the
+      # dedicated PG::Connection for LISTEN.
+      #
+      # Lifecycle:
+      #   Instance.new(...)  — allocates wiring, does NOT start threads
+      #   #start             — spawns listener/dispatcher/heartbeat in order
+      #   #register(conn)    — enqueues a ConnectMessage into the dispatch queue
+      #   #shutdown!         — sends shutdown sentinel to every connection and
+      #                        stops all threads in reverse order
+      #
+      # Dependency injection: every collaborator is constructor-injected so
+      # tests can swap in fakes without touching Pgbus.configuration. In
+      # production the module-level Streamer.current(...) builds all of the
+      # defaults from the configuration.
+      class Instance
+        attr_reader :registry, :listener, :dispatcher, :heartbeat, :dispatch_queue
+
+        def initialize(
+          client: Pgbus.client,
+          config: Pgbus.configuration,
+          pg_connection: nil,
+          logger: Pgbus.logger,
+          registry: nil,
+          dispatch_queue: nil
+        )
+          @client = client
+          @config = config
+          @logger = logger
+          @registry = registry || Registry.new
+          @dispatch_queue = dispatch_queue || Queue.new
+
+          @pg_connection = pg_connection || build_pg_connection
+          @listener = Listener.new(
+            pg_connection: @pg_connection,
+            dispatch_queue: @dispatch_queue,
+            health_check_ms: @config.streams_listen_health_check_ms,
+            logger: @logger
+          )
+          @dispatcher = Dispatcher.new(
+            client: @client,
+            registry: @registry,
+            listener: @listener,
+            dispatch_queue: @dispatch_queue,
+            logger: @logger,
+            config: @config
+          )
+          @heartbeat = Heartbeat.new(
+            registry: @registry,
+            dispatch_queue: @dispatch_queue,
+            interval: @config.streams_heartbeat_interval,
+            idle_timeout: @config.streams_idle_timeout,
+            logger: @logger
+          )
+          @started = false
+          @shutdown_mutex = Mutex.new
+        end
+
+        def start
+          return if @started
+
+          @started = true
+          @listener.start
+          @dispatcher.start
+          @heartbeat.start
+          self
+        end
+
+        # Enqueue a new SSE client. The dispatcher picks this up on its next
+        # iteration and runs the 5-step race-free replay sequence. The StreamApp
+        # calls this right after hijacking the socket.
+        #
+        # Guarded against the worker-shutdown race: if the request thread
+        # arrives here after `shutdown!` has flipped @started, we mark the
+        # connection dead and bail out instead of enqueueing a
+        # ConnectMessage. Otherwise the message would land on a dispatch
+        # queue that no one is draining, leaving the socket outside the
+        # registry and outside close_all_connections — the client would
+        # never see the pgbus:shutdown sentinel.
+        def register(connection)
+          return if connection.dead?
+
+          @shutdown_mutex.synchronize do
+            unless @started
+              connection.mark_dead!
+              return
+            end
+
+            @dispatch_queue << Dispatcher::ConnectMessage.new(connection: connection)
+          end
+        end
+
+        # Graceful shutdown for Puma worker restart. Order matters:
+        #   1. Heartbeat first (stop writing comments to connections we're
+        #      about to close)
+        #   2. Listener next (stop accepting new NOTIFYs)
+        #   3. Dispatcher next (drain the queue; it's now finite because
+        #      nothing else writes into it)
+        #   4. Send pgbus:shutdown sentinel to every connection and close
+        #      their sockets. We do this AFTER stopping the dispatcher so
+        #      no one else is writing to these IOs concurrently.
+        #
+        # Bounded by the configured write deadline per connection; a dead
+        # client drops instantly, a slow one stalls for at most write_deadline_ms.
+        def shutdown!
+          @shutdown_mutex.synchronize do
+            return unless @started
+
+            @started = false
+            safely { @heartbeat.stop }
+            safely { @listener.stop }
+            safely { @dispatcher.stop }
+            close_all_connections
+          end
+        end
+
+        private
+
+        def safely
+          yield
+        rescue StandardError => e
+          @logger.warn { "[Pgbus::Streamer::Instance] component stop raised: #{e.class}: #{e.message}" }
+        end
+
+        def close_all_connections
+          sentinel_bytes = Pgbus::Streams::Envelope.message(
+            id: 0,
+            event: "pgbus:shutdown",
+            data: '{"reason":"worker_restart"}'
+          )
+          deadline_ms = @config.streams_write_deadline_ms
+
+          @registry.each_connection do |connection|
+            # IoWriter holds the connection's mutex, so this write is
+            # serialised against any write the dispatcher/heartbeat
+            # might still be performing if their stop hadn't fully
+            # returned yet.
+            safely { IoWriter.write(connection, sentinel_bytes, deadline_ms: deadline_ms) }
+            safely { connection.close }
+          end
+        end
+
+        def build_pg_connection
+          require "pg" unless defined?(::PG::Connection)
+          opts = @config.connection_options
+          case opts
+          when String then ::PG.connect(opts)
+          when Hash   then ::PG.connect(**opts)
+          when Proc
+            # The Proc branch in connection_options typically returns
+            # ActiveRecord::Base.connection.raw_connection — a pooled
+            # AR connection. That's fatal for the streamer: LISTEN and
+            # wait_for_notify bind the session to the connection object,
+            # so if AR's pool hands the same raw_connection to another
+            # thread mid-wait, we get concurrent libpq calls → segfault
+            # or result corruption. Bail out with a clear error instead
+            # of silently segfaulting in production.
+            raise Pgbus::ConfigurationError,
+                  "Streamer cannot use a Proc-based connection_options. " \
+                  "Set Pgbus.configuration.database_url or connection_params " \
+                  "so the streamer owns its own dedicated PG::Connection for LISTEN/NOTIFY."
+          else
+            raise Pgbus::ConfigurationError,
+                  "Cannot build streamer PG connection from #{opts.class}"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pgbus/web/streamer/io_writer.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    module Streamer
+      # Non-blocking IO writer with a per-call deadline, serialised through the
+      # connection's own mutex. This is the bug-fix for puma/puma#576: a naive
+      # `io.write(bytes)` on a dead or slow SSE client deadlocks the dispatcher
+      # thread until the OS closes the socket (which can take minutes under a
+      # TCP keepalive). The message_bus gem hit this in production; we copy the
+      # pattern.
+      #
+      # The write loop uses write_nonblock + IO.select so a slow client at most
+      # stalls *its own* mutex-protected write for `deadline_ms`, never the
+      # dispatcher or heartbeat thread. When the deadline expires with bytes
+      # still pending, we return :blocked; the caller (Connection#enqueue or
+      # Connection#write_comment) translates that into mark_dead!, and the
+      # heartbeat sweep eventually unregisters the connection.
+      #
+      # Returns:
+      #   :ok       — all bytes written
+      #   :closed   — peer gone (EPIPE / ECONNRESET / IOError on closed IO)
+      #   :blocked  — deadline hit before all bytes could be written
+      module IoWriter
+        def self.write(connection, bytes, deadline_ms:)
+          connection.mutex.synchronize do
+            write_all(connection.io, bytes, deadline_ms)
+          end
+        end
+
+        def self.write_all(io, bytes, deadline_ms)
+          deadline = monotonic + (deadline_ms / 1000.0)
+          offset = 0
+          remaining = bytes.bytesize
+
+          while remaining.positive?
+            written = attempt_write(io, bytes, offset, remaining, deadline)
+            return written if written.is_a?(Symbol) # :blocked / :closed
+
+            offset += written
+            remaining -= written
+          end
+
+          :ok
+        end
+
+        def self.attempt_write(io, bytes, offset, remaining, deadline)
+          chunk = bytes.byteslice(offset, remaining)
+          io.write_nonblock(chunk)
+        rescue IO::WaitWritable
+          wait = deadline - monotonic
+          return :blocked if wait <= 0
+
+          # io.wait_writable is fiber-scheduler-friendly (Falcon v1.1) and
+          # functionally identical to IO.select on threaded Puma.
+          ready = io.wait_writable(wait)
+          return :blocked if ready.nil?
+
+          retry
+        rescue Errno::EPIPE, Errno::ECONNRESET, IOError
+          :closed
+        end
+
+        def self.monotonic
+          # Qualify ::Process because Pgbus::Process shadows it in this namespace.
+          ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+        end
+
+        private_class_method :write_all, :attempt_write, :monotonic
+      end
+    end
+  end
+end

data/lib/pgbus/web/streamer/listener.rb

@@ -0,0 +1,228 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    module Streamer
+      # Owns a single dedicated PG::Connection running LISTEN against every
+      # stream channel currently serving at least one SSE subscriber. On
+      # NOTIFY, posts a WakeMessage into the dispatch queue; the Dispatcher
+      # thread does the actual read_after + fanout.
+      #
+      # Threading:
+      #   - #start spawns ONE listener thread
+      #   - ensure_listening / remove_listening are called from the
+      #     dispatcher thread, which means the listener thread itself is
+      #     only running `wait_for_notify` — all LISTEN/UNLISTEN SQL goes
+      #     through a command queue that the listener thread drains between
+      #     notifies
+      #   - #stop joins the thread cleanly
+      #
+      # Health check: `wait_for_notify(timeout)` returns nil on timeout. When
+      # it does, the listener runs `SELECT 1` as a TCP keepalive. If that
+      # raises, the connection is reset (`conn.reset`) and every channel in
+      # `@listening_to` is re-LISTENed. This is the fix for design doc §11 #1
+      # (silently dropped LISTEN connections from NAT / PG restart / network
+      # blips).
+      #
+      # NOTIFY channel naming (from pgmq_v1.11.0.sql:1634):
+      #   PG_NOTIFY('pgmq.' || TG_TABLE_NAME || '.' || TG_OP, NULL)
+      # For a queue named `pgbus_stream_chat` the trigger table is
+      # `q_pgbus_stream_chat`, so the channel is `pgmq.q_pgbus_stream_chat.INSERT`.
+      class Listener
+        WakeMessage = Data.define(:queue_name)
+
+        CHANNEL_PREFIX = "pgmq.q_"
+        CHANNEL_SUFFIX = ".INSERT"
+
+        attr_reader :listening_to
+
+        def initialize(pg_connection:, dispatch_queue:, health_check_ms:, logger: Pgbus.logger)
+          @conn = pg_connection
+          @dispatch_queue = dispatch_queue
+          @health_check_ms = health_check_ms
+          @logger = logger
+          @listening_to = Set.new
+          @commands = Queue.new
+          @running = false
+          @thread = nil
+        end
+
+        def start
+          return if @running
+
+          @running = true
+          @thread = Thread.new { run_loop }
+          self
+        end
+
+        def stop
+          return unless @running
+
+          @running = false
+          @commands << [:stop]
+          # Interrupt the blocking wait_for_notify by closing the PG
+          # connection. Without this, the listener thread would sit
+          # inside wait_for_notify until a NOTIFY arrived, which may
+          # never happen. Closing the socket raises PG::Error inside
+          # wait_for_notify; our rescue clause sees @running == false
+          # on the next iteration and exits cleanly.
+          begin
+            @conn.close if @conn.respond_to?(:close)
+          rescue StandardError
+            # best effort — connection may already be gone
+          end
+          @thread&.join(5)
+          @thread = nil
+          self
+        end
+
+        # Synchronously enable LISTEN on a queue's NOTIFY channel.
+        # Blocks until the listener thread has actually executed the
+        # LISTEN SQL — this matters because Dispatcher#handle_connect
+        # needs the LISTEN to be active *before* it issues read_after,
+        # otherwise a broadcast committed in the gap would be neither
+        # in the read_after result nor delivered as a WakeMessage.
+        #
+        # The wait is bounded by the next wait_for_notify timeout
+        # (streams_listen_health_check_ms, default 250ms) plus a small
+        # margin so a stuck listener can't hang the dispatcher forever.
+        def ensure_listening(queue_name)
+          ack = Queue.new
+          @commands << [:listen, queue_name, ack]
+          ack.pop(timeout: ack_timeout)
+        end
+
+        # Asynchronous: lazy GC of LISTENs whose subscriber count
+        # has dropped to zero. No correctness path depends on this
+        # completing before the caller proceeds, so we don't block.
+        def remove_listening(queue_name)
+          @commands << [:unlisten, queue_name]
+        end
+
+        private
+
+        def run_loop
+          loop do
+            break unless @running
+
+            drain_commands
+            break unless @running
+
+            timeout_s = @health_check_ms / 1000.0
+            begin
+              @conn.wait_for_notify(timeout_s) do |channel, _pid, _payload|
+                handle_notify(channel)
+              end || run_health_check
+            rescue PG::Error => e
+              break unless @running
+
+              @logger.warn { "[Pgbus::Streamer::Listener] PG error (#{e.class}: #{e.message}) — reconnecting" }
+              reconnect!
+            end
+          end
+        ensure
+          safe_unlisten_all
+        end
+
+        def drain_commands
+          loop do
+            cmd = @commands.pop(true)
+            case cmd[0]
+            when :listen
+              ack = cmd[2]
+              begin
+                do_listen(cmd[1])
+              ensure
+                # Always ack so the caller is never left blocked, even
+                # if do_listen raised. The caller can still detect
+                # failure by checking @listening_to or by other means;
+                # we just promise not to deadlock the dispatcher.
+                ack&.push(:done)
+              end
+            when :unlisten then do_unlisten(cmd[1])
+            when :stop     then @running = false
+                                return
+            end
+          rescue ThreadError
+            # empty queue
+            return
+          end
+        end
+
+        # Caller wait budget for ensure_listening's ack. The listener
+        # thread will process the command at most one wait_for_notify
+        # cycle from now (bounded by health_check_ms); add a small
+        # safety margin so the dispatcher fails loud rather than
+        # hanging if the listener thread is dead.
+        def ack_timeout
+          (@health_check_ms / 1000.0) + 1.0
+        end
+
+        def do_listen(queue_name)
+          channel = channel_for(queue_name)
+          return if @listening_to.include?(channel)
+
+          @conn.exec(%(LISTEN "#{channel}"))
+          @listening_to.add(channel)
+        end
+
+        def do_unlisten(queue_name)
+          channel = channel_for(queue_name)
+          return unless @listening_to.include?(channel)
+
+          @conn.exec(%(UNLISTEN "#{channel}"))
+          @listening_to.delete(channel)
+        end
+
+        def handle_notify(channel)
+          queue_name = queue_name_from(channel)
+          return unless queue_name
+
+          @dispatch_queue << WakeMessage.new(queue_name: queue_name)
+        end
+
+        def run_health_check
+          @conn.exec("SELECT 1")
+        end
+
+        def reconnect!
+          @conn.reset
+          # Don't clear @listening_to until the new set is built. If a
+          # mid-loop LISTEN raises, we keep the original set so the
+          # next reconnect cycle still knows which channels need to
+          # come back. The previous version cleared first and lost
+          # any channels not yet retried on a transient error.
+          to_relisten = @listening_to.to_a
+          new_listening = Set.new
+          to_relisten.each do |channel|
+            @conn.exec(%(LISTEN "#{channel}"))
+            new_listening.add(channel)
+          end
+          @listening_to = new_listening
+        rescue PG::Error => e
+          @logger.error { "[Pgbus::Streamer::Listener] reconnect failed: #{e.class}: #{e.message}" }
+          sleep 0.5
+        end
+
+        def safe_unlisten_all
+          @listening_to.each do |channel|
+            @conn.exec(%(UNLISTEN "#{channel}"))
+          rescue PG::Error
+            # connection may be dead; nothing we can do
+          end
+          @listening_to.clear
+        end
+
+        def channel_for(queue_name)
+          "#{CHANNEL_PREFIX}#{queue_name}#{CHANNEL_SUFFIX}"
+        end
+
+        def queue_name_from(channel)
+          return nil unless channel.start_with?(CHANNEL_PREFIX) && channel.end_with?(CHANNEL_SUFFIX)
+
+          channel[CHANNEL_PREFIX.length..-(CHANNEL_SUFFIX.length + 1)]
+        end
+      end
+    end
+  end
+end

data/lib/pgbus/web/streamer/registry.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    module Streamer
+      # Worker-local in-memory registry of SSE connections indexed by stream
+      # name and connection id. Thread-safe via a single mutex. Reads return
+      # snapshots so iterators never hold the lock.
+      #
+      # Registry operations are O(1) under contention (mutex-protected hash
+      # lookups), and iteration is O(n) over a snapshot. The data structure is
+      # deliberately boring — the interesting parts of the streamer (LISTEN
+      # multiplexing, write scheduling, replay race handling) live elsewhere.
+      class Registry
+        def initialize
+          @mutex = Mutex.new
+          @by_stream = {}
+          @by_id = {}
+        end
+
+        def register(connection)
+          @mutex.synchronize do
+            existing = @by_id[connection.id]
+            return if existing.equal?(connection)
+
+            # Registration of a new object under an existing id is a
+            # defensive path — SecureRandom.hex(8) collisions are
+            # astronomical, but the Registry's invariant is "@by_stream
+            # only contains objects that are also in @by_id", so we
+            # must scrub the old entry from its stream index before
+            # overwriting. Otherwise connections_for(stream) would
+            # return orphaned objects and writes would go nowhere.
+            evict_from_stream(existing) if existing
+
+            @by_id[connection.id] = connection
+            (@by_stream[connection.stream_name] ||= Set.new).add(connection)
+          end
+        end
+
+        def unregister(connection)
+          @mutex.synchronize do
+            existing = @by_id.delete(connection.id)
+            return unless existing
+
+            set = @by_stream[existing.stream_name]
+            next unless set
+
+            set.delete(existing)
+            @by_stream.delete(existing.stream_name) if set.empty?
+          end
+        end
+
+        def lookup(id)
+          @mutex.synchronize { @by_id[id] }
+        end
+
+        # Returns a snapshot Array of connections on the given stream.
+        # Mutating the result has no effect on the registry.
+        def connections_for(stream_name)
+          @mutex.synchronize do
+            set = @by_stream[stream_name]
+            set ? set.to_a : []
+          end
+        end
+
+        # Snapshot list of stream names with at least one registered connection.
+        def streams
+          @mutex.synchronize { @by_stream.keys.dup }
+        end
+
+        def empty?(stream_name)
+          @mutex.synchronize do
+            set = @by_stream[stream_name]
+            set.nil? || set.empty?
+          end
+        end
+
+        def size
+          @mutex.synchronize { @by_id.size }
+        end
+
+        # Yields every registered connection across all streams. The iteration
+        # walks a snapshot so the block may safely call back into register/
+        # unregister without risk of deadlock or skipped items.
+        def each_connection(&)
+          snapshot = @mutex.synchronize { @by_id.values.dup }
+          snapshot.each(&)
+        end
+
+        private
+
+        # Must be called while holding @mutex.
+        def evict_from_stream(connection)
+          set = @by_stream[connection.stream_name]
+          return unless set
+
+          set.delete(connection)
+          @by_stream.delete(connection.stream_name) if set.empty?
+        end
+      end
+    end
+  end
+end

data/lib/pgbus/web/streamer.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    # The worker-local coordinator that owns SSE connections, one PG LISTEN
+    # session, and the dispatch/heartbeat threads. Lazily created on the
+    # first SSE connection to a Puma worker (or eagerly in tests). There is
+    # exactly one Instance per Puma worker process; the module-level
+    # accessors memoise it.
+    #
+    # This is NOT a Singleton in the GoF sense — tests are free to construct
+    # throwaway Instances directly and dependency-inject everything. The
+    # `current` / `reset!` helpers exist purely so the Rack StreamApp can
+    # share an instance across requests within a worker without passing it
+    # through every method call.
+    module Streamer
+      # Module-level mutex protecting `@current`. Without this, two
+      # concurrent first-callers inside a multi-threaded Puma worker
+      # can both build and start an Instance, leaking listener,
+      # dispatcher, and heartbeat threads plus a PG connection. Same
+      # gap would let `reset!` overlap teardown with a fresh
+      # replacement instance.
+      @current_mutex = Mutex.new
+
+      class << self
+        # Returns the worker-local instance, creating it on first call.
+        # `factory_opts` are passed to `Instance.new` the first time.
+        def current(**factory_opts)
+          @current_mutex.synchronize do
+            @current ||= Instance.new(**factory_opts).tap(&:start)
+          end
+        end
+
+        # Explicitly set the current instance — used by tests and by the
+        # Puma plugin to inject a pre-built instance.
+        def current=(instance)
+          @current_mutex.synchronize { @current = instance }
+        end
+
+        # Tear down the current instance and clear the slot. Called by the
+        # Puma shutdown hook (Phase 4.4) and by tests between examples.
+        def reset!
+          instance = nil
+          @current_mutex.synchronize do
+            instance = @current
+            @current = nil
+          end
+          instance&.shutdown!
+        end
+      end
+    end
+  end
+end

data/lib/pgbus.rb

@@ -39,6 +39,12 @@ module Pgbus
         )
         loader.ignore("#{__dir__}/generators")
         loader.ignore("#{__dir__}/active_job")
+        # lib/puma/plugin/pgbus_streams.rb is a Puma plugin — it's required
+        # explicitly by the user from config/puma.rb via `plugin :pgbus_streams`.
+        # Without this ignore, Zeitwerk scans lib/puma/ under the pgbus loader
+        # root and tries to autoload Puma::Plugin, which collides with the real
+        # Puma::Plugin class defined by the puma gem itself.
+        loader.ignore("#{__dir__}/puma")
         loader
       end
     end
@@ -78,10 +84,32 @@ module Pgbus
       @client ||= Client.new(configuration)
     end
 
+    # Entry point for the streams subsystem — `Pgbus.stream(name).broadcast(html)`
+    # or `Pgbus.stream(@order).current_msg_id`. Defined on Pgbus itself rather
+    # than inside lib/pgbus/streams.rb because that file is only Zeitwerk-loaded
+    # when Pgbus::Streams::Stream is first referenced — the chicken-and-egg
+    # problem means `Pgbus.stream(...)` would be undefined on the first call.
+    # Referencing Streams::Stream inside the method body forces Zeitwerk to
+    # load lib/pgbus/streams.rb lazily on first use, which is fine.
+    #
+    # Caches Stream instances by logical name so high-frequency callers
+    # (e.g. Turbo::StreamsChannel.broadcast_stream_to inside an
+    # after_update_commit callback firing 1000x/sec) don't allocate a new
+    # Stream + Mutex per broadcast. The cache is process-local; reset!
+    # clears it. The cache key is the resolved name string, not the raw
+    # streamables, so `Pgbus.stream(@order)` and `Pgbus.stream(@order)`
+    # in the same process return the same instance.
+    def stream(streamables)
+      name = Streams::Stream.name_from(streamables)
+      @stream_cache ||= Concurrent::Map.new
+      @stream_cache.compute_if_absent(name) { Streams::Stream.new(streamables) }
+    end
+
     def reset!
       @client&.close
       @client = nil
       @configuration = nil
+      @stream_cache = nil
     end
 
     # Discard the inherited PGMQ client after fork.