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

QUEUE_WAKE =

Sentinel pushed onto a token's queue to wake a waiter blocked in next_message. We cannot rely on Queue#close alone: on JRuby, close does NOT wake a pop() that is blocked with a timeout: -- neither the native Queue (Ruby >= 3.2 / JRuby 10) nor concurrent-ruby's RubyTimeoutQueue (Ruby < 3.2 / JRuby 9.4, whose timed pop only wakes on push) signals a timed waiter on close. CRuby's Queue#close does wake it, which is why this only ever bit JRuby. Pushing an explicit sentinel wakes the waiter immediately on every engine; next_message treats it as a timeout.

::Object.new

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.

• #subscribed_to?(nats) ⇒ Boolean

  True when the muxer's inbox subscription lives on this exact connection object.

• #token_ttl_seconds ⇒ Object

  Token TTL.

• #wake_and_close_queue(queue) ⇒ Object

  Wake any waiter blocked in next_message on this queue, then close it.

Constructor Details

#initialize ⇒ ResponseMuxer

Returns a new instance of ResponseMuxer.

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

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
  # The connection object the inbox subscription lives on. Compared by
  # identity in #start so a rebuilt connection (nats-pure fired on_close
  # and start_client_nats_connection made a fresh client) triggers a
  # restart instead of leaving the muxer subscribed to a dead connection.
  # An AtomicReference (not a plain ivar) so #start's healthy fast path
  # can read it without taking LOCK -- start runs once per RPC, and on
  # JRuby a per-request LOCK acquisition is real contention. Writes still
  # happen only while holding LOCK.
  @subscribed_nats = ::Concurrent::AtomicReference.new(nil)

  # 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 78

def cleanup(token)
  # Atomic remove-and-return; wake+close the queue to release any waiter.
  entry = @resp_map.delete(token)
  wake_and_close_queue(entry[:queue]) if entry
end

#cleanup_stale_tokens ⇒ Object

Periodic cleanup of stale tokens

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

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]}"
    # Wake any waiting thread, then close the queue.
    wake_and_close_queue(data[:queue])
  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 71

def dispatcher_count
  @dispatcher_count ||= begin
    default = ::RUBY_ENGINE == "jruby" ? ::Concurrent.processor_count : 1
    ::Protobuf::Nats.env_int("PB_NATS_RESPONSE_MUXER_DISPATCHERS", default, :min => 1)
  end
end

#logger ⇒ Object

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

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 63

def monotonic_now
  ::Protobuf::Nats.monotonic_time
end

#new_request ⇒ Object

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

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 99

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
    # QUEUE_WAKE is the sentinel pushed by wake_and_close_queue to wake a
    # timed pop on JRuby (where close alone does not); treat it as a
    # timeout so the caller fails over instead of returning garbage.
    if msg.nil? || msg.equal?(QUEUE_WAKE)
      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 159

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
  # The memoized connection is dropped when nats-pure fires on_close
  # (reconnect attempts exhausted). Raise the muxer's retryable error
  # instead of NoMethodError-on-nil so the client's transient-transport
  # retry path rebuilds the connection and tries again.
  if nats.nil?
    raise ::Protobuf::Nats::Errors::ResponseMuxer, "NATS connection unavailable (closed and not yet rebuilt) - cannot publish"
  end

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

#restart ⇒ Object

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

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
      drop_subscription_locked("during restart")

      # Stop the cleanup thread
      stop_cleanup_thread
    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 216

def start
  current_nats = ::Protobuf::Nats.client_nats_connection

  # Runs in Client#initialize, i.e. once per RPC, so the healthy path is
  # lock-free: a volatile read of the connection the inbox subscription
  # lives on. When set, also detect a replaced connection (nats-pure
  # fired on_close, on_close dropped the memoized client, and the next
  # request built a fresh one): our inbox subscription lived on the dead
  # connection, so without a rebuild every response would be lost and
  # every RPC would time out until the process restarted.
  subscribed = @subscribed_nats.get
  return if subscribed && (current_nats.nil? || subscribed.equal?(current_nats))

  # Slow path: not started, or the connection was replaced. Re-check
  # under LOCK (double-checked locking; the atomic read above may race a
  # concurrent start/restart).
  stale = false
  LOCK.synchronize do
    if _started?
      return if current_nats.nil? || @subscribed_nats.get.equal?(current_nats)
      stale = true
    end
  end

  if stale
    logger.warn "ResponseMuxer NATS connection was replaced; restarting the muxer on the new connection"
    restart
    return
  end

  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}.*")
      ::Protobuf::Nats.disable_subscription_byte_limit!(@resp_sub)
      @subscribed_nats.set(nats)
      @started = true
    rescue => e
      # Clean up partial state
      @resp_inbox_prefix = nil
      @resp_sub = nil
      @subscribed_nats.set(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 286

def started?
  LOCK.synchronize { _started? }
end

#stop ⇒ Object

Stop the cleanup thread

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

def stop
  LOCK.synchronize do
    stop_cleanup_thread
    @resp_handlers.each(&:kill)
    @resp_handlers.clear
    drop_subscription_locked("during stop")
  end
end

#subscribed_to?(nats) ⇒ Boolean

True when the muxer's inbox subscription lives on this exact connection object. Identity (not equality) is the point: a rebuilt connection to the same servers is still a different socket with no subscriptions.

Returns:

• (Boolean)

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

def subscribed_to?(nats)
  @subscribed_nats.get.equal?(nats)
end

#token_ttl_seconds ⇒ Object

Token TTL. Floors at TOKEN_TTL_SECONDS but stretches when the client's response_timeout is configured beyond it -- otherwise the cleanup thread would close a token's queue out from under a caller still legitimately waiting on a long response.

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

def token_ttl_seconds
  @token_ttl_seconds ||= [TOKEN_TTL_SECONDS, ::Protobuf::Nats.client_response_timeout + 60].max
end

#wake_and_close_queue(queue) ⇒ Object

Wake any waiter blocked in next_message on this queue, then close it. Pushing QUEUE_WAKE is what actually wakes a timed pop on JRuby (see the QUEUE_WAKE comment); close alone is insufficient there. Safe to call on an already-closed queue.

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

def wake_and_close_queue(queue)
  return unless queue
  begin
    queue.push(QUEUE_WAKE)
  rescue ::ClosedQueueError, ::ThreadError
    # Already closed by another path; a plain (untimed) waiter, if any,
    # was already woken by that close. Nothing more to do.
  end
  queue.close
end