omq 0.5.1 → 0.6.1

data/lib/omq/zmtp/engine.rb

@@ -14,22 +14,27 @@ module OMQ
       #
       attr_reader :socket_type
 
+
       # @return [Options] socket options
       #
       attr_reader :options
 
+
       # @return [Routing] routing strategy
       #
       attr_reader :routing
 
+
       # @return [String, nil] last bound endpoint
       #
       attr_reader :last_endpoint
 
+
       # @return [Integer, nil] last auto-selected TCP port
       #
       attr_reader :last_tcp_port
 
+
       # @param socket_type [Symbol] e.g. :REQ, :REP, :PAIR
       # @param options [Options]
       #
@@ -43,10 +48,24 @@ module OMQ
         @listeners            = []
         @tasks                = []
         @closed               = false
+        @closing              = false
         @last_endpoint        = nil
         @last_tcp_port        = nil
+        @peer_connected       = Async::Promise.new
+        @all_peers_gone       = Async::Promise.new
+        @reconnect_enabled    = true
+        @parent_task          = nil
+        @connection_promises  = {} # connection => Async::Promise
+        @fatal_error          = nil
       end
 
+
+      attr_reader :peer_connected, :all_peers_gone, :connections, :parent_task
+
+
+      attr_writer :reconnect_enabled
+
+
       # Binds to an endpoint.
       #
       # @param endpoint [String] e.g. "tcp://127.0.0.1:5555", "inproc://foo"
@@ -54,6 +73,7 @@ module OMQ
       # @raise [ArgumentError] on unsupported transport
       #
       def bind(endpoint)
+        capture_parent_task
         transport = transport_for(endpoint)
         listener = transport.bind(endpoint, self)
         @listeners << listener
@@ -61,12 +81,14 @@ module OMQ
         @last_tcp_port = extract_tcp_port(listener.endpoint)
       end
 
+
       # Connects to an endpoint.
       #
       # @param endpoint [String]
       # @return [void]
       #
       def connect(endpoint)
+        capture_parent_task
         @connected_endpoints << endpoint
         if endpoint.start_with?("inproc://")
           # Inproc connect is synchronous and instant
@@ -78,6 +100,7 @@ module OMQ
         end
       end
 
+
       # Disconnects from an endpoint. Closes connections to that endpoint
       # and stops auto-reconnection for it.
       #
@@ -95,6 +118,7 @@ module OMQ
         end
       end
 
+
       # Unbinds from an endpoint. Stops the listener and closes all
       # connections that were accepted on it.
       #
@@ -117,6 +141,7 @@ module OMQ
         end
       end
 
+
       # Called by a transport when an incoming connection is accepted.
       #
       # @param io [#read, #write, #close]
@@ -124,18 +149,20 @@ module OMQ
       # @return [void]
       #
       def handle_accepted(io, endpoint: nil)
-        setup_connection(io, as_server: true, endpoint: endpoint)
+        spawn_connection(io, as_server: true, endpoint: endpoint)
       end
 
+
       # Called by a transport when an outgoing connection is established.
       #
       # @param io [#read, #write, #close]
       # @return [void]
       #
       def handle_connected(io, endpoint: nil)
-        setup_connection(io, as_server: false, endpoint: endpoint)
+        spawn_connection(io, as_server: false, endpoint: endpoint)
       end
 
+
       # Called by inproc transport with a pre-validated DirectPipe.
       # Skips ZMTP handshake — just registers with routing strategy.
       #
@@ -146,25 +173,43 @@ module OMQ
         @connections << pipe
         @connection_endpoints[pipe] = endpoint if endpoint
         @routing.connection_added(pipe)
+        @peer_connected.resolve(pipe)
       end
 
+
       # Dequeues the next received message. Blocks until available.
       #
       # @return [Array<String>] message parts
+      # @raise if a background pump task crashed
       #
       def dequeue_recv
-        @routing.recv_queue.dequeue
+        raise @fatal_error if @fatal_error
+        msg = @routing.recv_queue.dequeue
+        raise @fatal_error if msg.nil? && @fatal_error
+        msg
       end
 
