protobuf-nats 0.13.0.pre4 → 0.13.0.pre5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0a264426e2c04180f86c1006cd19441e51c36bd5e736cc55252c898e620c2428
-  data.tar.gz: 861ed68ab3c036afb2ce8fcfff3248f699b7cb4989df945c9113e4a7c3652223
+  metadata.gz: bac03c214ceddc05b8ad2ac60e1db25f4179156a78a437f6338f9b1c3404dc8c
+  data.tar.gz: d26d42acfd6b9804093a368ebb08ea41df22896fa67f1d2e79f767c6f9aa7f00
 SHA512:
-  metadata.gz: 255fa1e7942741ee80f9e0aebbef0f2eed943b7be84ee95d7517dddb9fbfe692bd4dcad7c70a2d05aae953988fd203b8b6bb6b5c7ee295f084f1801f30e019b1
-  data.tar.gz: 18aa2c4a2a166f769ae080d4cac177de891a26fcd77251d007b2e382927bf468bf7541a86c1a84b0c7f26217bef7104618dcd04faf5698046fff81894ba66f8e
+  metadata.gz: b8992b0220a24d05c0c1f0e67f816ea06865817a8ba0d18dcab0ba2ac36a4b561d69a144e97a7ccf3362bd33f6c4ac9b5458777aec85d2846720f71645f6952f
+  data.tar.gz: ef9d92af39b7a59c6f7c14a47ccaada1807fcca53e3f54cbf044ff264d73b698aeaad31d51e07ad731ac865aacfb901ef498a7d63cc678d91f615ab2859e7662

data/bench/concurrency_bench.rb

