pgbus 0.17.2 → 0.18.0

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

@@ -14,10 +14,15 @@ module Pgbus
       # the client-side leg of the replay-race fix (§6.5 of the design doc).
       class Connection
         attr_reader :id, :stream_name, :io, :mutex, :context
+        # The client's per-tab presence seat (?seat=), or nil (issue #515).
+        attr_reader :seat
         # The presence member id this connection auto-joined as, or nil for
         # non-presence streams / anonymous connections. Set by the
         # Dispatcher on connect; read on disconnect and heartbeat touch.
         attr_accessor :presence_member
+        # True when presence_member joined as a seat row, false when as a
+        # member-keyed row (table not yet migrated). Leave follows the join.
+        attr_accessor :presence_seat
 
         # last_msg_id_sent is the ONE field whose writer can cross threads: with
         # streams_writer_threads > 0 (issue #321) the OutboundPump worker thread
@@ -31,7 +36,7 @@ module Pgbus
           @last_msg_id_sent.get
         end
 
-        def initialize(id:, stream_name:, io:, since_id:, writer:, write_deadline_ms:, context: nil)
+        def initialize(id:, stream_name:, io:, since_id:, writer:, write_deadline_ms:, context: nil, seat: nil)
           @id = id
           @stream_name = stream_name
           @io = io
@@ -50,6 +55,7 @@ module Pgbus
           # registry when evaluating visible_to predicates. Defaults to
           # nil for tests that don't need audience filtering.
           @context = context
+          @seat = seat
         end
 
         # deadline_ms defaults to the connection's own write deadline so every
@@ -82,20 +88,21 @@ module Pgbus
         end
 
         def write_comment(text)
-          bytes = Pgbus::Streams::Envelope.comment(text)
-          result = @writer.write(self, bytes, deadline_ms: @write_deadline_ms)
-          if result == :ok
-            @last_write_at = monotonic
-          else
-            mark_dead!
-          end
-          result
+          write_frame(Pgbus::Streams::Envelope.comment(text))
         end
 
         def write_sentinel(bytes)
           @writer.write(self, bytes, deadline_ms: @write_deadline_ms)
         end
 
+        # The pgbus:connected frame carrying this connection's id. Written by
+        # the Dispatcher once the connection is registered and its presence
+        # seat joined (issue #521), so a client that acts on the id right
+        # after the event finds its row.
+        def write_connected
+          write_frame(Pgbus::Streams::Envelope.connected(id: @id))
+        end
+
         def idle_for
           monotonic - @last_write_at
         end
@@ -136,6 +143,16 @@ module Pgbus
 
         private
 
+        def write_frame(bytes)
+          result = @writer.write(self, bytes, deadline_ms: @write_deadline_ms)
+          if result == :ok
+            @last_write_at = monotonic
+          else
+            mark_dead!
+          end
+          result
+        end
+
         # The SSE event name for a frame: the envelope's typed event when
         # present, else the default turbo-stream. Plain ReadAfter::Envelopes
         # (no event field) and StreamEnvelopes with a nil event both fall

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

@@ -13,7 +13,12 @@ module Pgbus
       # is fiber-safe under Falcon's scheduler.
       class FalconConnection
         attr_reader :id, :stream_name, :io, :mutex, :context
+        # The client's per-tab presence seat (?seat=), or nil (issue #515).
+        attr_reader :seat
         attr_accessor :presence_member
+        # True when presence_member joined as a seat row, false when as a
+        # member-keyed row (table not yet migrated). Leave follows the join.
+        attr_accessor :presence_seat
 
         # last_msg_id_sent mirrors Connection's cross-thread cursor (issue #321):
         # a Concurrent::AtomicReference so the OutboundPump worker's advance and
@@ -23,7 +28,7 @@ module Pgbus
           @last_msg_id_sent.get
         end
 
-        def initialize(id:, stream_name:, body:, since_id:, write_deadline_ms:, context: nil)
+        def initialize(id:, stream_name:, body:, since_id:, write_deadline_ms:, context: nil, seat: nil)
           @id = id
           @stream_name = stream_name
           @body = body
@@ -37,6 +42,7 @@ module Pgbus
           @created_at = monotonic
           @last_write_at = @created_at
           @context = context
+          @seat = seat
         end
 
         # deadline_ms is accepted for duck-type parity with Connection#enqueue
