Class: Protobuf::Nats::ResponseMuxer

Inherits:
Object
  • Object
show all
Defined in:
lib/protobuf/nats/response_muxer.rb

Constant Summary collapse

LOCK = 
::Mutex.new
MAX_RESPONSES_PER_TOKEN = 
10
TOKEN_TTL_SECONDS =

10 minutes

600
DEFAULT_RESPONSE_QUEUE_SIZE =

The shared response subscription is bounded by BOTH a message count and a byte ceiling; nats-pure drops (SlowConsumer) on whichever trips first, so the firehose is capped at min(count, bytes) instead of buffering unbounded protobuf payloads on the JVM heap (the 0.13.2 OOM). Dispatchers drain it to ~0, so these are burst headroom, not a working set.

The count is deliberately tighter than the ecosystem's per-subscription defaults (nats-pure 65,536; nats.go 500,000): those bound off-heap buffers, this is Ruby objects on the heap, and the byte cap is the real ceiling. The byte default stays aligned at 64 MiB (nats-pure/nats.go both use it). Override via PB_NATS_RESPONSE_MUXER_QUEUE_SIZE / _QUEUE_BYTES.

1024
DEFAULT_RESPONSE_QUEUE_BYTES =

64MiB

64 * 1024 * 1024
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 collapse

Constructor Details

#initialize ⇒ ResponseMuxer

Returns a new instance of ResponseMuxer.



40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
 
# File 'lib/protobuf/nats/response_muxer.rb', line 40

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)

  # High-water mark of the response queue depth since the last cleanup
  # cycle. Sampled in the dispatch loop, emitted+reset by the cleanup thread
  # (response_muxer.pending_queue_peak) so a burst between gauge samples is
  # still visible.
  @pending_queue_peak = ::Concurrent::AtomicFixnum.new(0)
end

Instance Method Details

#cleanup(token) ⇒ Object 



115
116
117
118
119
 
# File 'lib/protobuf/nats/response_muxer.rb', line 115

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



358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
 
# File 'lib/protobuf/nats/response_muxer.rb', line 358

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

  # Gauge the shared response firehose so a climbing backlog is visible
  # before it turns into timeouts/SlowConsumer drops. current == depth at
  # sample time; peak == high-water since the last cycle (reset here).
  ::Protobuf::Nats.instrument "response_muxer.pending_queue_size", pending_queue_size
  # Atomic read-and-reset of the high-water mark (AtomicFixnum has no
  # get_and_set): capture the prior value inside the update block.
  peak = 0
  @pending_queue_peak.update do |current_value|
    peak = current_value
    0 # set to 0
  end
  ::Protobuf::Nats.instrument "response_muxer.pending_queue_peak", peak
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.



91
92
93
94
95
96
 
# File 'lib/protobuf/nats/response_muxer.rb', line 91

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 



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

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.



83
84
85
 
# File 'lib/protobuf/nats/response_muxer.rb', line 83

def monotonic_now
  ::Protobuf::Nats.monotonic_time
end

#new_request ⇒ Object 



182
183
184
185
186
187
188
189
190
191
192
193
194
 
# File 'lib/protobuf/nats/response_muxer.rb', line 182

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 



136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
 
# File 'lib/protobuf/nats/response_muxer.rb', line 136

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

#pending_queue_size ⇒ Object

Current depth of the shared firehose; 0 before the muxer starts. Gauge for observability -- mirrors SuperSubscriptionManager#pending_queue_size.



111
112
113
 
# File 'lib/protobuf/nats/response_muxer.rb', line 111

def pending_queue_size
  @resp_sub&.pending_queue&.size || 0
end

#publish(subject, data, token) ⇒ Object 



196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
 
# File 'lib/protobuf/nats/response_muxer.rb', line 196

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

#response_queue_bytes ⇒ Object 



105
106
107
 
# File 'lib/protobuf/nats/response_muxer.rb', line 105

def response_queue_bytes
  ::Protobuf::Nats.env_int("PB_NATS_RESPONSE_MUXER_QUEUE_BYTES", DEFAULT_RESPONSE_QUEUE_BYTES, :min => 1)
end

#response_queue_size ⇒ Object

Message-count and byte caps for the shared response subscription (see DEFAULT_RESPONSE_QUEUE_SIZE / _BYTES). Read once each, in #start, so no memoization is needed.



101
102
103
 
# File 'lib/protobuf/nats/response_muxer.rb', line 101

def response_queue_size
  ::Protobuf::Nats.env_int("PB_NATS_RESPONSE_MUXER_QUEUE_SIZE", DEFAULT_RESPONSE_QUEUE_SIZE, :min => 1)
end

#restart ⇒ Object 



215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
 
# File 'lib/protobuf/nats/response_muxer.rb', line 215

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 



253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
 
# File 'lib/protobuf/nats/response_muxer.rb', line 253

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}.*")

      # The dispatch loop takes @resp_sub.synchronize to decrement
      # pending_size after each pop, which keeps the finite byte cap accurate.
      # nats-pure's Subscription includes MonitorMixin, so this always holds;
      # if it ever doesn't, nats-pure's internals changed in a way that would
      # break byte accounting (a growing counter that false-trips the limit
      # and drops every response). Fail loudly rather than degrade silently.
      unless @resp_sub.respond_to?(:synchronize)
        raise ::Protobuf::Nats::Errors::IncompatibleSubscription,
          "NATS subscription does not respond to #synchronize; cannot maintain pending_size byte accounting (nats-pure internals changed?)"
      end

      # Bound the firehose by both message count and bytes (see
      # DEFAULT_RESPONSE_QUEUE_SIZE / _BYTES).
      @resp_sub.pending_msgs_limit = response_queue_size
      @resp_sub.pending_bytes_limit = response_queue_bytes
      @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) 


338
339
340
 
# File 'lib/protobuf/nats/response_muxer.rb', line 338

def started?
  LOCK.synchronize { _started? }
end

#stop ⇒ Object

Stop the cleanup thread



399
400
401
402
403
404
405
406
 
# File 'lib/protobuf/nats/response_muxer.rb', line 399

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) 


345
346
347
 
# File 'lib/protobuf/nats/response_muxer.rb', line 345

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.



353
354
355
 
# File 'lib/protobuf/nats/response_muxer.rb', line 353

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.



125
126
127
128
129
130
131
132
133
134
 
# File 'lib/protobuf/nats/response_muxer.rb', line 125

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