@@ -0,0 +1,219 @@
+# Concurrency microbenchmarks for the protobuf-nats hot paths.
+#
+# Drives the REAL ResponseMuxer / Client subscription cache / ThreadPool so the
+# same file measures both the baseline and the patched implementations, on both
+# CRuby and JRuby. No NATS server required (NATS is faked for the muxer bench).
+#
+# Usage:
+#   bundle exec ruby -Ilib bench/concurrency_bench.rb
+#   BENCH_THREADS=1,4,8,16 BENCH_DURATION=4 BENCH_WARMUP=2 ruby -Ilib bench/concurrency_bench.rb
+#
+# Each cell runs for BENCH_DURATION seconds (after BENCH_WARMUP seconds of
+# untimed warmup) and reports aggregate ops/sec across all threads plus an
+# error count.
+
+require "bundler/setup"
+require "protobuf/nats"
+require "nats/client" # for NATS::Msg / NATS::Subscription
+
+::Protobuf::Logging.logger = ::Logger.new(nil)
+
+DURATION = Float(ENV.fetch("BENCH_DURATION", "4"))
+WARMUP   = Float(ENV.fetch("BENCH_WARMUP", "2"))
+THREADS  = ENV.fetch("BENCH_THREADS", "1,4,8,16").split(",").map(&:to_i)
+PAYLOAD  = ("x" * 64).freeze
+
+def mono
+  ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+end
+
+# Run `block` (which returns [ops, errors]) on `threads` workers. Counters are
+# only accumulated during the measured window (after an untimed warmup), using
+# atomics so the stop/measure flags are visible across threads on JRuby too.
+# Returns { ops_per_sec:, errors: }.
+def run_cell(threads, seconds, warmup)
+  go      = ::Concurrent::AtomicBoolean.new(false)
+  measure = ::Concurrent::AtomicBoolean.new(false)
+  stop    = ::Concurrent::AtomicBoolean.new(false)
+  totals  = ::Array.new(threads)
+
+  workers = threads.times.map do |i|
+    ::Thread.new do
+      ops = 0
+      errs = 0
+      sleep 0.0005 until go.true?
+      until stop.true?
+        o, e = yield
+        if measure.true?
+          ops += o
+          errs += e
+        end
+      end
+      totals[i] = [ops, errs]
+    end
+  end
+
+  go.make_true
+  sleep warmup
+  t0 = mono
+  measure.make_true
+  sleep seconds
+  stop.make_true
+  elapsed = mono - t0
+  workers.each(&:join)
+
+  ops = totals.sum { |t| t ? t[0] : 0 }
+  errs = totals.sum { |t| t ? t[1] : 0 }
+  { ops_per_sec: ops / elapsed, errors: errs }
+end
+
+def print_table(title, rows)
+  puts
+  puts title
+  puts "  threads |        ops/sec |  errors"
+  puts "  --------+----------------+--------"
+  rows.each do |t, r|
+    printf("  %7d | %14.0f | %7d\n", t, r[:ops_per_sec], r[:errors])
+  end
+end
+
+# --------------------------------------------------------------------------
+# Fake NATS connection for the muxer round-trip benchmark.
+# new_inbox + subscribe("<prefix>.*") feed the muxer; publish echoes a reply
+# onto the muxer's response subscription queue, simulating the server response.
+# --------------------------------------------------------------------------
+class BenchNats
+  def initialize
+    @inbox = 0
+    @resp_queue = nil
+  end
+
+  def new_inbox
+    @inbox += 1
+    "_INBOX.bench.#{@inbox}"
+  end
+
+  def subscribe(_subject, *_args)
+    sub = ::NATS::Subscription.new
+    sub.pending_queue = ::SizedQueue.new(8192)
+    @resp_queue = sub.pending_queue # muxer's response subscription
+    sub
+  end
+
+  def publish(_subject, data, reply_to = nil)
+    return unless reply_to && @resp_queue
+    @resp_queue.push(::NATS::Msg.new(:subject => reply_to, :data => data))
+  end
+
+  def flush(*); end
+end
+
+# --------------------------------------------------------------------------
+# A. ResponseMuxer round-trip (exercises #1 map lock, #2 dispatcher, #4, #6)
+# --------------------------------------------------------------------------
+def bench_muxer
+  rows = THREADS.map do |t|
+    ::Protobuf::Nats.client_nats_connection = BenchNats.new
+    muxer = ::Protobuf::Nats::ResponseMuxer.new
+    muxer.start
+
+    result = run_cell(t, DURATION, WARMUP) do
+      ops = 0
+      errs = 0
+      begin
+        req = muxer.new_request
+        req.publish("rpc.bench", PAYLOAD)
+        msg = req.next_message(5)
+        ops += 1 if msg
+      rescue => _e
+        errs += 1
+      ensure
+        req.cleanup if req
+      end
+      [ops, errs]
+    end
+
+    muxer.stop
+    [t, result]
+  end
+  print_table("A. ResponseMuxer round-trip (new_request -> publish -> next_message -> cleanup)", rows)
+end
+
+# --------------------------------------------------------------------------
+# B. Subscription-key cache (exercises #3: nested-Hash ||= vs Concurrent::Map).
+# Drives the real Client#cached_subscription_key. Each measured iteration clears
+# the shared class cache and races all threads to refill it (the cold-start
+# write race that triggers ConcurrentModificationError on JRuby + plain Hash).
+# --------------------------------------------------------------------------
+def bench_subscription_cache
+  # Build named dummy service classes + methods used as cache keys.
+  svc_classes = 20.times.map do |i|
+    name = "BenchSvc#{i}"
+    ::Object.const_set(name, Class.new) unless ::Object.const_defined?(name)
+    ::Object.const_get(name)
+  end
+  methods = [:create, :read, :update, :delete, :list]
+  combos = svc_classes.product(methods)
+
+  cache = ::Protobuf::Nats::Client.subscription_key_cache
+
+  rows = THREADS.map do |t|
+    iteration_lock = ::Mutex.new
+    result = run_cell(t, DURATION, WARMUP) do
+      ops = 0
+      errs = 0
+      # Clear occasionally to keep the write path hot (cold-fill race).
+      iteration_lock.synchronize { cache.clear } if rand(combos.size) == 0
+      combos.each do |klass, meth|
+        begin
+          client = ::Protobuf::Nats::Client.allocate
+          client.instance_variable_set(:@options, { :service => klass, :method => meth })
+          client.cached_subscription_key
+          ops += 1
+        rescue => _e
+          errs += 1
+        end
+      end
+      [ops, errs]
+    end
+    [t, result]
+  end
+  print_table("B. Subscription-key cache fill+read race (Client#cached_subscription_key)", rows)
+end
+
+# --------------------------------------------------------------------------
+# C. ThreadPool throughput (exercises #5: mutex counter vs AtomicFixnum)
+# --------------------------------------------------------------------------
+def bench_thread_pool
+  workers = (ENV["BENCH_POOL_WORKERS"] || "8").to_i
+  noop = ->{}
+  rows = THREADS.map do |t|
+    pool = ::Protobuf::Nats::ThreadPool.new(workers, :max_queue => 100_000)
+    result = run_cell(t, DURATION, WARMUP) do
+      ops = 0
+      errs = 0
+      if pool.push(&noop)
+        ops += 1
+      else
+        errs += 1 # pool full (backpressure); not a real error, tracked separately
+      end
+      [ops, errs]
+    end
+    pool.shutdown
+    [t, result]
+  end
+  print_table("C. ThreadPool push throughput (errors column = pool-full/backpressure)", rows)
+end
+
+puts "=" * 72
+puts "protobuf-nats concurrency bench"
+puts "engine=#{RUBY_ENGINE} #{RUBY_VERSION}  duration=#{DURATION}s warmup=#{WARMUP}s threads=#{THREADS.inspect}"
+puts "processor_count=#{::Concurrent.processor_count}"
+puts "=" * 72
+
+bench_muxer
+bench_subscription_cache
+bench_thread_pool
+
+puts
+puts "done."