+
+      # Pushes a nil sentinel into the recv queue, unblocking a
+      # pending {#dequeue_recv} with a nil return value.
+      #
+      def dequeue_recv_sentinel
+        @routing.recv_queue.push(nil)
+      end
+
+
       # Enqueues a message for sending. Blocks at HWM.
       #
       # @param parts [Array<String>]
       # @return [void]
+      # @raise if a background pump task crashed
       #
       def enqueue_send(parts)
+        raise @fatal_error if @fatal_error
         @routing.enqueue(parts)
       end
 
+
       # Starts a recv pump for a connection, or wires the inproc
       # fast path when the connection is a DirectPipe.
       #
@@ -180,17 +225,22 @@ module OMQ
           return nil
         end
 
-        Reactor.spawn_pump do
+        Reactor.spawn_pump(annotation: "recv pump") do
           loop do
             msg = conn.receive_message
-            msg = transform ? transform.call(msg) : msg
+            msg = transform ? transform.call(msg).freeze : msg
             recv_queue.enqueue(msg)
           end
-        rescue *CONNECTION_LOST
+        rescue Async::Stop
+          # normal shutdown
+        rescue ProtocolError, *CONNECTION_LOST
           connection_lost(conn)
+        rescue => error
+          signal_fatal_error(error)
         end
       end
 
+
       # Called when a connection is lost.
       #
       # @param connection [Connection]
@@ -202,41 +252,133 @@ module OMQ
         @routing.connection_removed(connection)
         connection.close
 
+        # Signal the connection task to exit.
+        done = @connection_promises.delete(connection)
+        done&.resolve(true)
+
+        # Resolve all_peers_gone once: had peers, now have none.
+        if @peer_connected.resolved? && @connections.empty?
+          @all_peers_gone.resolve(true)
+        end
+
         # Auto-reconnect if this was a connected (not bound) endpoint
-        if endpoint && @connected_endpoints.include?(endpoint) && !@closed
+        if endpoint && @connected_endpoints.include?(endpoint) && !@closed && !@closing && @reconnect_enabled
           schedule_reconnect(endpoint)
         end
       end
 
+
       # Closes all connections and listeners.
       #
       # @return [void]
       #
       def close
-        return if @closed
-        @closed = true
+        return if @closed || @closing
+        @closing = true
+
+        # Stop accepting new connections — but only if we already have
+        # peers to drain to. With zero connections the listeners must
+        # stay open so late-arriving peers can still receive queued
+        # messages during the linger period.
+        unless @connections.empty?
+          @listeners.each(&:stop)
+          @listeners.clear
+        end
 
         # Linger: wait for send queues to drain before closing.
         # linger=0 → close immediately, linger=nil → wait forever.
+        # @closed is set AFTER draining so reconnect tasks keep
+        # running during the linger period.
         linger = @options.linger
         if linger.nil? || linger > 0
           drain_timeout = linger # nil = wait forever, >0 = seconds
           drain_send_queues(drain_timeout)
         end
 
+        @closed = true
+
+        # Stop any remaining listeners.
+        @listeners.each(&:stop)
+        @listeners.clear
+
         # Close connections — causes pump tasks to get EOFError/IOError
         @connections.each(&:close)
         @connections.clear
-        @listeners.each(&:stop)
-        @listeners.clear
         # Stop any remaining pump tasks
         @routing.stop rescue nil
         @tasks.each { |t| t.stop rescue nil }
         @tasks.clear
       end
 
+
+      # Spawns a transient pump task with error propagation.
+      #
+      # Unexpected exceptions are caught and forwarded to
+      # {#signal_fatal_error} so blocked callers (send/recv)
+      # see the real error instead of deadlocking.
+      #
+      # @param annotation [String] task annotation for debugging
+      # @yield the pump loop body
+      # @return [Async::Task]
+      #
+      def spawn_pump_task(annotation:, &block)
+        @parent_task.async(transient: true, annotation: annotation) do
+          yield
+        rescue Async::Stop, ProtocolError, *CONNECTION_LOST
+          # normal shutdown / expected disconnect
+        rescue => error
+          signal_fatal_error(error)
+        end
+      end
+
+
+      # Wraps an unexpected pump error as {OMQ::SocketDeadError} and
+      # unblocks any callers waiting on the recv queue.
+      #
+      # Must be called from inside a rescue block so that +error+ is
+      # +$!+ and Ruby sets it as +#cause+ on the new exception.
+      #
+      # @param error [Exception]
+      #
+      def signal_fatal_error(error)
+        return if @closing || @closed
+        @fatal_error = begin
+          raise OMQ::SocketDeadError, "internal error killed #{@socket_type} socket"
+        rescue => wrapped
+          wrapped
+        end
+        @routing.recv_queue.enqueue(nil) rescue nil
+      end
+
+
       private
 