@@ -68,14 +74,12 @@ module Pgbus
         end
 
         def write_comment(text)
-          bytes = Pgbus::Streams::Envelope.comment(text)
-          result = write_to_body(bytes)
-          if result == :ok
-            @last_write_at = monotonic
-          else
-            mark_dead!
-          end
-          result
+          write_frame(Pgbus::Streams::Envelope.comment(text))
+        end
+
+        # See Connection#write_connected (issue #521).
+        def write_connected
+          write_frame(Pgbus::Streams::Envelope.connected(id: @id))
         end
 
         def write_sentinel(bytes)
@@ -114,6 +118,16 @@ module Pgbus
           event && !event.to_s.empty? ? event.to_s : Pgbus::Streams::DEFAULT_SSE_EVENT
         end
 
+        def write_frame(bytes)
+          result = write_to_body(bytes)
+          if result == :ok
+            @last_write_at = monotonic
+          else
+            mark_dead!
+          end
+          result
+        end
+
         def write_to_body(bytes)
           @mutex.synchronize do
             return :closed if @dead || @body.closed?

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

@@ -26,13 +26,18 @@ module Pgbus
       # the Connection's own mutex, so concurrent dispatcher + heartbeat
       # writes are safe.
       class Heartbeat
-        def initialize(registry:, dispatch_queue:, interval:, idle_timeout:, logger: Pgbus.logger, clock: nil)
+        # presence_expiry: true (streams_presence_patterns is set) posts the
+        # presence touch every tick, even with no presence connection here,
+        # so the dispatcher can expire seats left by other processes.
+        def initialize(registry:, dispatch_queue:, interval:, idle_timeout:, logger: Pgbus.logger, clock: nil,
+                       presence_expiry: false)
           @registry = registry
           @queue = dispatch_queue
           @interval = interval
           @idle_timeout = idle_timeout
           @logger = logger
           @clock = clock || -> { ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) }
+          @presence_expiry = presence_expiry
           @running = false
           @thread = nil
           @wake = ConditionVariable.new
@@ -122,9 +127,10 @@ module Pgbus
         # Posts one batched touch per tick so the dispatcher (which owns AR
         # connection release) refreshes last_seen_at for live presence
         # members, keeping them out of the sweeper. No message when no
-        # connection has presence — avoids waking the dispatcher for nothing.
+        # connection has presence and seat expiry is off — avoids waking the
+        # dispatcher for nothing.
         def enqueue_presence_touch(connections)
-          return if connections.empty?
+          return if connections.empty? && !@presence_expiry
 
           @queue << StreamEventDispatcher::PresenceTouchMessage.new(connections: connections)
         end

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

@@ -3,14 +3,15 @@
 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.
+      # Composes the Streamer's background threads (Listener, Dispatcher,
+      # Heartbeat, and the PeerCloseWatcher unless streams_peer_close_interval
+      # is nil) with the shared Registry and dispatch_queue. One Instance per
+      # Puma worker. Owns the lifecycle of those threads and the dedicated
+      # PG::Connection for LISTEN.
       #
       # Lifecycle:
       #   Instance.new(...)  — allocates wiring, does NOT start threads
-      #   #start             — spawns listener/dispatcher/heartbeat in order
+      #   #start             — spawns listener/dispatcher/heartbeat/watcher 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
@@ -21,7 +22,7 @@ module Pgbus
       # defaults from the configuration.
       class Instance
         attr_reader :registry, :listener, :dispatcher, :heartbeat, :dispatch_queue, :stream_counter, :pump,
-                    :autoscaler
+                    :autoscaler, :peer_close_watcher
 
         def initialize(
           client: Pgbus.client,
@@ -74,8 +75,10 @@ module Pgbus
             dispatch_queue: @dispatch_queue,
             interval: @config.streams_heartbeat_interval,
             idle_timeout: @config.streams_idle_timeout,
-            logger: @logger
+            logger: @logger,
+            presence_expiry: Array(@config.streams_presence_patterns).any?
           )
+          @peer_close_watcher = build_peer_close_watcher
           @started = false
           @shutdown_mutex = Mutex.new
         end
@@ -90,6 +93,7 @@ module Pgbus
           @listener.start
           @dispatcher.start
           @heartbeat.start
+          @peer_close_watcher&.start
           self
         end
 
@@ -118,7 +122,9 @@ module Pgbus
         end
 
         # Graceful shutdown for Puma worker restart. Order matters:
-        #   1. Heartbeat first (stop writing comments to connections we're
+        #   0. Peer-close watcher first (issue #519): stop reaping sockets
+        #      we're about to close ourselves; #stop wakes its select at once
+        #   1. Heartbeat next (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
@@ -149,6 +155,7 @@ module Pgbus
 
             @started = false
             threads = component_threads
+            safely { @peer_close_watcher&.stop }
             safely { @heartbeat.stop }
             safely { @listener.stop }
             safely { @dispatcher.stop }
@@ -241,6 +248,7 @@ module Pgbus
         def component_threads
           {
             "heartbeat" => @heartbeat.threads,
+            "peer_close_watcher" => @peer_close_watcher ? @peer_close_watcher.threads : [],
             "listener" => @listener.threads,
             "dispatcher" => @dispatcher.threads,
             "pump" => @pump ? @pump.threads : []
@@ -275,6 +283,15 @@ module Pgbus
           )
         end
 
+        # Notices a closed SSE client at once instead of on the next
+        # heartbeat write (issue #519). nil interval: heartbeat only.
+        def build_peer_close_watcher
+          interval = @config.streams_peer_close_interval
+          return nil unless interval
+
+          PeerCloseWatcher.new(registry: @registry, dispatch_queue: @dispatch_queue, interval: interval, logger: @logger)
+        end
+
         def build_pump
           threads = @config.streams_writer_threads
           return nil unless threads.positive?

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