data/bench/e2e_bench.rb

@@ -0,0 +1,78 @@
+# End-to-end throughput benchmark against a real nats-server.
+#
+# Starts a real protobuf-nats Server (warehouse service) in a background thread
+# and drives it with N client threads for a fixed duration, reporting req/sec.
+# Exercises the full patched stack: client ResponseMuxer + dispatchers, server
+# SuperSubscriptionManager + ThreadPool.
+#
+# Requires a nats-server already running on 127.0.0.1:4222.
+#
+# Usage:
+#   E2E_DURATION=10 E2E_THREADS=16 bundle exec ruby -Ilib bench/e2e_bench.rb
+
+ENV["PB_CLIENT_TYPE"] = "protobuf/nats/client"
+ENV["PB_SERVER_TYPE"] = "protobuf/nats/runner"
+# Skip the server's slow-start ramp so the benchmark starts promptly.
+ENV["PB_NATS_SERVER_SUBSCRIPTIONS_PER_RPC_ENDPOINT"] ||= "1"
+ENV["PB_NATS_SERVER_SLOW_START_DELAY"] ||= "0"
+
+require "./examples/warehouse/app"
+require "concurrent"
+
+::Protobuf::Logging.logger = ::Logger.new(nil)
+
+DURATION       = Float(ENV.fetch("E2E_DURATION", "10"))
+WARMUP         = Float(ENV.fetch("E2E_WARMUP", "3"))
+THREADS        = Integer(ENV.fetch("E2E_THREADS", "16"))
+SERVER_THREADS = Integer(ENV.fetch("E2E_SERVER_THREADS", "50"))
+
+def mono
+  ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+end
+
+puts "=" * 72
+puts "protobuf-nats e2e bench  engine=#{RUBY_ENGINE} #{RUBY_VERSION}"
+puts "client_threads=#{THREADS} server_threads=#{SERVER_THREADS} duration=#{DURATION}s warmup=#{WARMUP}s"
+puts "dispatchers=#{ENV['PB_NATS_RESPONSE_MUXER_DISPATCHERS'] || '(auto)'}"
+puts "=" * 72
+
+server = ::Protobuf::Nats::Server.new(:threads => SERVER_THREADS)
+server_thread = ::Thread.new { server.run }
+
+# Give the server time to connect + subscribe.
+sleep 2
+
+count   = ::Concurrent::AtomicFixnum.new(0)
+errors  = ::Concurrent::AtomicFixnum.new(0)
+measure = ::Concurrent::AtomicBoolean.new(false)
+stop    = ::Concurrent::AtomicBoolean.new(false)
+
+workers = THREADS.times.map do
+  ::Thread.new do
+    until stop.true?
+      begin
+        req = ::Warehouse::Shipment.new(:guid => ::SecureRandom.uuid, :sleep_time_ms => 0)
+        ::Warehouse::ShipmentService.client.create(req)
+        count.increment if measure.true?
+      rescue => _e
+        errors.increment if measure.true?
+      end
+    end
+  end
+end
+
+sleep WARMUP
+t0 = mono
+measure.make_true
+sleep DURATION
+stop.make_true
+elapsed = mono - t0
+workers.each(&:join)
+
+rps = count.value / elapsed
+puts
+printf("req/sec = %.0f   (completed=%d errors=%d in %.2fs)\n", rps, count.value, errors.value, elapsed)
+
+server.stop
+server_thread.join(5)
+puts "done."

data/lib/protobuf/nats/client.rb

@@ -1,5 +1,6 @@
 require 'securerandom'
 require "connection_pool"
+require "concurrent"
 require "protobuf/nats"
 require "protobuf/rpc/connectors/base"
 require "monitor"
@@ -13,7 +14,14 @@ module Protobuf
 
       RESPONSE_MUXER = ::Protobuf::Nats::ResponseMuxer.new
 
-      @subscription_key_cache = {}
+      # On JRuby (true parallelism) concurrent writes to a plain nested Hash can
+      # raise ConcurrentModificationError / corrupt the map, so the cache must be
+      # a Concurrent::Map. On CRuby the GVL makes plain-Hash reads/writes atomic
+      # (a racing `||=` at worst recomputes an identical value), and a plain Hash
+      # is meaningfully faster than Concurrent::Map, so we keep the Hash there.
+      CONCURRENT_SUBSCRIPTION_CACHE = (::RUBY_ENGINE == "jruby")
+
+      @subscription_key_cache = CONCURRENT_SUBSCRIPTION_CACHE ? ::Concurrent::Map.new : {}
       @subscription_pool_lock = ::Mutex.new
 
       # Structure to hold subscription and inbox to use within pool