+
+      # Saves the current Async task so connection subtrees can be
+      # spawned as siblings of the caller's task.
+      #
+      def capture_parent_task
+        @parent_task ||= Async::Task.current? ? Async::Task.current : nil
+      end
+
+
+      # Spawns an isolated connection task as a sibling of accept/reconnect
+      # tasks. All per-connection children (heartbeat, recv pump, reaper)
+      # live inside this task. When the connection dies, the entire subtree
+      # is cleaned up by Async.
+      #
+      def spawn_connection(io, as_server:, endpoint: nil)
+        task = @parent_task&.async(transient: true, annotation: "conn #{endpoint}") do
+          done = Async::Promise.new
+          setup_connection(io, as_server: as_server, endpoint: endpoint, done: done)
+          done.wait
+        rescue ProtocolError, *CONNECTION_LOST
+          # handshake failed or connection lost — subtree cleaned up
+        end
+        @tasks << task if task
+      end
+
+
       # Waits for the send queue to drain.
       #
       # @param timeout [Numeric, nil] max seconds to wait (nil = forever)
@@ -245,7 +387,8 @@ module OMQ
         return unless @routing.respond_to?(:send_queue)
         deadline = timeout ? Async::Clock.now + timeout : nil
 
-        until @routing.send_queue.empty?
+        until @routing.send_queue.empty? &&
+              (!@routing.respond_to?(:send_pump_idle?) || @routing.send_pump_idle?)
           if deadline
             remaining = deadline - Async::Clock.now
             break if remaining <= 0
@@ -254,12 +397,21 @@ module OMQ
         end
       end
 
