async-rabbitmq 0.4.0

data/lib/async_rabbitmq/frame_io.rb

@@ -0,0 +1,309 @@
+require "async"
+require "async/queue"
+require "async/limited_queue"
+require "async/semaphore"
+require "socket"
+require "amq/protocol"
+require_relative "errors"
+
+module AsyncRabbitMQ
+  # Handles low-level AMQP frame reading and writing over a native Ruby socket.
+  # async 2.x's fiber scheduler makes blocking socket reads yield automatically.
+  #
+  # Architecture:
+  #   Reader Task — reads frames, routes to per-channel Async::Queue.
+  #   Writer Task — drains a bounded Async::LimitedQueue, serializes all writes.
+  #   Heartbeat   — writes directly via Async::Semaphore (bypasses write queue).
+  #
+  # AMQ::Protocol::Frame.decode raises NotImplementedError (abstract stub), so
+  # frames are read in three steps: 7-byte header (decoded by
+  # AMQ::Protocol::Frame.decode_header), payload, 0xCE terminator.
+  class FrameIO
+    FRAME_HEADER_SIZE = 7
+    FRAME_TERMINATOR  = "\xCE".b.freeze
+    WRITE_QUEUE_LIMIT = 1024
+
+    # AMQP soft-error codes: broker closes the channel, connection stays alive.
+    # 311 content-too-large, 312 no-route, 313 no-consumers are channel-level errors.
+    # 403 access-refused, 404 not-found, 405 resource-locked, 406 precondition-failed.
+    SOFT_ERROR_CODES = [311, 312, 313, 403, 404, 405, 406].freeze
+
+    def initialize(socket, logger: Log.new)
+      @socket        = socket
+      @logger        = logger
+      @channels      = {}
+      @channel_sem   = Async::Semaphore.new(1)
+      @write_queue   = Async::LimitedQueue.new(WRITE_QUEUE_LIMIT)
+      @socket_sem    = Async::Semaphore.new(1)
+      @running            = false
+      @reader_task        = nil
+      @writer_task        = nil
+      @content_state      = {}
+      @on_frame           = nil   # optional proc called each time a frame is read
+      @blocked            = false
+      @unblocked_condition = nil
+    end
+
+    # Set a callback invoked immediately after each frame is successfully read.
+    # Used by Session to update its heartbeat timestamp.
+    def on_frame=(proc)
+      @on_frame = proc
+    end
+
+    def register_channel(channel_id)
+      @channel_sem.acquire do
+        q = @channels[channel_id] = Async::Queue.new
+        q
+      end
+    end
+
+    def unregister_channel(channel_id)
+      @channel_sem.acquire do
+        @channels.delete(channel_id)
+        @content_state.delete(channel_id)
+      end
+    end
+
+    def channel_queue(channel_id)
+      @channels[channel_id]
+    end
+
+    # +spawn+ makes the reader and writer tasks. It defaults to the calling
+    # task, which is only right for a throwaway connection: the session passes
+    # one that parents them at the reactor, so a short-lived task that happened
+    # to call connect cannot take the connection's IO down with it.
+    def start(task: Async::Task.current, spawn: ->(&block) { task.async(&block) })
+      @running = true
+      @writer_task = spawn.call { writer_loop }
+      @reader_task = spawn.call { reader_loop }
+    end
+
+    def stop
+      @running = false
+      @socket.close rescue nil
+      # Publishers parked on the connection.blocked gate can never proceed now.
+      @blocked = false
+      release_blocked_writers(ConnectionError.new(code: 0, text: "Connection closed while blocked by the broker"))
+      # Push nil sentinel to every channel queue FIRST so blocked poppers
+      # (e.g. wait_channel0_method) are unblocked before we cancel tasks.
+      # Cancelling the CURRENT task raises Async::Cancel (< Exception, not
+      # StandardError), which bypasses inline `rescue nil` and would skip
+      # the push if we did it after.
+      @channels.each_value { |q| q.push(nil) rescue nil }
+      @writer_task&.cancel rescue nil
+      @reader_task&.cancel rescue nil
+    end
+
+    # Enqueue raw frame bytes for writing. Yields the fiber when the queue is
+    # full. Content publishes (+publish: true+) additionally wait while the
+    # broker has the connection blocked (connection.blocked). Control frames
+    # (close, close-ok, acks, ...) are never gated, so shutdown cannot hang on
+    # a resource alarm. Raises ConnectionError if the connection is closed
+    # while a publish is parked.
+    def write_frame(data, publish: false)
+      if publish
+        while @blocked
+          @unblocked_condition ||= Async::Condition.new
+          outcome = @unblocked_condition.wait
+          raise outcome if outcome.is_a?(Exception)
+        end
+      end
+      @write_queue.push(data)
+    end
+
+    # Called by the Session channel-0 monitor when connection.blocked is received.
+    def set_blocked(reason)
+      @blocked = true
+      @logger.warn("Connection blocked: #{reason}")
+    end
+
+    # Called by the Session channel-0 monitor when connection.unblocked is received.
+    def set_unblocked
+      @blocked = false
+      release_blocked_writers
+      @logger.info("Connection unblocked")
+    end
+
+    def blocked?
+      @blocked
+    end
+
+    # How many writes are queued for the writer fiber but not yet on the socket.
+    # One entry is one write_frame call, so a batch published with
+    # basic_publish_batch counts once however many messages it carries. Used at
+    # close time to report what is about to be thrown away.
+    def pending_writes
+      @write_queue.size
+    end
+
+    # Write heartbeat directly to socket, bypassing the write queue so
+    # heartbeats are never starved by backpressure.
+    # Silently swallows IO errors — the reader_loop will detect and recover.
+    def write_heartbeat
+      heartbeat = AMQ::Protocol::HeartbeatFrame.encode
+      @socket_sem.acquire { @socket.write(heartbeat) }
+    rescue IOError, Errno::EPIPE, Errno::EBADF
+      # Socket is closed; reader_loop will raise and call trigger_recovery.
+    end
+
+    private
+
+    # Wake every publisher waiting on the blocked gate. With +outcome+ nil they
+    # re-check the gate; with an exception they raise it.
+    def release_blocked_writers(outcome = nil)
+      cond = @unblocked_condition
+      @unblocked_condition = nil
+      cond&.signal(outcome)
+    end
+
+    def writer_loop
+      while @running
+        data = @write_queue.pop
+        @socket_sem.acquire { @socket.write(data) }
+      end
+    rescue IOError, Errno::EPIPE, Errno::EBADF, Errno::ECONNRESET
+      # Socket closed — reader_loop will detect and call trigger_recovery.
+    rescue => e
+      # The writer is gone either way; without recovery every publisher would
+      # block forever once the write queue fills, with nothing to drain it.
+      @logger.error("FrameIO writer: #{e.class}: #{e.message}")
+      trigger_recovery(e) if @running
+    end
+
+    def reader_loop
+      while @running
+        frame = read_frame
+        @on_frame.call if @on_frame
+        dispatch(frame)
+      end
+    rescue EOFError, Errno::ECONNRESET, Errno::EPIPE, IOError, Errno::EBADF => e
+      if @running
+        @logger.warn("FrameIO reader closed: #{e.class}")
+        trigger_recovery(e)
+      else
+        # #stop closed the socket under us: an orderly shutdown, not a failure.
+        @logger.debug("FrameIO reader stopped (#{e.class})")
+      end
+    rescue AMQ::Protocol::Error => e
+      @logger.error("FrameIO AMQP error: #{e.class}: #{e.message}")
+      trigger_recovery(e)
+    rescue => e
+      @logger.error("FrameIO reader error: #{e.class}: #{e.message}")
+      trigger_recovery(e)
+    end
+
+    # Read one AMQP frame via manual 3-step IO.
+    # Returns an AMQ::Protocol::MethodFrame / HeaderFrame / BodyFrame / HeartbeatFrame.
+    def read_frame
+      header_bytes = read_exactly(FRAME_HEADER_SIZE)
+      # Raises AMQ::Protocol::FrameTypeError (an AMQ::Protocol::Error) on an
+      # unknown frame type, which reader_loop turns into recovery.
+      type, channel_id, payload_size = AMQ::Protocol::Frame.decode_header(header_bytes)
+      payload    = read_exactly(payload_size)
+      terminator = read_exactly(1)
+
+      unless terminator == FRAME_TERMINATOR
+        raise AMQ::Protocol::Error,
+              "Invalid frame terminator: #{terminator.inspect} (expected 0xCE)"
+      end
+
+      AMQ::Protocol::Frame::CLASSES[AMQ::Protocol::Frame::TYPES[type]].new(payload, channel_id)
+    end
+
+    # Read exactly n bytes, yielding the fiber at each blocking read.
+    def read_exactly(n)
+      return "".b if n == 0
+      buf = "".b
+      while buf.bytesize < n
+        chunk = @socket.read(n - buf.bytesize)
+        raise EOFError, "Socket closed mid-frame" if chunk.nil? || chunk.empty?
+        buf << chunk
+      end
+      buf
+    end
+
+    def dispatch(frame_obj)
+      channel_id = frame_obj.channel
+
+      case frame_obj
+      when AMQ::Protocol::MethodFrame
+        method = frame_obj.decode_payload
+
+        # Broker-initiated Connection::Close must be acknowledged immediately.
+        # Without a timely CloseOk the broker waits (up to ~30 s) before
+        # force-closing the TCP connection, e.g. when a vhost is deleted while
+        # a connection is still open.  We send CloseOk synchronously here,
+        # then route the method to channel 0 (so wait_channel0_method can also
+        # see it and raise a proper error if it is active), and finally trigger
+        # recovery so the session reconnects.
+        if channel_id == 0 && method.is_a?(AMQ::Protocol::Connection::Close)
+          begin
+            @socket_sem.acquire do
+              @socket.write(AMQ::Protocol::Connection::CloseOk.encode.encode)
+            end
+          rescue => e
+            @logger.warn("FrameIO: could not send Connection::CloseOk — #{e.message}")
+          end
+          route_to_channel(0, [:method, method])
+          trigger_recovery(
+            ConnectionError.new(code: method.reply_code, text: method.reply_text)
+          )
+          return
+        end
+
+        route_to_channel(channel_id, [:method, method])
+
+      when AMQ::Protocol::HeaderFrame
+        # A zero-length body has no body frames at all: neither RabbitMQ nor
+        # amq-protocol emits one. Route now, or the content would sit here
+        # forever and the next method frame would be taken as its body.
+        if frame_obj.body_size == 0
+          route_to_channel(channel_id, [:content, frame_obj, "".b])
+          return
+        end
+
+        @content_state[channel_id] = {
+          header:     frame_obj,
+          body_parts: [],
+          remaining:  frame_obj.body_size,
+        }
+
+      when AMQ::Protocol::BodyFrame
+        state = @content_state[channel_id]
+        unless state
+          @logger.warn("Body frame on channel #{channel_id} without content header — dropped")
+          return
+        end
+        chunk = frame_obj.payload
+        state[:body_parts] << chunk
+        state[:remaining]  -= chunk.bytesize
+
+        if state[:remaining] <= 0
+          body = state[:body_parts].join.b
+          @content_state.delete(channel_id)
+          route_to_channel(channel_id, [:content, state[:header], body])
+        end
+
+      when AMQ::Protocol::HeartbeatFrame
+        route_to_channel(0, [:heartbeat])
+
+      else
+        @logger.warn("Unknown frame type #{frame_obj.class} on channel #{channel_id} — dropped")
+      end
+    end
+
+    def route_to_channel(channel_id, message)
+      queue = @channels[channel_id]
+      if queue
+        queue.push(message)
+      else
+        @logger.debug("No queue for channel #{channel_id} — frame dropped (#{message.first})")
+      end
+    end
+
+    # Overridden by Session to initiate reconnect.
+    def trigger_recovery(error)
+      @running = false
+    end
+  end
+end