@@ -208,9 +216,15 @@ module Protobuf
         klass = @options[:service]
         method_name = @options[:method]
 
-        method_name_cache = self.class.subscription_key_cache[klass] ||= {}
-        method_name_cache[method_name] ||= begin
-          ::Protobuf::Nats.subscription_key(klass, method_name)
+        cache = self.class.subscription_key_cache
+        if CONCURRENT_SUBSCRIPTION_CACHE
+          method_name_cache = cache.compute_if_absent(klass) { ::Concurrent::Map.new }
+          method_name_cache.compute_if_absent(method_name) do
+            ::Protobuf::Nats.subscription_key(klass, method_name)
+          end
+        else
+          method_name_cache = cache[klass] ||= {}
+          method_name_cache[method_name] ||= ::Protobuf::Nats.subscription_key(klass, method_name)
         end
       end
 

data/lib/protobuf/nats/response_muxer.rb

@@ -3,7 +3,6 @@ require "connection_pool"
 require "protobuf/nats"
 require "protobuf/rpc/connectors/base"
 require "monitor"
-require "uuid7"
 require "protobuf/nats/uuidv7_helper"
 require "concurrent"
 require "concurrent/collection/timeout_queue"
@@ -16,11 +15,13 @@ module Protobuf
       TOKEN_TTL_SECONDS = 600 # 10 minutes
 
       def initialize
-        # Per-token response queues for lock-free message delivery
-        # Each token gets its own Queue, eliminating lock contention between different tokens
-        @resp_map = Hash.new { |h,k| h[k] = { } }
+        # Per-token response queues for lock-free message delivery. @resp_map is a
+        # Concurrent::Map so request threads and dispatcher threads can insert,
+        # look up, and delete tokens without serializing on a single mutex (on
+        # JRuby it is backed by java.util.concurrent.ConcurrentHashMap). Each
+        # value is a Hash { queue:, created_at: }.
+        @resp_map = ::Concurrent::Map.new
         @resp_handlers = []
-        @map_lock = ::Mutex.new  # Lightweight lock only for map structure changes
         @cleanup_thread = nil
         @shutdown = false
         @cleanup_mutex = ::Mutex.new
@@ -32,18 +33,38 @@ module Protobuf
         ::Protobuf::Logging.logger
       end
 
-      def cleanup(token)
-        @map_lock.synchronize do
-          # Close the queue to wake any waiting threads
-          queue = @resp_map.dig(token, :queue)
-          queue&.close
-          @resp_map.delete(token)
+      # Monotonic clock for token TTL accounting. Cheaper than Time.now (no Time
+      # object / timezone work per request) and immune to wall-clock jumps.
+      def monotonic_now
+        ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+      end
+
+      # Number of dispatcher threads draining the response subscription. On JRuby
+      # (true parallelism) a single dispatcher is a hard throughput ceiling, so we
+      # fan out to processor_count; on CRuby the GVL makes extra dispatchers
+      # pointless, so we stay at 1. Overridable via env for tuning/tests.
+      def dispatcher_count
+        @dispatcher_count ||= begin
+          if ::ENV.key?("PB_NATS_RESPONSE_MUXER_DISPATCHERS")
+            [::ENV["PB_NATS_RESPONSE_MUXER_DISPATCHERS"].to_i, 1].max
+          elsif ::RUBY_ENGINE == "jruby"
+            [::Concurrent.processor_count, 1].max
+          else
+            1
+          end
         end
       end
 
+      def cleanup(token)
+        # Atomic remove-and-return; close the queue to wake any waiting threads.
+        entry = @resp_map.delete(token)
+        entry[:queue]&.close if entry
+      end
+
       def next_message(token, timeout)
-        # Get the queue for this token with minimal locking
-        queue = @map_lock.synchronize { @resp_map.dig(token, :queue) }
+        # Lock-free get of the per-token queue.
+        entry = @resp_map[token]
+        queue = entry && entry[:queue]
 
         unless queue
           logger.warn "Token #{token} not found or already cleaned up during next_message"
@@ -83,22 +104,16 @@ module Protobuf
         end
       end
 
-      def new_uuidv7
-        UUID7.generate
-      end
-
       def new_request
         # Use UUIDv7 so we can figure out what time a message was originally created in-memory.
