pgbus 0.5.1 → 0.6.1

data/lib/pgbus/web/stream_app.rb

@@ -0,0 +1,179 @@
+# frozen_string_literal: true
+
+require "rack/request"
+
+module Pgbus
+  module Web
+    # Rack app mounted at /pgbus/streams. Not a Rails controller — this
+    # bypasses the entire Rails middleware stack for the streaming path,
+    # which avoids several classes of bug (ActionDispatch::Cookies leaking
+    # into long-lived requests, ActiveRecord::QueryCache holding the AR
+    # connection open, flash modifications after hijack, etc.). It also
+    # removes the temptation to use ActionController::Live, which has a
+    # well-documented tendency to tie up Puma threads (see puma#1009,
+    # puma#938, puma#569).
+    #
+    # Call flow (happy path, Puma 6.1+ hijack):
+    #   1. Parse the signed name from PATH_INFO
+    #   2. Verify it via Pgbus::Streams::SignedName (raises → 404)
+    #   3. Run the authorize hook (raises → 403)
+    #   4. Parse the cursor from ?since= or Last-Event-ID
+    #   5. Check streams_max_connections per worker (over → 503)
+    #   6. Check rack.hijack? (missing → 501)
+    #   7. Hijack, write the HTTP response line + SSE headers + opening
+    #      comment directly to the socket
+    #   8. Build a Connection, hand to Streamer.current.register(...)
+    #   9. Return [-1, {}, []] (Puma's async protocol)
+    #
+    # On errors, returns a normal [status, headers, body] response so
+    # Rack / the reverse proxy can log and retry. The whole thing is
+    # designed so a reviewer can read call/2 top-to-bottom and see the
+    # full request lifecycle.
+    class StreamApp
+      PATH_PREFIX = "/pgbus/streams"
+      private_constant :PATH_PREFIX
+
+      def initialize(streamer: nil, config: nil, logger: nil, authorize: nil)
+        @streamer_override = streamer
+        @config_override   = config
+        @logger_override   = logger
+        @authorize         = authorize || ->(_env, _stream_name) { true }
+      end
+
+      def call(env)
+        request = Rack::Request.new(env)
+        return not_found("only GET is supported") unless request.get?
+
+        signed_name = extract_signed_name(env)
+        return not_found("missing signed stream name") if signed_name.nil?
+
+        stream_name = verify!(signed_name)
+        return not_found("invalid signed stream name") if stream_name.nil?
+
+        authorize_result = @authorize.call(env, stream_name)
+        return forbidden unless authorize_result
+
+        # If authorize returned a non-boolean value (e.g. a User model),
+        # treat it as the connection context for audience filtering.
+        # A bare `true` means "authorized but no context" — filters that
+        # depend on a context will fail-closed on these connections.
+        context = authorize_result unless authorize_result == true
+
+        cursor = parse_cursor(env, request)
+        return bad_request("invalid cursor: #{cursor}") if cursor.is_a?(String)
+
+        return unsupported_server unless env["rack.hijack?"]
+
+        return over_capacity if streamer.registry.size >= config.streams_max_connections
+
+        hijack_and_register(env, stream_name: stream_name, since_id: cursor, context: context)
+      rescue StandardError => e
+        logger.error { "[Pgbus::StreamApp] #{e.class}: #{e.message}" }
+        server_error
+      end
+
+      private
+
+      def extract_signed_name(env)
+        path = env["PATH_INFO"].to_s
+        # PATH_INFO when mounted is relative to the mount point, so it's
+        # just "/<signed_name>". When not mounted (e.g. tests calling call
+        # directly) it's "/pgbus/streams/<signed_name>". Handle both.
+        path = path.sub(PATH_PREFIX, "")
+        path = path[1..] if path.start_with?("/") # strip leading slash
+        return nil if path.nil? || path.empty?
+
+        path
+      end
+
+      def verify!(token)
+        Pgbus::Streams::SignedName.verify!(token)
+      rescue Pgbus::Streams::SignedName::InvalidSignedName,
+             Pgbus::Streams::SignedName::MissingSecret
+        nil
+      end
+
+      def parse_cursor(env, request)
+        Pgbus::Streams::Cursor.parse(
+          query_since: request.params["since"],
+          last_event_id: env["HTTP_LAST_EVENT_ID"]
+        )
+      rescue Pgbus::Streams::Cursor::InvalidCursor => e
+        e.message
+      end
+
+      def hijack_and_register(env, stream_name:, since_id:, context: nil)
+        # Rack hijack is a one-time transition per the Rack spec. Call
+        # once, use the return value, and fall back to the side-effect
+        # `rack.hijack_io` variable only if the return was nil (some
+        # older Rack versions populate the side-effect var instead of
+        # returning the IO).
+        hijack_result = env["rack.hijack"].call
+        io = hijack_result || env["rack.hijack_io"]
+
+        write_headers(io, stream_name: stream_name, since_id: since_id)
+
+        connection = Pgbus::Web::Streamer::Connection.new(
+          id: SecureRandom.hex(8),
+          stream_name: stream_name,
+          io: io,
+          since_id: since_id,
+          writer: Pgbus::Web::Streamer::IoWriter,
+          write_deadline_ms: config.streams_write_deadline_ms,
+          context: context
+        )
+        streamer.register(connection)
+
+        [-1, {}, []]
+      end
+
+      def write_headers(io, stream_name:, since_id:)
+        io.write(Pgbus::Streams::Envelope.http_response_headers)
+        io.write(Pgbus::Streams::Envelope.retry_directive(2_000))
+        io.write(Pgbus::Streams::Envelope.comment("pgbus stream open since_id=#{since_id} stream=#{stream_name}"))
+      end
+
+      def streamer
+        @streamer_override || Pgbus::Web::Streamer.current
+      end
+
+      def config
+        @config_override || Pgbus.configuration
+      end
+
+      def logger
+        @logger_override || Pgbus.logger
+      end
+
+      # Error responses — plain text, no body parsing, no Rails. The client
+      # is EventSource which doesn't render error bodies, but we include
+      # short text for operator debugging.
+
+      def not_found(reason)
+        [404, { "content-type" => "text/plain" }, ["pgbus: #{reason}"]]
+      end
+
+      def forbidden
+        [403, { "content-type" => "text/plain" }, ["pgbus: forbidden"]]
+      end
+
+      def bad_request(reason)
+        [400, { "content-type" => "text/plain" }, ["pgbus: #{reason}"]]
+      end
+
+      def over_capacity
+        [503, { "content-type" => "text/plain", "retry-after" => "2" }, ["pgbus: streamer over capacity"]]
+      end
+
+      def unsupported_server
+        [501,
+         { "content-type" => "text/plain" },
+         ["pgbus streams require Puma 6.1+ or Falcon — current Rack server does not provide rack.hijack"]]
+      end
+
+      def server_error
+        [500, { "content-type" => "text/plain" }, ["pgbus: internal error"]]
+      end
+    end
+  end
+end

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

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+module Pgbus
+  module Web
+    module Streamer
+      # Wraps a single hijacked SSE client socket with its own cursor state,
+      # per-io mutex, and liveness flag. Owns no threads — the Dispatcher and
+      # Heartbeat threads call #enqueue / #write_comment on Connection instances
+      # directly, and the per-io mutex in IoWriter serialises concurrent writes.
+      #
+      # Cursor semantics: `last_msg_id_sent` is strictly monotonic. `enqueue`
+      # filters envelopes with `msg_id > last_msg_id_sent` and advances the
+      # cursor only for envelopes that actually wrote successfully. This is
+      # the client-side leg of the replay-race fix (§6.5 of the design doc).
+      class Connection
+        attr_reader :id, :stream_name, :io, :mutex, :last_msg_id_sent, :context
+
+        def initialize(id:, stream_name:, io:, since_id:, writer:, write_deadline_ms:, context: nil)
+          @id = id
+          @stream_name = stream_name
+          @io = io
+          @last_msg_id_sent = since_id.to_i
+          @writer = writer
+          @write_deadline_ms = write_deadline_ms
+          @mutex = Mutex.new
+          @dead = false
+          @closed = false
+          @created_at = monotonic
+          @last_write_at = @created_at
+          # Context is whatever the StreamApp's authorize hook returned
+          # (a truthy non-boolean value). Typically a user model or a
+          # session hash. The Dispatcher passes it to the Filters
+          # registry when evaluating visible_to predicates. Defaults to
+          # nil for tests that don't need audience filtering.
+          @context = context
+        end
+
+        def enqueue(envelopes)
+          written = []
+          envelopes.each do |envelope|
+            next if envelope.msg_id <= @last_msg_id_sent
+
+            bytes = Pgbus::Streams::Envelope.message(
+              id: envelope.msg_id,
+              event: "turbo-stream",
+              data: envelope.payload
+            )
+
+            result = @writer.write(self, bytes, deadline_ms: @write_deadline_ms)
+            if result == :ok
+              @last_msg_id_sent = envelope.msg_id
+              @last_write_at = monotonic
+              written << envelope
+            else
+              mark_dead!
+              break
+            end
+          end
+          written
+        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
+        end
+
+        def idle_for
+          monotonic - @last_write_at
+        end
+
+        def dead?
+          @dead
+        end
+
+        def mark_dead!
+          @dead = true
+        end
+
+        # Idempotent socket close for use by Instance#shutdown! and the
+        # heartbeat idle reaper. Wraps the respond_to? / closed? dance
+        # so callers don't need to know about StringIO-in-tests vs real
+        # Socket-in-prod or about the mark_dead! ordering.
+        #
+        # Takes the same mutex as IoWriter.write so it can't fire
+        # mid-write — otherwise the write loop could hit a half-closed
+        # socket and corrupt the `last_msg_id_sent` cursor by marking
+        # the connection dead between successful writes. The rescue
+        # narrows to IO-related exceptions; unrelated errors (bugs in
+        # the fake IO used by tests, nil-dereferences, etc.) should
+        # still propagate so the test suite catches them.
+        def close
+          @mutex.synchronize do
+            return if @closed
+
+            @closed = true
+            mark_dead!
+            return unless @io.respond_to?(:close)
+
+            @io.close unless @io.respond_to?(:closed?) && @io.closed?
+          end
+        rescue IOError, SystemCallError => e
+          Pgbus.logger&.debug { "[Pgbus::Streamer::Connection] close failed: #{e.class}: #{e.message}" }
+        end
+
+        private
+
+        def monotonic
+          # Qualify ::Process because Pgbus::Process already exists as a
+          # namespace for worker/supervisor/consumer and would otherwise
+          # shadow the top-level constant here.
+          ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+        end
+      end
+    end
+  end
+end