data/lib/async_rabbitmq/log.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module AsyncRabbitMQ
+  # The default logger: warnings and errors, one line each, to $stderr.
+  #
+  # Anything that responds to +debug+, +info+, +warn+ and +error+ can be passed
+  # as +logger:+ instead, and that is the expected thing to do in an
+  # application: Ruby's Logger, Rails.logger, a Console wrapper, your own
+  # object. The gem deliberately does not depend on the logger gem, which
+  # stopped being a default gem in Ruby 4.0, so the choice of logging stack
+  # stays with the application.
+  #
+  #   AsyncRabbitMQ::Session.new(logger: Rails.logger)
+  #   AsyncRabbitMQ::Session.new(logger: Logger.new($stdout, level: Logger::INFO))
+  #   AsyncRabbitMQ::Session.new(logger: AsyncRabbitMQ::Log.silent)
+  #   AsyncRabbitMQ::Session.new(logger: AsyncRabbitMQ::Log.new($stdout, level: :debug))
+  class Log
+    LEVELS = { debug: 0, info: 1, warn: 2, error: 3, fatal: 4 }.freeze
+    DEFAULT_LEVEL = :warn
+
+    attr_reader :level
+
+    # +io+ nil discards everything.
+    def initialize(io = $stderr, level: DEFAULT_LEVEL)
+      @io = io
+      self.level = level
+    end
+
+    # Accepts a symbol (:debug) or Logger's integer constants (Logger::DEBUG).
+    def level=(value)
+      @level = value.is_a?(Integer) ? (LEVELS.key(value) || DEFAULT_LEVEL) : value.to_sym
+      @threshold = LEVELS.fetch(@level, LEVELS[DEFAULT_LEVEL])
+    end
+
+    # Discards everything. Useful in tests and in one-shot scripts.
+    def self.silent
+      new(nil)
+    end
+
+    LEVELS.each_key do |name|
+      define_method(name) do |message = nil, &block|
+        return nil if @io.nil? || LEVELS.fetch(name) < @threshold
+
+        text = message.nil? && block ? block.call : message
+        @io.puts("#{name.to_s.upcase.ljust(5)} async-rabbitmq: #{text}")
+        nil
+      end
+
+      define_method(:"#{name}?") { !@io.nil? && LEVELS.fetch(name) >= @threshold }
+    end
+  end
+end