-        token = new_uuidv7 # nats.new_inbox with nuid is not threadsafe.
-
-        @map_lock.synchronize do
-          # Create a dedicated queue for this token
-          # TimeoutQueue provides native timeout support for efficient blocking
-          @resp_map[token] = {
-            queue: Concurrent::Collection::TimeoutQueue.new,
-            created_at: Time.now
-          }
-        end
+        token = UUIDv7Helper.generate # nats.new_inbox with nuid is not threadsafe.
+
+        # Create a dedicated queue for this token. Concurrent::Map#[]= is atomic,
+        # so no surrounding lock is required.
+        @resp_map[token] = {
+          queue: ::Concurrent::Collection::TimeoutQueue.new,
+          created_at: monotonic_now
+        }
 
         ResponseMuxerRequest.new(self, token)
       end
@@ -193,117 +208,12 @@ module Protobuf
         # Start the cleanup thread
         start_cleanup_thread
 
+        # Top up the dispatcher pool to dispatcher_count. Prunes dead threads
+        # first so self-healing restarts converge to the target count instead of
+        # multiplying threads.
         LOCK.synchronize do
           @resp_handlers.select!(&:alive?)
-          @resp_handlers << Thread.new do
-            # Unique thread name for debugging
-            Thread.current.name = "response-muxer-#{Thread.current.object_id}"
-            begin
-              # Reset crash count on successful start
-              @crash_count = 0
-
-              loop do
-                begin
-                  # --- Start of per-message block ---
-                  msg = @resp_sub.pending_queue.pop
-
-                  # ACK means the message has been picked up and put into the waiting thread_pool
-                  next if msg.nil?
-
-                  # Decrease pending size since consumed already
-                  # NOTE: This is outside the lock since it's just updating metrics
-                  @resp_sub.pending_size -= msg.data.size if @resp_sub
-
-                  # Validate message subject before processing
-                  unless msg.subject.is_a?(String) && msg.subject.include?('.')
-                    ::ActiveSupport::Notifications.instrument "client.invalid_message.protobuf-nats", 1
-
-                    logger.warn "Received message with invalid subject: #{msg.subject}. Dropping."
-                    next
-                  end
-
-                  # example(random data):
-                  # _INBOX.{random_data}.{random_data_msg_id}
-                  token = msg.subject.split('.').last
-
-                  logger.debug { "token: #{token}, resp_map.keys:#{@resp_map.keys}" } if logger.debug?
-
-                  # Get the queue for this token with minimal locking
-                  queue = @map_lock.synchronize do
-                    unless @resp_map.key?(token)
-                      # Try to decode the UUIDv7 timestamp to calculate message age
-                      delay_seconds = UUIDv7Helper.age_in_seconds(token)
-
-                      ::ActiveSupport::Notifications.instrument "client.unexpected_message.protobuf-nats", delay_seconds || 1
-
-                      if delay_seconds
-                        logger.warn "Received unexpected message (#{delay_seconds.round(3)}s old). MSG.subject=#{msg.subject}. RESP_SUBJ.subject=#{@resp_sub.subject rescue 'unknown'}. Dropping unexpected message."
-                      else
-                        logger.warn "Received unexpected message. MSG.subject=#{msg.subject}. RESP_SUBJ.subject=#{@resp_sub.subject rescue 'unknown'}. Dropping unexpected message."
-                      end
-                      nil
-                    else
-                      @resp_map[token][:queue]
-                    end
-                  end
-
-                  # Skip if token wasn't found
-                  next unless queue
-
-                  # Push message onto the queue - this is lock-free and thread-safe
-                  # The Queue implementation handles all synchronization internally
-                  begin
-                    # Check queue size to prevent memory bloat
-                    if queue.size >= MAX_RESPONSES_PER_TOKEN
-                      logger.warn "Token #{token} has #{queue.size} queued responses. Possible duplicate messages or slow consumer. Dropping message."
-                      next
-                    end
-
-                    queue.push(msg)
-                  rescue ThreadError => e
-                    # Queue was closed (cleanup happened) - this is fine, just drop the message
-                    logger.debug "Queue closed for token #{token}, dropping message"
-                  end
-
-                  # --- End of per-message block ---
-                rescue => per_message_error
-                  # ThreadError is fatal, it means the queue is closed and the loop cannot continue.
-                  raise if per_message_error.is_a?(::ThreadError)
-
-                  # Log the error for the specific message, but DON'T kill the thread.
-                  logger.error("ResponseMuxer failed to process a message. Error: #{per_message_error.message}")
-                  ::Protobuf::Nats.notify_error_callbacks(per_message_error)
-                end
-              end
-            rescue => fatal_error
-              # This block is now only for truly fatal errors that kill the loop itself.
-              logger.error("ResponseMuxer thread crashed fatally. Error: #{fatal_error.message}")
-              ::Protobuf::Nats.notify_error_callbacks(fatal_error)
-
-              # --- Self-healing logic ---
-              @crash_count = (@crash_count || 0) + 1
-              # Exponential backoff, e.g., 1, 4, 9, 16s... capped at 60s.
-              sleep_duration = [(@crash_count**2), 60].min
-              logger.warn("Waiting #{sleep_duration}s before attempting to restart ResponseMuxer.")
-              sleep sleep_duration
-              # --- End of self-healing logic ---
-
-              # After sleeping, reset the state and try to start again.
-              LOCK.synchronize do
-                if @resp_sub
-                  begin
-                    @resp_sub.unsubscribe
-                  rescue => e
-                    logger.warn "Failed to unsubscribe old response muxer subscription during self-healing: #{e.message}"
-                  ensure
-                    @resp_sub = nil
-                  end
-                end
-                @started = false
-              end
-              start
-            end
-          end
+          @resp_handlers << spawn_dispatcher while @resp_handlers.size < dispatcher_count
         end
       end
 
