Class: Protobuf::Nats::ResponseMuxer

Inherits:

Object

• Object
• Protobuf::Nats::ResponseMuxer

Defined in:

lib/protobuf/nats/response_muxer.rb

Constant Summary

LOCK =

::Mutex.new

MAX_RESPONSES_PER_TOKEN =

10

TOKEN_TTL_SECONDS =

10 minutes

600

DISABLED_PENDING_BYTES_LIMIT =

Slow-consumer protection on the response subscription is by message count (the SizedQueue depth, tracked accurately for free by nats-pure); the byte-based limit is disabled (see #disable_byte_limit!) so we don’t have to mirror nats-pure’s pending_size accounting on the dispatch hot path. INFINITY makes nats-pure’s ‘pending_size >= pending_bytes_limit` check always false.

::Float::INFINITY

Instance Method Summary

• #cleanup(token) ⇒ Object
• #cleanup_stale_tokens ⇒ Object

  Periodic cleanup of stale tokens.

• #dispatcher_count ⇒ Object

  Number of dispatcher threads draining the response subscription.

• #initialize ⇒ ResponseMuxer constructor

  A new instance of ResponseMuxer.

• #logger ⇒ Object
• #monotonic_now ⇒ Object

  Monotonic clock for token TTL accounting (single source of truth in Protobuf::Nats.monotonic_time).

• #new_request ⇒ Object
• #next_message(token, timeout) ⇒ Object
• #publish(subject, data, token) ⇒ Object
• #restart ⇒ Object
• #start ⇒ Object
• #started? ⇒ Boolean
• #stop ⇒ Object

  Stop the cleanup thread.

Constructor Details

#initialize ⇒ ResponseMuxer

Returns a new instance of ResponseMuxer.

# File 'lib/protobuf/nats/response_muxer.rb', line 17

def initialize
  # 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 = []
  @cleanup_thread = nil
  @shutdown = false
  @cleanup_mutex = ::Mutex.new
  @cleanup_cv = ::ConditionVariable.new
  @restarting = false  # Flag to prevent concurrent restarts

  # Shared self-healing backoff counter for the dispatcher pool. Atomic so
  # concurrent dispatchers don't lose updates when several crash at once,
  # and it decays back to zero once a dispatcher is healthy again (see
  # run_dispatch_loop), so a later transient crash restarts the backoff
  # from 1s instead of staying pinned at the cap.
  @crash_count = ::Concurrent::AtomicFixnum.new(0)
end

Instance Method Details

#cleanup(token) ⇒ Object

# File 'lib/protobuf/nats/response_muxer.rb', line 73

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

#cleanup_stale_tokens ⇒ Object

Periodic cleanup of stale tokens

# File 'lib/protobuf/nats/response_muxer.rb', line 241

def cleanup_stale_tokens
  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

  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
    ::Protobuf::Nats.instrument "response_muxer.stale_tokens_cleaned", stale_count
  end
end

#dispatcher_count ⇒ Object

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.

# File 'lib/protobuf/nats/response_muxer.rb', line 61

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

#logger ⇒ Object

# File 'lib/protobuf/nats/response_muxer.rb', line 47

def logger
  ::Protobuf::Logging.logger
end

#monotonic_now ⇒ Object

Monotonic clock for token TTL accounting (single source of truth in Protobuf::Nats.monotonic_time). Immune to wall-clock jumps.

# File 'lib/protobuf/nats/response_muxer.rb', line 53

def monotonic_now
  ::Protobuf::Nats.monotonic_time
end

#new_request ⇒ Object

# File 'lib/protobuf/nats/response_muxer.rb', line 122

def new_request
  # Use UUIDv7 so we can figure out what time a message was originally created in-memory.
  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

#next_message(token, timeout) ⇒ Object

# File 'lib/protobuf/nats/response_muxer.rb', line 79

def next_message(token, timeout)
  # 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"
    raise ::NATS::Timeout
  end

  # Handle edge case: zero or negative timeout
  if timeout && timeout <= 0
    raise ::NATS::Timeout
  end

  # Use TimeoutQueue's native timeout support for efficient, lock-free waiting per token
  # Each token has its own queue, eliminating contention between different requests
  begin
    # TimeoutQueue.pop(non_block, timeout: seconds)
    # - With timeout: blocks until message arrives or timeout expires (returns nil on timeout)
    # - Without timeout (nil): blocks indefinitely until message arrives
    msg = if timeout
            queue.pop(false, timeout: timeout)
          else
            queue.pop(false)
          end

    # Queue.pop returns nil when:
    # 1. The queue is closed
    # 2. The timeout expires
    unless msg
      logger.warn "Queue closed or timeout for token #{token} during next_message"
      raise ::NATS::Timeout
    end

    msg
  rescue ThreadError
    # Queue was closed - treat as timeout
    logger.warn "Queue closed for token #{token} during next_message"
    raise ::NATS::Timeout
  end
end

#publish(subject, data, token) ⇒ Object

# File 'lib/protobuf/nats/response_muxer.rb', line 136

def publish(subject, data, token)
  # Validate muxer started before publish
  unless @resp_inbox_prefix
    raise ::Protobuf::Nats::Errors::ResponseMuxer, "ResponseMuxer not started - cannot publish"
  end

  nats = Protobuf::Nats.client_nats_connection
  reply_to = "#{@resp_inbox_prefix}.#{token}"
  nats.publish(subject, data, reply_to)
end

#restart ⇒ Object

# File 'lib/protobuf/nats/response_muxer.rb', line 147

def restart
  logger.debug "restarting response_muxer"

  # Prevent concurrent restarts - only one restart at a time
  LOCK.synchronize do
    if @restarting
      logger.warn "Restart already in progress, skipping concurrent restart request"
      return
    end
    @restarting = true
  end

  # Yield so other restart callers spawned around the same time get a
  # chance to reach the @restarting check above and skip. Without this,
  # CRuby's GVL can let the current thread run the entire restart to
  # completion (clearing @restarting) before sibling threads even enter
  # the method, defeating the concurrent-restart guard.
  Thread.pass

  begin
    # Stop the existing muxer first, if it's running
    LOCK.synchronize do
      @resp_handlers.each(&:kill)
      @resp_handlers.clear
      if @resp_sub
        begin
          @resp_sub.unsubscribe
        rescue => e
          logger.warn "Failed to unsubscribe old response muxer subscription: #{e.message}"
        ensure
          # Always set to nil, even if unsubscribe raises
          @resp_sub = nil
        end
      end

      # Stop the cleanup thread
      stop_cleanup_thread

      @started = false
    end

    # Then start it fresh.
    start
  ensure
    # Always clear the restarting flag
    LOCK.synchronize { @restarting = false }
  end
end

#start ⇒ Object

# File 'lib/protobuf/nats/response_muxer.rb', line 196

def start
  return if started?
  LOCK.synchronize do
    # We check this twice in case another thread was waiting for the lock to
    # start this party. Use the unlocked check to prevent deadlocks.
    return if _started?

    nats = ::Protobuf::Nats.client_nats_connection
    return if nats.nil?

    # Clean up partial state on exception
    begin
      @resp_inbox_prefix = nats.new_inbox

      # Subscribe to our per-instance inbox
      @resp_sub = nats.subscribe("#{@resp_inbox_prefix}.*")
      disable_byte_limit!(@resp_sub)
      @started = true
    rescue => e
      # Clean up partial state
      @resp_inbox_prefix = nil
      @resp_sub = nil
      @started = false
      logger.error "Failed to start ResponseMuxer: #{e.message}"
      raise
    end
  end

  # 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 << spawn_dispatcher while @resp_handlers.size < dispatcher_count
  end
end

#started? ⇒ Boolean

Returns:

• (Boolean)

# File 'lib/protobuf/nats/response_muxer.rb', line 236

def started?
  LOCK.synchronize { _started? }
end

#stop ⇒ Object

Stop the cleanup thread

# File 'lib/protobuf/nats/response_muxer.rb', line 269

def stop
  LOCK.synchronize do
    stop_cleanup_thread
    @resp_handlers.each(&:kill)
    @resp_handlers.clear
    if @resp_sub
      begin
        @resp_sub.unsubscribe
      rescue => e
        logger.warn "Failed to unsubscribe during stop: #{e.message}"
      ensure
        @resp_sub = nil
      end
    end
    @started = false
  end
end