data/lib/async_rabbitmq/notifier.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+module AsyncRabbitMQ
+  # Structured events for metrics, tracing and debugging.
+  #
+  # Nothing is emitted until something subscribes. Every call site goes through
+  # Instrumented#instrument, which checks for subscribers before running the
+  # block that builds the payload, so an unwatched session pays one method call
+  # and one array check per event and allocates nothing.
+  #
+  #   session.on_event { |name, payload| logger.info("#{name} #{payload}") }
+  #   session.on_event("message.")     { |name, payload| statsd.increment(name) }
+  #   session.on_event("channel.rpc")  { |_, p| histogram.record(p[:duration]) }
+  #   session.on_event(/recovery/)     { |name, _| pager.note(name) }
+  #
+  # A pattern is nil for everything, a String for one event name or, ending in
+  # a dot, a prefix, or a Regexp. Subscribe returns a handle for #unsubscribe.
+  #
+  # A subscriber that raises is logged and skipped: instrumentation must never
+  # take the connection down with it.
+  class Notifier
+    # The events the client emits. Names are API: they do not change without a
+    # major version. Payload keys are documented in README.md.
+    EVENTS = %w[
+      connection.open
+      connection.closed
+      connection.lost
+      connection.blocked
+      connection.unblocked
+      recovery.attempt
+      recovery.succeeded
+      recovery.exhausted
+      heartbeat.sent
+      channel.open
+      channel.closed
+      channel.rpc
+      consumer.registered
+      consumer.cancelled
+      message.published
+      message.confirmed
+      message.returned
+      message.consumed
+    ].freeze
+
+    def initialize(logger: nil)
+      @subscribers = []
+      @logger      = logger
+    end
+
+    def subscribe(pattern = nil, &block)
+      raise ArgumentError, "a subscriber block is required" unless block
+
+      @subscribers << [pattern, block]
+      block
+    end
+
+    # Remove a subscriber using the handle subscribe returned.
+    def unsubscribe(handle)
+      @subscribers.reject! { |(_, block)| block.equal?(handle) }
+      nil
+    end
+
+    def subscribed?
+      !@subscribers.empty?
+    end
+
+    def subscriber_count
+      @subscribers.size
+    end
+
+    # Called by Instrumented#instrument once the payload exists.
+    def publish(name, payload)
+      # Snapshot: a subscriber that unsubscribes itself (or another) mutates
+      # @subscribers under the iterator, which skips the next one along.
+      @subscribers.dup.each do |(pattern, block)|
+        next unless match?(pattern, name)
+
+        begin
+          block.call(name, payload)
+        rescue => e
+          @logger&.error("Event subscriber for #{name} raised #{e.class}: #{e.message}")
+        end
+      end
+      nil
+    end
+
+    private
+
+    def match?(pattern, name)
+      case pattern
+      when nil    then true
+      when String then pattern.end_with?(".") ? name.start_with?(pattern) : name == pattern
+      when Regexp then pattern.match?(name)
+      else             pattern === name
+      end
+    end
+  end
+
+  # Mixed into Session and Channel, both of which hold a @notifier.
+  module Instrumented
+    private
+
+    # Emit +name+ with the payload the block returns. The block runs only when
+    # something is subscribed, so building the payload costs nothing otherwise.
+    def instrument(name)
+      notifier = @notifier
+      return unless notifier&.subscribed?
+
+      notifier.publish(name, yield)
+    end
+
+    # For call sites that need to time the work: returns nil when nobody is
+    # listening, so the caller can skip the clock reads too.
+    def instrument_clock
+      @notifier&.subscribed? ? Process.clock_gettime(Process::CLOCK_MONOTONIC) : nil
+    end
+
+    def instrument_elapsed(started)
+      Process.clock_gettime(Process::CLOCK_MONOTONIC) - started
+    end
+  end
+end