@@ -313,25 +223,29 @@ module Protobuf
 
       # Periodic cleanup of stale tokens
       def cleanup_stale_tokens
-        cutoff = Time.now - TOKEN_TTL_SECONDS
-
-        @map_lock.synchronize do
-          stale_count = 0
-          @resp_map.delete_if do |token, data|
-            if data[:created_at] && data[:created_at] < cutoff
-              stale_count += 1
-              logger.warn "Cleaning up stale token #{token} created at #{data[:created_at]}"
-              # Close the queue to wake any waiting threads
-              data[:queue]&.close
-              true
-            else
-              false
-            end
-          end
+        cutoff = monotonic_now - TOKEN_TTL_SECONDS
+
+        # Collect stale tokens first, then delete. Concurrent::Map iteration does
+        # not hold a global lock, so request threads are never blocked across this
+        # O(n) scan (unlike the previous single-mutex implementation).
+        stale_tokens = []
+        @resp_map.each_pair do |token, data|
+          created_at = data[:created_at]
+          stale_tokens << token if created_at && created_at < cutoff
+        end
 
-          if stale_count > 0
-            ::ActiveSupport::Notifications.instrument "response_muxer.stale_tokens_cleaned.protobuf-nats", stale_count
-          end
+        stale_count = 0
+        stale_tokens.each do |token|
+          data = @resp_map.delete(token)
+          next unless data
+          stale_count += 1
+          logger.warn "Cleaning up stale token #{token} created at #{data[:created_at]}"
+          # Close the queue to wake any waiting threads
+          data[:queue]&.close
+        end
+
+        if stale_count > 0
+          ::ActiveSupport::Notifications.instrument "response_muxer.stale_tokens_cleaned.protobuf-nats", stale_count
         end
       end
 
@@ -360,6 +274,124 @@ module Protobuf
         !!@started
       end
 