@@ -0,0 +1,183 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    module Streamer
+      # Notices an SSE client that went away the moment its socket says so,
+      # instead of on the next heartbeat write (issue #519).
+      #
+      # The server never reads a hijacked SSE socket after the request, so a
+      # closed tab used to stay registered (and present) until a heartbeat
+      # write failed, up to streams_heartbeat_interval later. Often longer:
+      # the first write after the client's FIN still succeeds. The client
+      # sends nothing after its request, so the socket only turns readable
+      # when the peer closes or resets it. This thread waits on every live
+      # socket with IO.select and probes the readable ones with a
+      # non-blocking read. EOF or a reset means the client is gone: the
+      # connection is marked dead, closed, and handed to the dispatcher's
+      # disconnect path, the same DisconnectMessage the heartbeat posts. That
+      # path is idempotent, so a heartbeat racing the watcher still leaves
+      # the seat once.
+      #
+      # Only connections with a real socket are watched: Puma's and Falcon's
+      # rack.hijack paths. Falcon's experimental streaming-body path
+      # (FalconConnection) exposes no socket and keeps the heartbeat.
+      #
+      # Half-close: a client that shuts down only its write side (SHUT_WR)
+      # reads as EOF and is reaped. EventSource, fetch streams and reverse
+      # proxies never half-close an SSE request, so EOF means the client left.
+      #
+      # The probe takes the connection's write mutex with try_lock: a TLS
+      # wrapper (Puma's MiniSSL) is not safe to read while IoWriter writes
+      # through it. A socket mid-write is skipped and probed on a later pass.
+      #
+      # `interval` bounds how long a new connection waits to join the
+      # watched set; a close on a watched socket wakes the select at once,
+      # and #stop wakes it through a self-pipe.
+      class PeerCloseWatcher
+        # Stray bytes from a client are read and discarded in chunks of this.
+        DRAIN_BYTES = 4096
+
+        # How long #stop waits for the thread. It is woken at once, so this
+        # only bounds a reap stuck closing a socket (see Heartbeat#stop).
+        STOP_JOIN_SECONDS = 5
+
+        def initialize(registry:, dispatch_queue:, interval: 1.0, logger: Pgbus.logger)
+          @registry = registry
+          @queue = dispatch_queue
+          @interval = interval
+          @logger = logger
+          @running = false
+          @thread = nil
+          @wake_reader, @wake_writer = IO.pipe
+        end
+
+        def start
+          return if @running
+
+          @wake_reader, @wake_writer = IO.pipe if @wake_reader.closed?
+          @running = true
+          @thread = Thread.new { run_loop }
+          self
+        end
+
+        # See Heartbeat#threads (issue #443).
+        def threads
+          [@thread].compact
+        end
+
+        def stop
+          return unless @running
+
+          @running = false
+          wake
+          @thread&.join(STOP_JOIN_SECONDS)
+          @thread = nil
+          @wake_reader.close
+          @wake_writer.close
+          self
+        end
+
+        # One wait-and-probe pass, at most `interval` long. Public for tests.
+        def tick
+          by_socket = watched_sockets
+          ready, = IO.select([@wake_reader, *by_socket.keys], nil, nil, @interval)
+          Array(ready).each do |socket|
+            socket.equal?(@wake_reader) ? drain_wake : probe(by_socket[socket])
+          end
+        rescue IOError, Errno::EBADF => e
+          # A socket was closed under the select (shutdown, a dispatcher
+          # write that failed). The next pass rebuilds the set without it.
+          @logger.debug { "[Pgbus::Streamer::PeerCloseWatcher] select interrupted: #{e.class}: #{e.message}" }
+        end
+
+        private
+
+        def wake
+          @wake_writer.write_nonblock(".", exception: false)
+        rescue IOError
+          nil
+        end
+
+        def drain_wake
+          @wake_reader.read_nonblock(64, exception: false)
+        end
+
+        # One socket's probe never ends the pass for the others.
+        def probe(connection)
+          return unless connection.mutex.try_lock
+
+          closed = begin
+            peer_closed?(connection)
+          ensure
+            connection.mutex.unlock
+          end
+          reap(connection) if closed
+        rescue StandardError => e
+          @logger.error { "[Pgbus::Streamer::PeerCloseWatcher] probe of #{connection.id} raised: #{e.class}: #{e.message}" }
+        end
+
+        def run_loop
+          while @running
+            begin
+              tick
+            rescue StandardError => e
+              @logger.error { "[Pgbus::Streamer::PeerCloseWatcher] tick raised: #{e.class}: #{e.message}" }
+              sleep @interval
+            end
+          end
+        end
+
+        # select needs a real IO: Puma's TLS socket and Falcon's hijacked
+        # stream answer #to_io with the underlying socket.
+        def watched_sockets
+          sockets = {}
+          @registry.each_connection do |connection|
+            next if connection.dead?
+
+            socket = socket_of(connection)
+            sockets[socket] = connection if socket
+          end
+          sockets
+        end
+
+        def socket_of(connection)
+          io = connection.io
+          return nil unless io.respond_to?(:to_io) && io.respond_to?(:read_nonblock)
+
+          socket = io.to_io
+          socket if socket.is_a?(::BasicSocket) && !socket.closed?
+        end
+
+        def peer_closed?(connection)
+          # A nil read is EOF; :wait_readable (TLS, a spurious wake) and stray
+          # bytes leave the connection alone.
+          connection.io.read_nonblock(DRAIN_BYTES, exception: false).nil?
+        rescue IO::WaitReadable
+          # A socket wrapper that raises instead of honouring exception:
+          # false (EAGAIN is a SystemCallError, so this rescue comes first).
+          false
+        rescue IOError, SystemCallError, *tls_errors
+          true
+        end
+
+        # A TLS socket whose client vanished without close_notify raises its
+        # library's SSL error ("unexpected eof while reading") on read.
+        def tls_errors
+          errors = []
+          errors << ::OpenSSL::SSL::SSLError if defined?(::OpenSSL::SSL::SSLError)
+          errors << ::Puma::MiniSSL::SSLError if defined?(::Puma::MiniSSL::SSLError)
+          errors
+        end
+
+        def reap(connection)
+          return if connection.dead?
+
+          connection.mark_dead!
+          connection.close
+          @queue << StreamEventDispatcher::DisconnectMessage.new(connection: connection)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    module Streamer
+      # Connection-driven presence (issues #169, #515), run on the
+      # StreamEventDispatcher thread, which owns AR connection release.
+      #
+      # On a seat table (Pgbus::Streams::Presence.seats_supported?) every SSE
+      # connection is its own row, keyed by connection id and carrying the
+      # client's seat; two tabs of one member are two rows. On an older
+      # table it falls back to one member-keyed row, as before.
+      #
+      # The heartbeat posts a PresenceTouchMessage each tick; #touch refreshes
+      # the live seats, re-joins any whose row was expired while the
+      # connection stayed open, and expires silent seats across all streams
+      # (rows of crashed processes included) at most every min_ttl / 2 after
+      # a successful sweep.
+      #
+      # Each connection remembers whether it joined as a seat or as a
+      # member-keyed row, so a table migrated while it is open still has
+      # its row removed on leave (and converted to a seat on the next touch).
+      #
+      # Every failure is logged and swallowed: a presence DB hiccup must not
+      # knock a live SSE connection out of the registry.
+      class PresenceTracker
+        def initialize(config:, presence_provider:, logger:, sweeper: nil, clock: nil)
+          @config = config
+          @presence_provider = presence_provider
+          @logger = logger
+          @sweeper = sweeper || method(:expire_seats)
+          @clock = clock || -> { ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) }
+          @last_sweep_at = nil
+        end
+
+        def join(connection, stream = connection.stream_name)
+          return unless @config&.stream_presence?(stream)
+
+          member = @config.presence_member_for(connection.context)
+          return unless member
+
+          presence = @presence_provider.call(stream)
+          seat = seats?(presence)
+          join_member(presence, connection, member, seat)
+          connection.presence_member = member[:id] if connection.respond_to?(:presence_member=)
+          connection.presence_seat = seat if connection.respond_to?(:presence_seat=)
+        rescue StandardError => e
+          log_failure("join", e)
+        end
+
+        def leave(connection, stream = connection.stream_name)
+          member_id = presence_member_of(connection)
+          return unless member_id
+
+          # Cleared first: a heartbeat touch already queued behind this
+          # disconnect must not re-join the seat it removes.
+          connection.presence_member = nil if connection.respond_to?(:presence_member=)
+          presence = @presence_provider.call(stream)
+          if joined_as_seat?(connection, presence)
+            presence.leave(connection_id: connection.id)
+          else
+            presence.leave(member_id: member_id)
+          end
+        rescue StandardError => e
+          log_failure("leave", e)
+        end
+
+        def touch(connections)
+          connections.select { |conn| presence_member_of(conn) }.group_by(&:stream_name).each do |stream, conns|
+            touch_stream(stream, conns)
+          rescue StandardError => e
+            log_failure("touch", e)
+          end
+          sweep_if_due
+        end
+
+        private
+
+        def join_member(presence, connection, member, seat)
+          metadata = member[:metadata] || {}
+          if seat
+            presence.join(member_id: member[:id], metadata: metadata,
+                          connection_id: connection.id, seat: seat_of(connection))
+          else
+            presence.join(member_id: member[:id], metadata: metadata)
+          end
+        end
+
+        def touch_stream(stream, conns)
+          presence = @presence_provider.call(stream)
+          unless seats?(presence)
+            conns.each { |conn| presence.touch(member_id: presence_member_of(conn)) }
+            return
+          end
+
+          live = presence.touch(connection_id: conns.map(&:id))
+          conns.reject { |conn| live.include?(conn.id) }.each { |conn| rejoin_as_seat(presence, conn, stream) }
+        end
+
+        # A connection that joined as a member-keyed row before the table
+        # was migrated drops that row before it takes a seat.
+        def rejoin_as_seat(presence, connection, stream)
+          presence.leave(member_id: presence_member_of(connection)) unless joined_as_seat?(connection, presence)
+          join(connection, stream)
+        end
+
+        def sweep_if_due
+          return if Array(@config.streams_presence_patterns).empty?
+
+          at = now
+          return if @last_sweep_at && at - @last_sweep_at < @config.presence_min_ttl / 2.0
+
+          @sweeper.call
+          @last_sweep_at = at
+        rescue StandardError => e
+          log_failure("expiry", e)
+        end
+
+        def expire_seats
+          Pgbus::Streams::Presence.expire_seats!(
+            ttl_for: ->(stream) { @config.presence_ttl_for(stream) }, min_ttl: @config.presence_min_ttl
+          )
+        end
+
+        def seats?(presence)
+          presence.respond_to?(:seats?) && presence.seats?
+        end
+
+        # The mode recorded at join; a connection without the accessor
+        # follows the table's current schema.
+        def joined_as_seat?(connection, presence)
+          return connection.presence_seat if connection.respond_to?(:presence_seat)
+
+          seats?(presence)
+        end
+
+        def seat_of(connection)
+          connection.respond_to?(:seat) ? connection.seat : nil
+        end
+
+        def presence_member_of(connection)
+          connection.respond_to?(:presence_member) ? connection.presence_member : nil
+        end
+
+        def now
+          @clock.call
+        end
+
+        def log_failure(action, error)
+          @logger.error { "[Pgbus::Streamer::PresenceTracker] presence #{action} failed: #{error.class}: #{error.message}" }
+        end
+      end
+    end
+  end
+end