omq-cli 0.2.0 → 0.3.1

data/lib/omq/cli/base_runner.rb

@@ -2,10 +2,16 @@
 
 module OMQ
   module CLI
+    # Template runner base class for all socket-type CLI runners.
+    # Subclasses override {#run_loop} to implement socket-specific behaviour.
     class BaseRunner
+      # @return [Config] frozen CLI configuration
+      # @return [Object] the OMQ socket instance
       attr_reader :config, :sock
 
 
+      # @param config [Config] frozen CLI configuration
+      # @param socket_class [Class] OMQ socket class to instantiate (e.g. OMQ::PUSH)
       def initialize(config, socket_class)
         @config = config
         @klass  = socket_class
@@ -13,26 +19,18 @@ module OMQ
       end
 
 
+      # Runs the full lifecycle: socket setup, peer wait, BEGIN/END blocks, and the main loop.
+      #
+      # @param task [Async::Task] the parent async task
+      # @return [void]
       def call(task)
-        @sock = create_socket
-        attach_endpoints
-        setup_curve
-        setup_subscriptions
-        compile_expr
-
-        if config.transient
-          start_disconnect_monitor(task)
-          Async::Task.current.yield  # let monitor start waiting
-        end
-
+        setup_socket
+        maybe_start_transient_monitor(task)
         sleep(config.delay) if config.delay && config.recv_only?
         wait_for_peer if needs_peer_wait?
-
-        @sock.instance_exec(&@send_begin_proc) if @send_begin_proc
-        @sock.instance_exec(&@recv_begin_proc) if @recv_begin_proc
+        run_begin_blocks
         run_loop(task)
-        @sock.instance_exec(&@send_end_proc) if @send_end_proc
-        @sock.instance_exec(&@recv_end_proc) if @recv_end_proc
+        run_end_blocks
       ensure
         @sock&.close
       end
@@ -46,34 +44,59 @@ module OMQ
         raise NotImplementedError
       end
 
+
       # ── Socket creation ─────────────────────────────────────────────
 
 
+      def setup_socket
+        @sock = create_socket
+        attach_endpoints unless config.parallel
+        setup_curve
+        setup_subscriptions
+        compile_expr
+      end
+
+
       def create_socket
-        sock_opts              = { linger: config.linger }
-        sock_opts[:conflate]   = true if config.conflate && %w[pub radio].include?(config.type_name)
-        sock                   = @klass.new(**sock_opts)
-        sock.recv_timeout      = config.timeout if config.timeout
-        sock.send_timeout      = config.timeout if config.timeout
-        sock.reconnect_interval  = config.reconnect_ivl if config.reconnect_ivl
-        sock.heartbeat_interval  = config.heartbeat_ivl if config.heartbeat_ivl
-        sock.identity            = config.identity if config.identity
-        sock.router_mandatory  = true if config.type_name == "router"
-        sock
+        SocketSetup.build(@klass, config)
       end
 
 
       def attach_endpoints
-        config.binds.each do |url|
-          @sock.bind(url)
-          log "Bound to #{@sock.last_endpoint}"
-        end
-        config.connects.each do |url|
-          @sock.connect(url)
-          log "Connecting to #{url}"
-        end
+        SocketSetup.attach(@sock, config, verbose: config.verbose)
+      end
+
+
+      # ── Transient disconnect monitor ────────────────────────────────
+
+
+      def maybe_start_transient_monitor(task)
+        return unless config.transient
+        @transient_monitor = TransientMonitor.new(@sock, config, task, method(:log))
+        Async::Task.current.yield  # let monitor start waiting
       end
 
+
+      def transient_ready!
+        @transient_monitor&.ready!
+      end
+
+
+      # ── BEGIN / END blocks ──────────────────────────────────────────
+
+
+      def run_begin_blocks
+        @sock.instance_exec(&@send_begin_proc) if @send_begin_proc
+        @sock.instance_exec(&@recv_begin_proc) if @recv_begin_proc
+      end
+
+
+      def run_end_blocks
+        @sock.instance_exec(&@send_end_proc) if @send_end_proc
+        @sock.instance_exec(&@recv_end_proc) if @recv_end_proc
+      end
+
+
       # ── Peer wait with grace period ─────────────────────────────────
 
 
@@ -86,46 +109,29 @@ module OMQ
         with_timeout(config.timeout) do
           @sock.peer_connected.wait
           log "Peer connected"