-      def setup_connection(io, as_server:, endpoint: nil)
+
+      # Performs the ZMTP handshake, starts heartbeating, and registers
+      # the new connection with the routing strategy.
+      #
+      # @param io [#read, #write, #close] underlying transport stream
+      # @param as_server [Boolean] whether we are the ZMTP server side
+      # @param endpoint [String, nil] endpoint for reconnection tracking
+      # @param done [Async::Promise, nil] resolved when the connection is lost
+      #
+      def setup_connection(io, as_server:, endpoint: nil, done: nil)
         conn = Connection.new(
           io,
-          socket_type:       @socket_type.to_s,
-          identity:          @options.identity,
-          as_server:         as_server,
+          socket_type:        @socket_type.to_s,
+          identity:           @options.identity,
+          as_server:          as_server,
           mechanism:          @options.mechanism,
           heartbeat_interval: @options.heartbeat_interval,
           heartbeat_ttl:      @options.heartbeat_ttl,
@@ -270,12 +422,21 @@ module OMQ
         conn.start_heartbeat
         @connections << conn
         @connection_endpoints[conn] = endpoint if endpoint
+        @connection_promises[conn]  = done if done
         @routing.connection_added(conn)
+        @peer_connected.resolve(conn)
       rescue ProtocolError, *CONNECTION_LOST
         conn&.close
         raise
       end
 
+
+      # Spawns a background task that reconnects to the given endpoint
+      # with exponential back-off based on the reconnect_interval option.
+      #
+      # @param endpoint [String] endpoint to reconnect to
+      # @param delay [Numeric, nil] initial delay in seconds (defaults to reconnect_interval)
+      #
       def schedule_reconnect(endpoint, delay: nil)
         ri = @options.reconnect_interval
         if ri.is_a?(Range)
@@ -286,7 +447,7 @@ module OMQ
           max_delay = nil
         end
 
-        @tasks << Reactor.spawn_pump do
+        @tasks << Reactor.spawn_pump(annotation: "reconnect #{endpoint}") do
           loop do
             break if @closed
             sleep delay if delay > 0
@@ -314,6 +475,7 @@ module OMQ
         end
       end
 
+
       def extract_tcp_port(endpoint)
         return nil unless endpoint&.start_with?("tcp://")
         port = endpoint.split(":").last.to_i

data/lib/omq/zmtp/options.rb

@@ -24,13 +24,14 @@ module OMQ
         @heartbeat_interval    = nil   # seconds, nil = disabled
         @heartbeat_ttl         = nil   # seconds, nil = use heartbeat_interval
         @heartbeat_timeout     = nil   # seconds, nil = use heartbeat_interval
-        @max_message_size       = nil  # bytes, nil = unlimited
-        @mechanism              = Mechanism::Null.new
+        @max_message_size      = nil  # bytes, nil = unlimited
+        @conflate              = false
+        @mechanism             = Mechanism::Null.new
       end
 
       attr_accessor :send_hwm,  :recv_hwm,
                     :linger,    :identity,
-                    :router_mandatory,
+                    :router_mandatory,  :conflate,
                     :read_timeout,          :write_timeout,
                     :reconnect_interval,
                     :heartbeat_interval,    :heartbeat_ttl,    :heartbeat_timeout,

data/lib/omq/zmtp/reactor.rb

@@ -24,13 +24,13 @@ module OMQ
         #
         # @return [#stop] a stoppable handle
         #
-        def spawn_pump(&block)
+        def spawn_pump(annotation: nil, &block)
           if Async::Task.current?
-            Async(transient: true, &block)
+            Async(transient: true, annotation: annotation, &block)
           else
             handle = PumpHandle.new
             ensure_started
-            @work_queue.push([:spawn, block, handle])
+            @work_queue.push([:spawn, block, handle, annotation])
             handle
           end
         end
@@ -81,6 +81,10 @@ module OMQ
 
         private
 
+        # Runs the shared Async reactor loop, dispatching work items.
+        #
+        # @param ready [Thread::Queue] signaled once the reactor is accepting work
+        #
         def run_reactor(ready)
           Async do |task|
             ready.push(true)
@@ -88,8 +92,8 @@ module OMQ
               item = @work_queue.dequeue
               case item[0]
               when :spawn
-                _, block, handle = item
-                async_task = task.async(transient: true, &block)
+                _, block, handle, annotation = item
+                async_task = task.async(transient: true, annotation: annotation, &block)
                 handle.task = async_task
               when :run
                 _, block, result_queue = item
@@ -113,6 +117,7 @@ module OMQ
         #
         attr_accessor :task
 
+
         # Stops the pump task.
         #
         # @return [void]

data/lib/omq/zmtp/routing/channel.rb

@@ -14,7 +14,8 @@ module OMQ
           @connection = nil
           @recv_queue = Async::LimitedQueue.new(engine.options.recv_hwm)
           @send_queue = Async::LimitedQueue.new(engine.options.send_hwm)
-          @tasks      = []
+          @tasks          = []
+          @send_pump_idle = true
         end
 
         # @return [Async::LimitedQueue]
@@ -56,10 +57,15 @@ module OMQ
 
         private
 
+        def send_pump_idle? = @send_pump_idle
+
+
         def start_send_pump(conn)
-          @send_pump = Reactor.spawn_pump do
+          @send_pump = @engine.spawn_pump_task(annotation: "send pump") do
             loop do
+              @send_pump_idle = true
               batch = [@send_queue.dequeue]
+              @send_pump_idle = false
               Routing.drain_send_queue(@send_queue, batch)
               batch.each { |parts| conn.write_message(parts) }
               conn.flush

data/lib/omq/zmtp/routing/fan_out.rb

@@ -12,6 +12,8 @@ module OMQ
       # their #initialize.
       #
       module FanOut
+        attr_reader :subscriber_joined
+
         private
 
         def init_fan_out(engine)
@@ -19,6 +21,9 @@ module OMQ
           @subscriptions      = {} # connection => Set of prefixes
           @send_queue         = Async::LimitedQueue.new(engine.options.send_hwm)
           @send_pump_started  = false
+          @send_pump_idle     = true
+          @conflate           = engine.options.conflate
+          @subscriber_joined  = Async::Promise.new
         end
 
         # @return [Boolean] whether the connection is subscribed to the topic
@@ -38,6 +43,7 @@ module OMQ
         #
         def on_subscribe(conn, prefix)
           @subscriptions[conn] << prefix
+          @subscriber_joined.resolve(conn) unless @subscriber_joined.resolved?
         end
 
         # Called when a cancel command is received from a peer.
@@ -50,23 +56,48 @@ module OMQ
           @subscriptions[conn]&.delete(prefix)
         end
 
+        # @return [Boolean] true when the send pump is idle (not sending a batch)
+        def send_pump_idle? = @send_pump_idle
+
+
         def start_send_pump
           @send_pump_started = true
-          @tasks << Reactor.spawn_pump do
+          @tasks << @engine.spawn_pump_task(annotation: "send pump") do
             loop do
+              @send_pump_idle = true
               batch = [@send_queue.dequeue]
+              @send_pump_idle = false
               Routing.drain_send_queue(@send_queue, batch)
 
               written = Set.new
-              batch.each do |parts|
-                topic = parts.first || "".b
-                @connections.each do |conn|
-                  next unless subscribed?(conn, topic)
+
+              if @conflate
+                # Keep only the last matching message per connection.
+                latest = {} # conn => parts
+                batch.each do |parts|
+                  topic = parts.first || "".b
+                  @connections.each do |conn|
+                    next unless subscribed?(conn, topic)
+                    latest[conn] = parts
+                  end
+                end
+                latest.each do |conn, parts|
                   begin
                     conn.write_message(parts)
                     written << conn
                   rescue *ZMTP::CONNECTION_LOST
-                    # connection dead — will be cleaned up
+                  end
+                end
+              else
+                batch.each do |parts|
+                  topic = parts.first || "".b
+                  @connections.each do |conn|
+                    next unless subscribed?(conn, topic)
+                    begin
+                      conn.write_message(parts)
+                      written << conn
+                    rescue *ZMTP::CONNECTION_LOST
+                    end
                   end
                 end
               end
@@ -74,14 +105,13 @@ module OMQ
               written.each do |conn|
                 conn.flush
               rescue *ZMTP::CONNECTION_LOST
-                # connection dead — will be cleaned up
               end
             end
           end
         end
 
         def start_subscription_listener(conn)
-          @tasks << Reactor.spawn_pump do
+          @tasks << Reactor.spawn_pump(annotation: "recv pump") do
             loop do
               frame = conn.read_frame
               next unless frame.command?

data/lib/omq/zmtp/routing/pair.rb

@@ -17,7 +17,8 @@ module OMQ
           @connection = nil
           @recv_queue = Async::LimitedQueue.new(engine.options.recv_hwm)
           @send_queue = Async::LimitedQueue.new(engine.options.send_hwm)
-          @tasks      = []
+          @tasks          = []
+          @send_pump_idle = true
         end
 
         # @return [Async::LimitedQueue]
@@ -59,10 +60,15 @@ module OMQ
 
         private
 
+        def send_pump_idle? = @send_pump_idle
+
+
         def start_send_pump(conn)
-          @send_pump = Reactor.spawn_pump do
+          @send_pump = @engine.spawn_pump_task(annotation: "send pump") do
             loop do
+              @send_pump_idle = true
               batch = [@send_queue.dequeue]
+              @send_pump_idle = false
               Routing.drain_send_queue(@send_queue, batch)
               batch.each { |parts| conn.write_message(parts) }
               conn.flush

data/lib/omq/zmtp/routing/peer.rb

@@ -21,6 +21,7 @@ module OMQ
           @connections_by_routing_id  = {}
           @tasks                      = []
           @send_pump_started          = false
+          @send_pump_idle             = true
         end
 
         # @return [Async::LimitedQueue]
@@ -59,11 +60,16 @@ module OMQ
 
         private
 
+        def send_pump_idle? = @send_pump_idle
+
+
         def start_send_pump
           @send_pump_started = true
-          @tasks << Reactor.spawn_pump do
+          @tasks << @engine.spawn_pump_task(annotation: "send pump") do
             loop do
+              @send_pump_idle = true
               batch = [@send_queue.dequeue]
+              @send_pump_idle = false
               Routing.drain_send_queue(@send_queue, batch)
 
               written = Set.new