+      # Spawn a single dispatcher thread. Multiple dispatchers safely share the
+      # one @resp_sub.pending_queue (Queue is thread-safe) and route via the
+      # lock-free @resp_map.
+      def spawn_dispatcher
+        Thread.new do
+          # Unique thread name for debugging
+          Thread.current.name = "response-muxer-#{Thread.current.object_id}"
+          begin
+            # Reset crash count on successful start
+            @crash_count = 0
+            run_dispatch_loop
+          rescue => fatal_error
+            # Only truly fatal errors that kill the loop reach here (ThreadError
+            # from the shared pending_queue being closed).
+            logger.error("ResponseMuxer thread crashed fatally. Error: #{fatal_error.message}")
+            ::Protobuf::Nats.notify_error_callbacks(fatal_error)
+
+            # --- Self-healing logic ---
+            @crash_count = (@crash_count || 0) + 1
+            # Exponential backoff, e.g., 1, 4, 9, 16s... capped at 60s.
+            sleep_duration = [(@crash_count**2), 60].min
+            logger.warn("Waiting #{sleep_duration}s before attempting to restart ResponseMuxer.")
+            sleep sleep_duration
+            # --- End of self-healing logic ---
+
+            # After sleeping, reset the state and try to start again.
+            LOCK.synchronize do
+              if @resp_sub
+                begin
+                  @resp_sub.unsubscribe
+                rescue => e
+                  logger.warn "Failed to unsubscribe old response muxer subscription during self-healing: #{e.message}"
+                ensure
+                  @resp_sub = nil
+                end
+              end
+              @started = false
+            end
+            start
+          end
+        end
+      end
+
+      def run_dispatch_loop
+        loop do
+          begin
+            # --- Start of per-message block ---
+            msg = @resp_sub.pending_queue.pop
+
+            # ACK means the message has been picked up and put into the waiting thread_pool
+            next if msg.nil?
+
+            # Decrease pending size since consumed already.
+            # NOTE: advisory metric only; with multiple dispatchers this is a
+            # benign lost-update race on the NATS subscription's counter.
+            @resp_sub.pending_size -= msg.data.size if @resp_sub
+
+            dispatch_message(msg)
+            # --- End of per-message block ---
+          rescue => per_message_error
+            # ThreadError is fatal, it means the queue is closed and the loop cannot continue.
+            raise if per_message_error.is_a?(::ThreadError)
+
+            # Log the error for the specific message, but DON'T kill the thread.
+            logger.error("ResponseMuxer failed to process a message. Error: #{per_message_error.message}")
+            ::Protobuf::Nats.notify_error_callbacks(per_message_error)
+          end
+        end
+      end
+
+      def dispatch_message(msg)
+        # Validate message subject before processing
+        unless msg.subject.is_a?(String) && msg.subject.include?('.')
+          ::ActiveSupport::Notifications.instrument "client.invalid_message.protobuf-nats", 1
+
+          logger.warn "Received message with invalid subject: #{msg.subject}. Dropping."
+          return
+        end
+
+        # example(random data):
+        # _INBOX.{random_data}.{random_data_msg_id}
+        token = msg.subject.split('.').last
+
+        logger.debug { "token: #{token}, resp_map.keys:#{@resp_map.keys}" } if logger.debug?
+
+        # Lock-free get of the per-token queue.
+        entry = @resp_map[token]
+        queue = entry && entry[:queue]
+
+        unless queue
+          # Try to decode the UUIDv7 timestamp to calculate message age
+          delay_seconds = UUIDv7Helper.age_in_seconds(token)
+
+          ::ActiveSupport::Notifications.instrument "client.unexpected_message.protobuf-nats", delay_seconds || 1
+
+          if delay_seconds
+            logger.warn "Received unexpected message (#{delay_seconds.round(3)}s old). MSG.subject=#{msg.subject}. RESP_SUBJ.subject=#{@resp_sub.subject rescue 'unknown'}. Dropping unexpected message."
+          else
+            logger.warn "Received unexpected message. MSG.subject=#{msg.subject}. RESP_SUBJ.subject=#{@resp_sub.subject rescue 'unknown'}. Dropping unexpected message."
+          end
+          return
+        end
+
+        # Push message onto the queue - this is lock-free and thread-safe.
+        begin
+          # Check queue size to prevent memory bloat
+          if queue.size >= MAX_RESPONSES_PER_TOKEN
+            logger.warn "Token #{token} has #{queue.size} queued responses. Possible duplicate messages or slow consumer. Dropping message."
+            return
+          end
+
+          queue.push(msg)
+        rescue ThreadError
+          # Queue was closed (cleanup happened) - this is fine, just drop the message
+          logger.debug "Queue closed for token #{token}, dropping message"
+        end
+      end
+
       def start_cleanup_thread
         # Only start if not already running
         return if @cleanup_thread&.alive?

data/lib/protobuf/nats/thread_pool.rb

@@ -1,10 +1,15 @@
+require "concurrent"
+
 module Protobuf
   module Nats
     class ThreadPool
 
       def initialize(size, opts = {})
         @queue = ::Queue.new
-        @active_work = 0
+        # Lock-free counter of in-flight work. Replaces a mutex-guarded integer so
+        # that N workers running in true parallel (JRuby) don't serialize on every
+        # task completion.
+        @active_work = ::Concurrent::AtomicFixnum.new(0)
 
         # Callbacks
         @error_cb = lambda do |error|
@@ -14,14 +19,14 @@ module Protobuf
         end
 
         # Synchronization
-        @mutex = ::Mutex.new
+        @mutex = ::Mutex.new      # guards the @workers array only
         @cb_mutex = ::Mutex.new
 
         # Let's get this party started
         queue_size = opts[:max_queue].to_i || 0
         @max_size = size + queue_size
         @max_workers = size
-        @shutting_down = false
+        @shutting_down = ::Concurrent::AtomicBoolean.new(false)
         @workers = []
         supervise_workers
       end
@@ -32,7 +37,7 @@ module Protobuf
 
       # Thread-safe access to check if the pool is full.
       def full?
-        @mutex.synchronize { @active_work >= @max_size }
+        @active_work.value >= @max_size
       end
 
       def max_size