-          if %w[pub xpub].include?(config.type_name)
-            @sock.subscriber_joined.wait
-            log "Subscriber joined"
-          end
-
-          # Grace period: when multiple peers may be connecting (bind or
-          # multiple connect URLs), wait one reconnect interval so
-          # latecomers finish their handshake before we start sending.
-          if config.binds.any? || config.connects.size > 1
-            ri = @sock.options.reconnect_interval
-            sleep(ri.is_a?(Range) ? ri.begin : ri)
-          end
+          wait_for_subscriber
+          apply_grace_period
         end
       end
 
-      # ── Transient disconnect monitor ────────────────────────────────
 
-
-      def start_disconnect_monitor(task)
-        @transient_barrier = Async::Promise.new
-        task.async do
-          @transient_barrier.wait
-          @sock.all_peers_gone.wait unless @sock.connection_count == 0
-          log "All peers disconnected, exiting"
-          @sock.reconnect_enabled = false
-          if config.send_only?
-            task.stop
-          else
-            @sock.close_read
-          end
-        end
+      def wait_for_subscriber
+        return unless %w[pub xpub].include?(config.type_name)
+        @sock.subscriber_joined.wait
+        log "Subscriber joined"
       end
 
 
-      def transient_ready!
-        if config.transient && !@transient_barrier.resolved?
-          @transient_barrier.resolve(true)
-        end
+      # Grace period: when multiple peers may be connecting (bind or
+      # multiple connect URLs), wait one reconnect interval so
+      # latecomers finish their handshake before we start sending.
+      def apply_grace_period
+        return unless config.binds.any? || config.connects.size > 1
+        ri = @sock.options.reconnect_interval
+        sleep(ri.is_a?(Range) ? ri.begin : ri)
       end
 
+
       # ── Timeout helper ──────────────────────────────────────────────
 
 
@@ -137,50 +143,22 @@ module OMQ
         end
       end
 
+
       # ── Socket setup ────────────────────────────────────────────────
 
 
       def setup_subscriptions
-        case config.type_name
-        when "sub"
-          prefixes = config.subscribes.empty? ? [""] : config.subscribes
-          prefixes.each { |p| @sock.subscribe(p) }
-        when "dish"
-          config.joins.each { |g| @sock.join(g) }
-        end
+        SocketSetup.setup_subscriptions(@sock, config)
+      end
+
+
+      def setup_subscriptions_on(sock)
+        SocketSetup.setup_subscriptions(sock, config)
       end
 
 
       def setup_curve
-        server_key_z85 = config.curve_server_key || ENV["OMQ_SERVER_KEY"]
-        server_mode    = config.curve_server || (ENV["OMQ_SERVER_PUBLIC"] && ENV["OMQ_SERVER_SECRET"])
-
-        return unless server_key_z85 || server_mode
-
-        crypto = CLI.load_curve_crypto(config.curve_crypto || ENV["OMQ_CURVE_CRYPTO"], verbose: config.verbose)
-        require "protocol/zmtp/mechanism/curve"
-
-        if server_key_z85
-          server_key = Protocol::ZMTP::Z85.decode(server_key_z85)
-          client_key = crypto::PrivateKey.generate
-          @sock.mechanism = Protocol::ZMTP::Mechanism::Curve.client(
-            client_key.public_key.to_s, client_key.to_s,
-            server_key: server_key, crypto: crypto
-          )
-        elsif server_mode
-          if ENV["OMQ_SERVER_PUBLIC"] && ENV["OMQ_SERVER_SECRET"]
-            server_pub = Protocol::ZMTP::Z85.decode(ENV["OMQ_SERVER_PUBLIC"])
-            server_sec = Protocol::ZMTP::Z85.decode(ENV["OMQ_SERVER_SECRET"])
-          else
-            key        = crypto::PrivateKey.generate
-            server_pub = key.public_key.to_s
-            server_sec = key.to_s
-          end
-          @sock.mechanism = Protocol::ZMTP::Mechanism::Curve.server(
-            server_pub, server_sec, crypto: crypto
-          )
-          $stderr.puts "OMQ_SERVER_KEY='#{Protocol::ZMTP::Z85.encode(server_pub)}'"
-        end
+        SocketSetup.setup_curve(@sock, config)
       end
 
 
@@ -189,28 +167,14 @@ module OMQ
 
       def run_send_logic
         n = config.count
-        i = 0
         sleep(config.delay) if config.delay
         if config.interval
-          i += send_tick
-          unless @send_tick_eof || (n && n > 0 && i >= n)
-            Async::Loop.quantized(interval: config.interval) do
-              i += send_tick
-              break if @send_tick_eof || (n && n > 0 && i >= n)
-            end
-          end
+          run_interval_send(n)
         elsif config.data || config.file
           parts = eval_send_expr(read_next)
           send_msg(parts) if parts
         elsif stdin_ready?
