omq 0.5.0 → 0.6.0

data/lib/omq/zmtp/connection.rb

@@ -13,14 +13,17 @@ module OMQ
       #
       attr_reader :peer_socket_type
 
+
       # @return [String] peer's identity (from READY handshake)
       #
       attr_reader :peer_identity
 
+
       # @return [Object] transport IO (#read, #write, #close)
       #
       attr_reader :io
 
+
       # @param io [#read, #write, #close] transport IO
       # @param socket_type [String] our socket type name (e.g. "REQ")
       # @param identity [String] our identity
@@ -51,6 +54,7 @@ module OMQ
         @max_message_size    = max_message_size
       end
 
+
       # Performs the full ZMTP handshake via the configured mechanism.
       #
       # @return [void]
@@ -77,25 +81,44 @@ module OMQ
         end
       end
 
-      # Sends a multi-frame message.
+
+      # Sends a multi-frame message (write + flush).
       #
       # @param parts [Array<String>] message frames
       # @return [void]
       #
       def send_message(parts)
         @mutex.synchronize do
-          parts.each_with_index do |part, i|
-            more = i < parts.size - 1
-            if @mechanism.encrypted?
-              @io.write(@mechanism.encrypt(part.b, more: more))
-            else
-              @io.write(Codec::Frame.new(part, more: more).to_wire)
-            end
-          end
+          write_frames(parts)
+          @io.flush
+        end
+      end
+
+
+      # Writes a multi-frame message to the buffer without flushing.
+      # Call {#flush} after batching writes.
+      #
+      # @param parts [Array<String>] message frames
+      # @return [void]
+      #
+      def write_message(parts)
+        @mutex.synchronize do
+          write_frames(parts)
+        end
+      end
+
+
+      # Flushes the write buffer to the underlying IO.
+      #
+      # @return [void]
+      #
+      def flush
+        @mutex.synchronize do
           @io.flush
         end
       end
 
+
       # Receives a multi-frame message.
       # PING/PONG commands are handled automatically by #read_frame.
       #
@@ -110,12 +133,13 @@ module OMQ
             yield frame if block_given?
             next
           end
-          frames << frame.body
+          frames << frame.body.freeze
           break unless frame.more?
         end
-        frames
+        frames.freeze
       end
 
+
       # Starts the heartbeat sender task. Call after handshake.
       #
       # @return [#stop, nil] the heartbeat task, or nil if disabled
@@ -123,7 +147,7 @@ module OMQ
       def start_heartbeat
         return nil unless @heartbeat_interval
         @last_received_at = monotonic_now
-        @heartbeat_task = Reactor.spawn_pump do
+        @heartbeat_task = Reactor.spawn_pump(annotation: "heartbeat") do
           loop do
             sleep @heartbeat_interval
             # Send PING with TTL
@@ -145,6 +169,7 @@ module OMQ
         end
       end
 
+
       # Sends a command.
       #
       # @param command [Codec::Command]
@@ -161,6 +186,7 @@ module OMQ
         end
       end
 
+
       # Reads one frame from the wire. Handles PING/PONG automatically.
       # When using an encrypted mechanism, MESSAGE commands are decrypted
       # back to ZMTP frames transparently.
@@ -200,6 +226,7 @@ module OMQ
         end
       end
 
+
       # Closes the connection.
       #
       # @return [void]
@@ -211,16 +238,36 @@ module OMQ
         # already closed
       end
 
+
       private
 
+
+      # Writes message parts as ZMTP frames, encrypting if needed.
+      #
+      # @param parts [Array<String>] message frames
+      #
+      def write_frames(parts)
+        parts.each_with_index do |part, i|
+          more = i < parts.size - 1
+          if @mechanism.encrypted?
+            @io.write(@mechanism.encrypt(part.b, more: more))
+          else
+            @io.write(Codec::Frame.new(part, more: more).to_wire)
+          end
+        end
+      end
+
+
       def touch_heartbeat
         @last_received_at = monotonic_now if @heartbeat_interval
       end
 
+
       def monotonic_now
         Async::Clock.now
       end
 
+
       # Sends one frame to the wire.
       #
       # @param frame [Codec::Frame]

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

@@ -12,11 +12,9 @@ module OMQ
     # tasks — mirroring libzmq's IO thread architecture.
     #
     module Reactor
-      @work_queue = Thread::Queue.new
+      @work_queue = Async::Queue.new
       @thread     = nil
       @mutex      = Mutex.new
-      @wake_r     = nil
-      @wake_w     = nil
 
       class << self
         # Spawns a pump task (recv loop, send loop, accept loop).
@@ -26,14 +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])
-            @wake_w.write_nonblock(".") rescue nil
+            @work_queue.push([:spawn, block, handle, annotation])
             handle
           end
         end
@@ -52,7 +49,6 @@ module OMQ
             result_queue = Thread::Queue.new
             ensure_started
             @work_queue.push([:run, block, result_queue])
-            @wake_w.write_nonblock(".") rescue nil
             status, value = result_queue.pop
             raise value if status == :error
             value
@@ -66,9 +62,8 @@ module OMQ
         def ensure_started
           @mutex.synchronize do
             return if @thread&.alive?
-            @wake_r, @wake_w = IO.pipe
             ready = Thread::Queue.new
-            @thread = Thread.new { run_reactor(ready, @wake_r) }
+            @thread = Thread.new { run_reactor(ready) }
             @thread.name = "omq-io"
             ready.pop
           end
@@ -80,42 +75,35 @@ module OMQ
         #
         def stop!
           @work_queue.push([:stop])
-          @wake_w&.write_nonblock(".") rescue nil
           @thread&.join(2)
           @thread = nil
-          @wake_r&.close rescue nil
-          @wake_w&.close rescue nil
-          @wake_r = nil
-          @wake_w = nil
         end
 
         private
 
-        def run_reactor(ready, wake_r)
+        # 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)
             loop do
-              # Wait for wakeup signal (non-blocking for Async scheduler)
-              wake_r.wait_readable
-              wake_r.read_nonblock(256) rescue nil
-
-              # Drain all pending work items
-              while (item = @work_queue.pop(true) rescue nil)
-                case item[0]
-                when :spawn
-                  _, block, handle = item
-                  async_task = task.async(transient: true, &block)
-                  handle.task = async_task
-                when :run
-                  _, block, result_queue = item
-                  task.async do
-                    result_queue.push([:ok, block.call])
-                  rescue => e
-                    result_queue.push([:error, e])
-                  end
-                when :stop
-                  return
+              item = @work_queue.dequeue
+              case item[0]
+              when :spawn
+                _, block, handle, annotation = item
+                async_task = task.async(transient: true, annotation: annotation, &block)
+                handle.task = async_task
+              when :run
+                _, block, result_queue = item
+                task.async do
+                  result_queue.push([:ok, block.call])
+                rescue => e
+                  result_queue.push([:error, e])
                 end
+              when :stop
+                return
               end
             end
           end
@@ -129,6 +117,7 @@ module OMQ
         #
         attr_accessor :task
 
+
         # Stops the pump task.
         #
         # @return [void]