@@ -41,33 +46,33 @@ module Protobuf
 
       # This method is now thread-safe.
       def push(&work_cb)
-        @mutex.synchronize do
-          # Re-check conditions inside the lock to guarantee safety.
-          return false if @active_work >= @max_size
-          return false if @shutting_down
-
-          @queue << [:work, work_cb]
-          @active_work += 1
+        return false if @shutting_down.true?
+
+        # Optimistically claim a slot; back off if we exceeded the cap. This admits
+        # work only while active_work < max_size, matching the original guard, but
+        # without holding a mutex across the enqueue.
+        if @active_work.increment > @max_size
+          @active_work.decrement
+          return false
         end
 
-        # Supervise outside the lock to avoid holding it during thread creation.
+        @queue << [:work, work_cb]
+
+        # Supervise outside any lock-held section to avoid holding it during thread creation.
         supervise_workers
         true
       end
 
       # This method is now thread-safe.
       def shutdown
-        @mutex.synchronize do
-          return if @shutting_down # Prevent sending stop messages multiple times
-          @shutting_down = true
-        end
+        # CAS ensures the poison pills are pushed exactly once.
+        return unless @shutting_down.make_true
 
-        # Pushing poison pills can happen outside the lock.
         @max_workers.times { @queue << [:stop, nil] }
       end
 
       def kill
-        @shutting_down = true
+        @shutting_down.make_true
         @workers.map(&:kill)
       end
 
@@ -88,7 +93,7 @@ module Protobuf
 
       # Thread-safe access to the current active work size.
       def size
-        @mutex.synchronize { @active_work }
+        @active_work.value
       end
 
     private
@@ -98,7 +103,7 @@ module Protobuf
       end
 
       def prune_dead_workers
-        # This must be called inside a mutex block.
+        # This must be called inside @mutex.
         @workers = @workers.select(&:alive?)
       end
 
@@ -126,7 +131,7 @@ module Protobuf
             rescue => error
               @cb_mutex.synchronize { @error_cb.call(error) }
             ensure
-              @mutex.synchronize { @active_work -= 1 }
+              @active_work.decrement
             end
           end
         end

data/lib/protobuf/nats/uuidv7_helper.rb

@@ -1,6 +1,29 @@
 module Protobuf
   module Nats
     class UUIDv7Helper
+      # Generate a UUIDv7 string without a CSPRNG. Callers that only need a
+      # 48-bit millisecond timestamp prefix (so #age_in_seconds can report a
+      # value) plus enough randomness to stay unique among concurrent generators
+      # don't need SecureRandom: its gen_random call dominated per-request CPU
+      # and garbage (measured ~6.8us/op and 4 GC-triggering allocations). A
+      # per-thread non-cryptographic Random halves both. The layout still matches
+      # RFC 9562 UUIDv7 (version 7 + RFC 4122 variant bits).
+      #
+      # @return [String] a UUIDv7 string (e.g. "01234567-89ab-7def-8123-456789abcdef")
+      def self.generate
+        unix_ts_ms = ::Process.clock_gettime(::Process::CLOCK_REALTIME, :millisecond) & 0xffffffffffff
+        rng = (::Thread.current[:pb_nats_uuid_rng] ||= ::Random.new)
+        format(
+          "%08x-%04x-%04x-%04x-%04x%08x",
+          (unix_ts_ms >> 16) & 0xffffffff,   # 32 high bits of the ms timestamp
+          unix_ts_ms & 0xffff,               # 16 low bits of the ms timestamp
+          (0x7000 | rng.rand(0x1000)),       # version 7 + 12 random bits
+          (0x8000 | rng.rand(0x4000)),       # RFC 4122 variant + 14 random bits
+          rng.rand(0x10000),                 # 16 random bits
+          rng.rand(0x100000000)              # 32 random bits
+        )
+      end
+
       # Extract the Unix timestamp (in seconds) from a UUIDv7 string
       # Returns nil if the UUID cannot be parsed
       #

data/lib/protobuf/nats/version.rb

@@ -1,5 +1,5 @@
 module Protobuf
   module Nats
-    VERSION = "0.13.0.pre4"
+    VERSION = "0.13.0.pre5"
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: protobuf-nats
 version: !ruby/object:Gem::Version
-  version: 0.13.0.pre4
+  version: 0.13.0.pre5
 platform: ruby
 authors:
 - Brandon Dewitt
@@ -200,7 +200,9 @@ files:
 - Rakefile
 - bench/bench.md
 - bench/client.rb
+- bench/concurrency_bench.rb
 - bench/console.rb
+- bench/e2e_bench.rb
 - bench/real_client.rb
 - bench/real_client.sh
 - bench/real_client_threaded.rb