-          loop do
-            parts = read_next
-            break unless parts
-            parts = eval_send_expr(parts)
-            send_msg(parts) if parts
-            i += 1
-            break if n && n > 0 && i >= n
-          end
+          run_stdin_send(n)
         elsif @send_eval_proc
           parts = eval_send_expr(nil)
           send_msg(parts) if parts
@@ -218,9 +182,38 @@ module OMQ
       end
 
 
+      def run_interval_send(n)
+        i = send_tick
+        return if @send_tick_eof || (n && n > 0 && i >= n)
+        Async::Loop.quantized(interval: config.interval) do
+          i += send_tick
+          break if @send_tick_eof || (n && n > 0 && i >= n)
+        end
+      end
+
+
+      def run_stdin_send(n)
+        i = 0
+        loop do
+          parts = read_next
+          break unless parts
+          parts = eval_send_expr(parts)
+          send_msg(parts) if parts
+          i += 1
+          break if n && n > 0 && i >= n
+        end
+      end
+
+
       def send_tick
         raw = read_next_or_nil
-        if raw.nil? && !@send_eval_proc
+        if raw.nil?
+          if @send_eval_proc && !@stdin_ready
+            # Pure generator mode: no stdin, eval produces output from nothing.
+            parts = eval_send_expr(nil)
+            send_msg(parts) if parts
+            return 1
+          end
           @send_tick_eof = true
           return 0
         end
@@ -233,14 +226,52 @@ module OMQ
       def run_recv_logic
         n = config.count
         i = 0
-        loop do
-          parts = recv_msg
-          break if parts.nil?
-          parts = eval_recv_expr(parts)
-          output(parts)
-          i += 1
-          break if n && n > 0 && i >= n
+        if config.interval
+          run_interval_recv(n)
+        else
+          loop do
+            parts = recv_msg
+            break if parts.nil?
+            parts = eval_recv_expr(parts)
+            output(parts)
+            i += 1
+            break if n && n > 0 && i >= n
+          end
+        end
+      end
+
+
+      def run_interval_recv(n)
+        i = recv_tick
+        return if i == 0
+        return if n && n > 0 && i >= n
+        Async::Loop.quantized(interval: config.interval) do
+          i += recv_tick
+          break if @recv_tick_eof || (n && n > 0 && i >= n)
+        end
+      end
+
+
+      def recv_tick
+        parts = recv_msg
+        if parts.nil?
+          @recv_tick_eof = true
+          return 0
         end
+        parts = eval_recv_expr(parts)
+        output(parts)
+        1
+      end
+
+
+      # Parallel recv-eval: delegates to ParallelRecvRunner.
+      #
+      def run_parallel_recv(task)
+        # @sock was created by call() before run_loop; close it now so it doesn't
+        # steal messages from the N worker sockets ParallelRecvRunner creates.
+        @sock&.close
+        @sock = nil
+        ParallelRecvRunner.new(@klass, config, @fmt, method(:output)).run(task)
       end
 
 
@@ -257,6 +288,7 @@ module OMQ
         end
       end
 
+
       # ── Message I/O ─────────────────────────────────────────────────
 
 
@@ -286,23 +318,32 @@ module OMQ
 
 
       def read_next
+        config.data || config.file ? read_inline_data : read_stdin_input
+      end
+
+
+      def read_inline_data
         if config.data
           @fmt.decode(config.data + "\n")
-        elsif config.file
+        else
           @file_data ||= (config.file == "-" ? $stdin.read : File.read(config.file)).chomp
           @fmt.decode(@file_data + "\n")
-        elsif config.format == :msgpack
+        end
+      end
+
+
+      def read_stdin_input
+        case config.format
+        when :msgpack
           @fmt.decode_msgpack($stdin)
-        elsif config.format == :marshal
+        when :marshal
           @fmt.decode_marshal($stdin)
-        elsif config.format == :raw
+        when :raw
           data = $stdin.read
-          return nil if data.nil? || data.empty?
-          [data]
+          data.nil? || data.empty? ? nil : [data]
         else
           line = $stdin.gets
-          return nil if line.nil?
-          @fmt.decode(line)
+          line.nil? ? nil : @fmt.decode(line)
         end
       end
 
@@ -320,10 +361,10 @@ module OMQ
       def read_next_or_nil
         if config.data || config.file
           read_next
-        elsif @send_eval_proc
-          nil
+        elsif stdin_ready?
+          read_stdin_input
         else