data/lib/async_rabbitmq/pool.rb

@@ -0,0 +1,43 @@
+begin
+  require "async/pool"
+rescue LoadError
+  raise LoadError, "AsyncRabbitMQ::Pool needs the async-pool gem, which this gem does not " \
+                   "depend on: add gem \"async-pool\" to your Gemfile"
+end
+require_relative "session"
+
+module AsyncRabbitMQ
+  # Opt-in connection pool backed by Async::Pool::Controller.
+  #
+  # Each pooled resource is a connected Session. A Session advertises its
+  # negotiated channel_max as its +concurrency+, so the controller hands the
+  # same Session to many fibers at once and only opens another connection
+  # when every existing one is saturated or unusable.
+  #
+  # Usage:
+  #   pool = AsyncRabbitMQ::Pool.new(max: 5, host: "localhost")
+  #   pool.acquire { |session| session.with_channel { |ch| ch.basic_publish(...) } }
+  #   pool.close
+  class Pool
+    def initialize(max: 4, **session_opts)
+      @session_opts = session_opts
+      constructor   = -> { Session.new(**@session_opts).tap(&:connect) }
+      @pool         = Async::Pool::Controller.new(constructor, limit: max)
+    end
+
+    # Acquire a connected Session. With a block, the Session is released when
+    # the block returns; without one, the caller must call #release.
+    def acquire(&block)
+      @pool.acquire(&block)
+    end
+
+    def release(session)
+      @pool.release(session)
+    end
+
+    # Close every pooled Session.
+    def close
+      @pool.close
+    end
+  end
+end

