omq-cli 0.2.0 → 0.3.0

data/lib/omq/cli/expression_evaluator.rb

@@ -0,0 +1,156 @@
+# frozen_string_literal: true
+
+module OMQ
+  module CLI
+    # Compiles and evaluates a single Ruby expression string for use in
+    # --recv-eval / --send-eval. Handles BEGIN{}/END{} block extraction,
+    # proc compilation, and result normalisation.
+    #
+    # One instance per direction (send or recv).
+    #
+    class ExpressionEvaluator
+      attr_reader :begin_proc, :end_proc, :eval_proc
+
+      # Sentinel: eval proc returned the context object, meaning it already
+      # sent the reply itself.
+      SENT = Object.new.freeze
+
+
+      # @param src [String, nil]  the raw expression string (may include BEGIN{}/END{})
+      # @param format [Symbol]    the active format, used to normalise results
+      # @param fallback_proc [Proc, nil]  registered OMQ.outgoing/incoming handler;
+      #   used only when +src+ is nil (no inline expression)
+      #
+      def initialize(src, format:, fallback_proc: nil)
+        @format = format
+
+        if src
+          expr, begin_body, end_body = extract_blocks(src)
+          @begin_proc = eval("proc { #{begin_body} }") if begin_body # rubocop:disable Security/Eval
+          @end_proc   = eval("proc { #{end_body} }")   if end_body   # rubocop:disable Security/Eval
+          if expr && !expr.strip.empty?
+            @eval_proc = eval("proc { $_ = $F&.first; #{expr} }") # rubocop:disable Security/Eval
+          end
+        elsif fallback_proc
+          @eval_proc = proc { |msg|
+            $_ = msg&.first
+            fallback_proc.call(msg)
+          }
+        end
+      end
+
+
+      # Runs the eval proc against +parts+ using +context+ as self.
+      # Returns the normalised result Array, nil (filter/skip), or SENT.
+      #
+      def call(parts, context)
+        return parts unless @eval_proc
+
+        $F     = parts
+        result = context.instance_exec(parts, &@eval_proc)
+        return nil  if result.nil?
+        return SENT if result.equal?(context)
+        return [result] if @format == :marshal
+
+        case result
+        when Array
+          result
+        when String
+          [result]
+        else
+          [result.to_str]
+        end
+      rescue => e
+        $stderr.puts "omq: eval error: #{e.message} (#{e.class})"
+        exit 3
+      end
+
+
+      # Normalises an eval result to nil (skip) or an Array of strings.
+      # Used inside Ractor worker blocks where instance methods are unavailable.
+      #
+      def self.normalize_result(result)
+        case result
+        when nil
+          nil
+        when Array
+          result
+        when String
+          [result]
+        else
+          [result.to_s]
+        end
+      end
+
+
+      # Compiles begin/end/eval procs inside a Ractor from a raw expression
+      # string. Returns [begin_proc, end_proc, eval_proc], any may be nil.
+      #
+      # Must be called inside the Ractor block (Procs are not Ractor-shareable).
+      #
+      def self.compile_inside_ractor(src)
+        return [nil, nil, nil] unless src
+
+        extract = ->(expr, kw) {
+          s = expr.index(/#{kw}\s*\{/)
+          return [expr, nil] unless s
+          ci = expr.index("{", s)
+          depth = 1
+          j = ci + 1
+          while j < expr.length && depth > 0
+            depth += 1 if expr[j] == "{"
+            depth -= 1 if expr[j] == "}"
+            j += 1
+          end
+          [expr[0...s] + expr[j..], expr[(ci + 1)..(j - 2)]]
+        }
+
+        expr, begin_body = extract.(src, "BEGIN")
+        expr, end_body   = extract.(expr, "END")
+
+        begin_proc = eval("proc { #{begin_body} }") if begin_body # rubocop:disable Security/Eval
+        end_proc   = eval("proc { #{end_body} }")   if end_body   # rubocop:disable Security/Eval
+        eval_proc  = nil
+        if expr && !expr.strip.empty?
+          ractor_expr = expr.gsub(/\$F\b/, "__F")
+          eval_proc   = eval("proc { |__F| $_ = __F&.first; #{ractor_expr} }") # rubocop:disable Security/Eval
+        end
+
+        [begin_proc, end_proc, eval_proc]
+      end
+
+
+      private
+
+
+      def extract_blocks(expr)
+        expr, begin_body = extract_block(expr, "BEGIN")
+        expr, end_body   = extract_block(expr, "END")
+        [expr, begin_body, end_body]
+      end
+
+
+      def extract_block(expr, keyword)
+        start = expr.index(/#{keyword}\s*\{/)
+        return [expr, nil] unless start
+
+        i     = expr.index("{", start)
+        depth = 1
+        j     = i + 1
+        while j < expr.length && depth > 0
+          case expr[j]
+          when "{"
+            depth += 1
+          when "}"
+            depth -= 1
+          end
+          j += 1
+        end
+
+        body    = expr[(i + 1)..(j - 2)]
+        trimmed = expr[0...start] + expr[j..]
+        [trimmed, body]
+      end
+    end
+  end
+end

data/lib/omq/cli/formatter.rb

@@ -5,12 +5,18 @@ module OMQ
     # Handles encoding/decoding messages in the configured format,
     # plus optional Zstandard compression.
     class Formatter
+      # @param format [Symbol] wire format (:ascii, :quoted, :raw, :jsonl, :msgpack, :marshal)
+      # @param compress [Boolean] whether to apply Zstandard compression per frame
       def initialize(format, compress: false)
         @format   = format
         @compress = compress
       end
 
 
+      # Encodes message parts into a printable string for output.
+      #
+      # @param parts [Array<String>] message frames
+      # @return [String] formatted output line
       def encode(parts)
         case @format
         when :ascii
@@ -19,7 +25,7 @@ module OMQ
           parts.map { |p| p.b.dump[1..-2] }.join("\t") + "\n"
         when :raw
           parts.each_with_index.map do |p, i|
-            ZMTP::Codec::Frame.new(p.to_s, more: i < parts.size - 1).to_wire
+            Protocol::ZMTP::Codec::Frame.new(p.to_s, more: i < parts.size - 1).to_wire
           end.join
         when :jsonl
           JSON.generate(parts) + "\n"
@@ -31,6 +37,10 @@ module OMQ
       end
 
 
+      # Decodes a formatted input line into message parts.
+      #
+      # @param line [String] input line (newline-terminated)
+      # @return [Array<String>] message frames
       def decode(line)
         case @format
         when :ascii, :marshal
@@ -47,6 +57,10 @@ module OMQ
       end
 
 
+      # Decodes one Marshal object from the given IO stream.
+      #
+      # @param io [IO] input stream
+      # @return [Object, nil] deserialized object, or nil on EOF
       def decode_marshal(io)
         Marshal.load(io)
       rescue EOFError, TypeError
@@ -54,6 +68,10 @@ module OMQ
       end
 
 
+      # Decodes one MessagePack object from the given IO stream.
+      #
+      # @param io [IO] input stream
+      # @return [Object, nil] deserialized object, or nil on EOF
       def decode_msgpack(io)
         @msgpack_unpacker ||= MessagePack::Unpacker.new(io)
         @msgpack_unpacker.read
@@ -62,11 +80,19 @@ module OMQ
       end
 
 
+      # Compresses each frame with Zstandard if compression is enabled.
+      #
+      # @param parts [Array<String>] message frames
+      # @return [Array<String>] optionally compressed frames
       def compress(parts)
         @compress ? parts.map { |p| Zstd.compress(p) } : parts
       end
 
 
+      # Decompresses each frame with Zstandard if compression is enabled.
+      #
+      # @param parts [Array<String>] possibly compressed message frames
+      # @return [Array<String>] decompressed frames
       def decompress(parts)
         @compress ? parts.map { |p| Zstd.decompress(p) } : parts
       end

data/lib/omq/cli/pair.rb

@@ -2,12 +2,20 @@
 
 module OMQ
   module CLI
+    # Runner for PAIR, DEALER, and CHANNEL sockets (bidirectional messaging).
     class PairRunner < BaseRunner
       private
 
 
       def run_loop(task)
-        receiver = task.async do
+        receiver = recv_async(task)
+        sender   = task.async { run_send_logic }
+        wait_for_loops(receiver, sender)
+      end
+
+
+      def recv_async(task)
+        task.async do
           n = config.count
           i = 0
           loop do
@@ -19,12 +27,6 @@ module OMQ
             break if n && n > 0 && i >= n
           end
         end
-
-        sender = task.async do
-          run_send_logic
-        end
-
-        wait_for_loops(receiver, sender)
       end
     end
   end

data/lib/omq/cli/parallel_recv_runner.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+module OMQ
+  module CLI
+    # Manages N OMQ::Ractor workers for parallel recv-eval.
+    #
+    # Each worker gets its own input socket connecting to the external
+    # endpoints; ZMQ distributes work naturally. Results are collected via
+    # an inproc PULL back to the main task for output.
+    #
+    class ParallelRecvRunner
+      # @param klass [Class] OMQ socket class for worker input sockets
+      # @param config [Config] frozen CLI configuration
+      # @param fmt [Formatter] message formatter
+      # @param output_fn [Method] callable for writing output
+      def initialize(klass, config, fmt, output_fn)
+        @klass     = klass
+        @config    = config
+        @fmt       = fmt
+        @output_fn = output_fn
+      end
+
+
+      # Spawns N Ractor workers, distributes incoming messages, and collects results.
+      #
+      # @param task [Async::Task] parent async task
+      # @return [void]
+      def run(task)
+        cfg       = @config
+        n_workers = cfg.parallel
+        inproc    = "inproc://omq-out-#{object_id}"
+
+        # Pack worker config into a shareable Hash passed via omq.data —
+        # Ruby 4.0 forbids Ractor blocks from closing over outer locals.
+        worker_data = ::Ractor.make_shareable({
+          recv_src:  cfg.recv_expr,
+          fmt_sym:   cfg.format,
+          fmt_compr: cfg.compress,
+        })
+
+        # Create N input sockets in the main Async context
+        input_socks = n_workers.times.map do
+          sock = SocketSetup.build(@klass, cfg)
+          SocketSetup.setup_subscriptions(sock, cfg)
+          cfg.connects.each { |url| sock.connect(url) }
+          sock
+        end
+
+        with_timeout(cfg.timeout) { input_socks.each { |s| s.peer_connected.wait } }
+
+        # Inproc collector: one bound PULL to receive all worker output
+        collector = OMQ::PULL.new(linger: cfg.linger)
+        collector.recv_timeout = cfg.timeout if cfg.timeout
+        collector.bind(inproc)
+
+        # N output sockets connecting to the collector
+        output_socks = n_workers.times.map do
+          s = OMQ::PUSH.new(linger: cfg.linger)
+          s.connect(inproc)
+          s
+        end
+
+        workers = n_workers.times.map do |i|
+          OMQ::Ractor.new(input_socks[i], output_socks[i], serialize: false, data: worker_data) do |omq|
+            pull_p, push_p = omq.sockets
+            d = omq.data
+
+            # Re-compile expression inside Ractor (Procs are not shareable)
+            begin_proc, end_proc, eval_proc =
+              OMQ::CLI::ExpressionEvaluator.compile_inside_ractor(d[:recv_src])
+
+            formatter = OMQ::CLI::Formatter.new(d[:fmt_sym], compress: d[:fmt_compr])
+            # Use a dedicated context object so @ivar expressions in BEGIN/END/eval
+            # work inside Ractors (self in a Ractor is shareable; Object.new is not).
+            _ctx = Object.new
+            _ctx.instance_exec(&begin_proc) if begin_proc
+
+            if eval_proc
+              loop do
+                parts = pull_p.receive
+                break if parts.nil?
+                parts = OMQ::CLI::ExpressionEvaluator.normalize_result(
+                  _ctx.instance_exec(formatter.decompress(parts), &eval_proc)
+                )
+                next if parts.nil?
+                push_p << parts unless parts.empty?
+              end
+            else
+              loop do
+                parts = pull_p.receive
+                break if parts.nil?
+                push_p << formatter.decompress(parts)
+              end
+            end
+
+            if end_proc
+              out = OMQ::CLI::ExpressionEvaluator.normalize_result(
+                _ctx.instance_exec(&end_proc)
+              )
+              push_p << out if out && !out.empty?
+            end
+          end
+        end
+
+
+        # Collect loop: drain inproc PULL → output
+        n_count = cfg.count
+        if n_count && n_count > 0
+          n_count.times do
+            parts = collector.receive
+            break if parts.nil?
+            @output_fn.call(parts)
+          end
+        else
+          loop do
+            parts = collector.receive
+            break if parts.nil?
+            @output_fn.call(parts)
+          end
+        end
+
+
+        # Inject nil into each worker's input port so it exits its loop
+        # without waiting for recv_timeout (workers don't self-terminate
+        # when the collect loop exits on count).
+        workers.each do |w|
+          w.close
+        rescue Ractor::RemoteError => e
+          $stderr.puts "omq: Ractor error: #{e.cause&.message || e.message}"
+        end
+      ensure
+        input_socks&.each(&:close)
+        output_socks&.each(&:close)
+        collector&.close
+      end
+
+
+      private
+
+
+      def with_timeout(seconds)
+        if seconds
+          Async::Task.current.with_timeout(seconds) { yield }
+        else
+          yield
+        end
+      end
+    end
+  end
+end