-          read_next
+          nil
         end
       end
 
@@ -334,116 +375,50 @@ module OMQ
         $stdout.flush
       end
 
-      # ── Routing helpers ─────────────────────────────────────────────
-
-
-      def display_routing_id(id)
-        if id.bytes.all? { |b| b >= 0x20 && b <= 0x7E }
-          id
-        else
-          "0x#{id.unpack1("H*")}"
-        end
-      end
-
-
-      def resolve_target(target)
-        if target.start_with?("0x")
-          [target[2..].delete(" ")].pack("H*")
-        else
-          target
-        end
-      end
 
       # ── Eval ────────────────────────────────────────────────────────
 
 
       def compile_expr
-        compile_one_expr(:send, config.send_expr)
-        compile_one_expr(:recv, config.recv_expr)
-        @send_eval_proc ||= wrap_registered_proc(OMQ.outgoing_proc)
-        @recv_eval_proc ||= wrap_registered_proc(OMQ.incoming_proc)
-      end
-
-
-      def wrap_registered_proc(block)
-        return unless block
-        proc do |msg|
-          $_ = msg&.first
-          block.call(msg)
-        end
+        @send_evaluator = compile_evaluator(config.send_expr, fallback: OMQ.outgoing_proc)
+        @recv_evaluator = compile_evaluator(config.recv_expr, fallback: OMQ.incoming_proc)
+        assign_send_aliases
+        assign_recv_aliases
       end
 
 
-      def compile_one_expr(direction, src)
-        return unless src
-        expr, begin_body, end_body = extract_blocks(src)
-        instance_variable_set(:"@#{direction}_begin_proc", eval("proc { #{begin_body} }")) if begin_body
-        instance_variable_set(:"@#{direction}_end_proc",   eval("proc { #{end_body} }"))   if end_body
-        if expr && !expr.strip.empty?
-          instance_variable_set(:"@#{direction}_eval_proc", eval("proc { $_ = $F&.first; #{expr} }"))
-        end
+      def compile_evaluator(src, fallback:)
+        ExpressionEvaluator.new(src, format: config.format, fallback_proc: fallback)
       end
 
 
-      def extract_blocks(expr)
-        begin_body = end_body = nil
-        expr, begin_body = extract_block(expr, "BEGIN")
-        expr, end_body   = extract_block(expr, "END")
-        [expr, begin_body, end_body]
+      def assign_send_aliases
+        # Keep ivar aliases — subclasses check these directly
+        @send_begin_proc = @send_evaluator.begin_proc
+        @send_eval_proc  = @send_evaluator.eval_proc
+        @send_end_proc   = @send_evaluator.end_proc
       end
 
 
-      def extract_block(expr, keyword)
-        start = expr.index(/#{keyword}\s*\{/)
-        return [expr, nil] unless start
-
-        # Find the opening brace
-        i = expr.index("{", start)
-        depth = 1
-        j     = i + 1
-        while j < expr.length && depth > 0
-          case expr[j]
-          when "{" then depth += 1
-          when "}" then depth -= 1
-          end
-          j += 1
-        end
-
-        body    = expr[(i + 1)..(j - 2)]
-        trimmed = expr[0...start] + expr[j..]
-        [trimmed, body]
+      def assign_recv_aliases
+        @recv_begin_proc = @recv_evaluator.begin_proc
+        @recv_eval_proc  = @recv_evaluator.eval_proc
+        @recv_end_proc   = @recv_evaluator.end_proc
       end
 
 
-      SENT = Object.new.freeze # sentinel: eval already sent the reply
-
       def eval_send_expr(parts)
-        return parts unless @send_eval_proc
-        run_eval(@send_eval_proc, parts)
+        @send_evaluator.call(parts, @sock)
       end
 
 
       def eval_recv_expr(parts)
-        return parts unless @recv_eval_proc
-        run_eval(@recv_eval_proc, parts)
+        @recv_evaluator.call(parts, @sock)
       end
 
 
-      def run_eval(eval_proc, parts)
-        $F = parts
-        result = @sock.instance_exec(parts, &eval_proc)
-        return nil if result.nil?
-        return SENT if result.equal?(@sock)
-        return [result] if config.format == :marshal
-        case result
-        when Array  then result
-        when String then [result]
-        else             [result.to_str]
-        end
-      rescue => e
-        $stderr.puts "omq: eval error: #{e.message} (#{e.class})"
-        exit 3
-      end
+      SENT = ExpressionEvaluator::SENT
+
 
       # ── Logging ─────────────────────────────────────────────────────