data/lib/async_rabbitmq/queue.rb

@@ -0,0 +1,88 @@
+module AsyncRabbitMQ
+  # Represents an AMQP queue bound to a channel.
+  # Created via channel.queue("name", durable: true).
+  class Queue
+    # Queue types RabbitMQ knows, for Channel#durable_queue(name, type).
+    module Types
+      CLASSIC = "classic"
+      QUORUM  = "quorum"
+      STREAM  = "stream"
+    end
+
+    attr_reader :name, :message_count, :consumer_count
+
+    def initialize(name, message_count, consumer_count, channel, durable: false, exclusive: false, auto_delete: false)
+      @name           = name
+      @message_count  = message_count
+      @consumer_count = consumer_count
+      @channel        = channel
+      @durable        = durable
+      @exclusive      = exclusive
+      @auto_delete    = auto_delete
+    end
+
+    def durable?
+      @durable
+    end
+
+    def exclusive?
+      @exclusive
+    end
+
+    def auto_delete?
+      @auto_delete
+    end
+
+    # Internal: topology recovery re-declared this server-named queue under a
+    # new broker-generated name.
+    def update_name_to(new_name)
+      @name = new_name
+    end
+
+    def server_named?
+      @name.start_with?("amq.gen-")
+    end
+
+    # Re-declare with passive: true to refresh message_count and consumer_count.
+    def status
+      q = @channel.queue(@name, passive: true)
+      @message_count  = q.message_count
+      @consumer_count = q.consumer_count
+      { message_count: @message_count, consumer_count: @consumer_count }
+    end
+
+    # Synchronously fetch one message: [delivery_info, header, body], or nil if
+    # the queue is empty. Defaults to manual acknowledgement like Bunny: a
+    # message fetched and then dropped by the caller is requeued, not lost.
+    def pop(manual_ack: true)
+      @channel.basic_get(@name, manual_ack: manual_ack)
+    end
+    alias get pop
+
+    def subscribe(manual_ack: false, **opts, &block)
+      @channel.basic_consume(@name, manual_ack: manual_ack, **opts, &block)
+    end
+
+    def publish(payload, routing_key: @name, **opts)
+      @channel.basic_publish(payload, routing_key: routing_key, **opts)
+    end
+
+    def bind(exchange:, routing_key: "", arguments: {})
+      @channel.queue_bind(@name, exchange: exchange, routing_key: routing_key, arguments: arguments)
+      self
+    end
+
+    def unbind(exchange:, routing_key: "", arguments: {})
+      @channel.queue_unbind(@name, exchange: exchange, routing_key: routing_key, arguments: arguments)
+      self
+    end
+
+    def delete(if_unused: false, if_empty: false)
+      @channel.queue_delete(@name, if_unused: if_unused, if_empty: if_empty)
+    end
+
+    def purge
+      @channel.queue_purge(@name)
+    end
+  end
+end