Aug26 more improvements

CHANGELOG.md

@@ -1,5 +1,21 @@
 ## Changelog
 
+### 0.13.3.pre1
+Correctness fixes for concurrency bugs found while reviewing the 0.13.1/0.13.2 changes. Every fix ships with a spec verified to fail against the previous code. No API or configuration changes.
+
+#### Byte accounting
+- Fixed permanent upward drift in the `ByteBoundedQueue` byte counter (new in 0.13.2). Bytes were counted *after* the enqueue, so a consumer could pop an item and subtract its bytes first; `pop`'s clamp at zero swallowed that subtraction, and the producer's increment then applied to an item that was already gone. The counter only ratcheted up, so a long-running server eventually reached the 128 MiB ceiling and dropped **every** request while still looking healthy. The server pops the shared queue from `processor_count` handler threads on JRuby, so the race was live on every message. Bytes are now counted before the enqueue and rolled back if it does not happen.
+
+#### Shutdown
+- Work accepted just as shutdown begins is no longer stranded. `ThreadPool#push` checks the shutdown flag and then enqueues, so `shutdown` could slip its poison pills between those two steps; workers took a pill and exited, leaving an already-ACKed request to hang until the client's response timeout (60s default). Workers now drain work queued behind their pill, and `wait_for_termination` drains once more after the last worker exits. Admission and shutdown are still not atomic (`#push` is lock-free by design), but nothing enqueued before the pool reports termination is dropped. The final drain also skips past orphan poison pills: a worker killed by a non-`StandardError` is not replaced during shutdown, so its pill outlives every worker and used to stop the drain with work still behind it.
+- Removed the `Timeout.timeout(10)` wrapper around `SuperSubscriptionManager#shutdown` in `Server#run`. 0.13.1 removed `Timeout` from the manager because its async `Thread#raise`, fired while a thread holds the `SizedQueue` mutex, makes JRuby unwind through the held mutex; the wrapper reintroduced that hazard one frame up. It could genuinely fire: `#shutdown` is not bounded by 10s (one 1s push deadline per handler, then a 5s join, then 1s kill-joins). `#shutdown` already self-bounds, so the wrapper is gone along with the now-dead `require "timeout"`.
+
+#### Self-healing
+- A response-muxer dispatcher that crashes now tears down only the subscription it actually died on. Dispatchers that crash together wake on staggered backoffs (1s, then 4s), so a late one destroyed the subscription an earlier one had just rebuilt and, via `fail_inflight_requests`, cancelled every request already waiting on it. A dispatcher that finds its subscription already healed now spawns its replacement directly; calling `start` there returned at the fast path and left the pool one dispatcher short per crash.
+
+#### Observability
+- `UUIDv7Helper.extract_timestamp` validates the whole token instead of just its length. `String#to_i(16)` stops at the first non-hex character and returns 0 rather than raising, so a foreign reply token parsed as epoch 0 and reported a ~56-year age into the `client.unexpected_message` gauge. Both the dashed and compact (dash-free) UUIDv7 forms are still accepted.
+
 ### 0.13.2
 Bounds the RPC transport's in-memory buffering to prevent the JVM-heap OOM introduced by the JNats → nats-pure migration. Both the client response muxer and the server intake queue are now capped by message count **and** total bytes, dropping (with client retry) rather than buffering unbounded protobuf payloads on the heap.
 

lib/protobuf/nats.rb

@@ -1,7 +1,7 @@
 require "protobuf/nats/version"
 
 require "protobuf"
-# We don't need this, but the CLI attempts to terminate.
+# Unused here, but the protobuf CLI calls ServiceDirectory#stop on shutdown.
 require "protobuf/rpc/service_directory"
 
 require "nats/io/client"
@@ -42,15 +42,15 @@ def self.config
       end
     end
 
-    # Eagerly load the yml config.
+    # Load the YAML config now.
     config
 
-    # We will always log  an error.
+    # Always log an error.
     def self.error_callbacks
       @error_callbacks ||= [lambda { |error| log_error(error) }]
     end
 
-    # Eagerly load the yml config.
+    # Load the default error callback now.
     error_callbacks
 
     def self.on_error(&block)
@@ -59,10 +59,8 @@ def self.on_error(&block)
       nil
     end
 
-    # Single instrumentation entry point. Appends the gem's `.protobuf-nats`
-    # suffix so callers don't repeat it (and can't typo it). Supports both the
-    # value form `instrument("server.x", 5)` and the block form
-    # `instrument("client.request_duration") { ... }`.
+    # Single entry point for instrumentation. Adds the `.protobuf-nats`
+    # suffix so callers do not repeat it.
     def self.instrument(event, payload = {}, &block)
       ::ActiveSupport::Notifications.instrument("#{event}.protobuf-nats", payload, &block)
     end
@@ -79,40 +77,34 @@ def self.notify_error_callbacks(error)
       nil
     end
 
-    # Bounded, single-thread executor for running error callbacks OFF hot/shared
-    # threads (notably nats-pure's read/flush thread via on_error). A slow user
-    # callback must not stall message processing for every subject. The queue is
-    # bounded and over-capacity notifications are discarded (they're advisory).
+    # Runs error callbacks off nats-pure's read/flush thread, so a slow
+    # callback cannot stall message processing. Drops callbacks over the
+    # queue limit; they are advisory only.
     ERROR_CALLBACK_EXECUTOR = ::Concurrent::ThreadPoolExecutor.new(
       :min_threads => 0,
       :max_threads => 1,
       :max_queue => 1024,
       :fallback_policy => :discard
     )
 
-    # Count of error callbacks discarded because the bounded executor was
-    # saturated. Lets a flood of dropped callbacks during an incident be observed
-    # instead of vanishing silently.
+    # Counts error callbacks dropped when the queue is full. Makes a flood
+    # of drops visible during an incident.
     ERROR_CALLBACK_DROP_COUNT = ::Concurrent::AtomicFixnum.new(0)
 
     def self.error_callback_drop_count
       ERROR_CALLBACK_DROP_COUNT.value
     end
 
     def self.notify_error_callbacks_async(error)
-      # #post returns false when the job is rejected. With the :discard fallback
-      # policy the job is silently dropped (returning false) rather than raising,
-      # so the false return is the only drop signal to handle.
+      # #post returns false on rejection. The :discard policy drops the
+      # job instead of raising, so false is the only drop signal.
       accepted = ERROR_CALLBACK_EXECUTOR.post { notify_error_callbacks(error) }
       record_dropped_error_callback unless accepted
       nil
     end
 
-    # Record a discarded error callback. Kept cheap -- this runs on nats-pure's
-    # read/flush thread, so it must NOT format/log the error synchronously (the
-    # whole point of the async path). The atomic counter is the durable signal;
-    # the instrument gauge emits a discrete event for dashboards (drops only
-    # happen under a severe flood, so a notification per drop is acceptable).
+    # Records a dropped callback. Runs on nats-pure's read/flush thread, so
+    # do NOT format or log the error here; that would defeat the async path.
     def self.record_dropped_error_callback
       ERROR_CALLBACK_DROP_COUNT.increment
       instrument("error_callback_dropped", 1)
@@ -133,16 +125,16 @@ def self.start_client_nats_connection
       GET_CONNECTED_MUTEX.synchronize do
         break true if @client_nats_connection
 
-        # NOTE: nats-pure has no :disable_reconnect_buffer option (it was a
-        # jnats concept). During a reconnect nats-pure buffers publishes and,
-        # if the connection is fully closed, raises ConnectionClosedError --
-        # both of which the client's transient-error retry path now handles.
+        # nats-pure has no :disable_reconnect_buffer option (a jnats
+        # concept). It buffers publishes during reconnect, and raises
+        # ConnectionClosedError if the connection fully closes. The
+        # client's retry path handles both cases.
         options = config.connection_options
 
         client = NatsClient.new
 
-        # Register lifecycle callbacks BEFORE connecting so a disconnect or
-        # error during the initial handshake is still observed.
+        # Register lifecycle callbacks before connecting, so the handshake
+        # is also observed.
         client.on_disconnect do
           logger.warn("Client NATS connection was disconnected")
         end
@@ -153,28 +145,26 @@ def self.start_client_nats_connection
 
         client.on_close do
           logger.warn("Client NATS connection was closed")
-          # A close is terminal for this client object (nats-pure only reconnects
-          # via on_disconnect/on_reconnect; on_close means it gave up). Drop the
-          # memoized reference so the next start_client_nats_connection rebuilds a
-          # fresh connection instead of reusing a permanently-dead one. In-flight
-          # callers keep their own local reference; only new calls rebuild.
+          # A close is terminal (nats-pure only reconnects via
+          # on_disconnect/on_reconnect). Clear the memo, so the next call
+          # builds a fresh connection. Callers with a local reference keep it.
           @client_nats_connection = nil
         end
 
         client.on_error do |error|
-          # Runs on nats-pure's read/flush thread -- offload so a slow callback
-          # can't stall message processing.
+          # Runs on nats-pure's read/flush thread; offload it so a slow
+          # callback cannot stall message processing.
           notify_error_callbacks_async(error)
         end
 
         begin
           client.connect(options)
-          # Ensure we have a valid connection to the NATS server.
+          # Confirm the connection is valid.
           client.flush(5)
         rescue => e
-          # A failed handshake can leave nats-pure's reader/flusher threads
-          # running on a half-open client; close it so we don't leak them, then
-          # surface the failure (the next call will retry with a fresh client).
+          # A failed handshake can leave nats-pure's threads running on a
+          # half-open client. Close it to avoid a leak, then raise; the
+          # next call retries with a fresh client.
           client.close rescue nil
           raise e
         end
@@ -185,17 +175,16 @@ def self.start_client_nats_connection
       end
     end
 
-    # Monotonic clock for durations/ages; immune to wall-clock (NTP) jumps.
-    # Single source of truth shared by the client muxer and server pools.
+    # Monotonic clock for durations and ages, immune to wall-clock (NTP)
+    # jumps. Shared by the client muxer and server pools.
     def self.monotonic_time
       ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
     end
 
-    # Strict integer parsing for env overrides. String#to_i silently turns a
-    # malformed value ("5s", "abc") into 0 -- which for a timeout means "fail
-    # every request instantly". Log loudly and fall back to the default
-    # instead. Values below `min` (when given) are rejected the same way, so
-    # range policy lives here rather than ad hoc at each call site.
+    # Strict integer parsing for env overrides. String#to_i silently turns
+    # a bad value ("5s", "abc") into 0, which fails every request instantly
+    # for a timeout. Log an error and use the default instead. Also
+    # rejects a value below `min`.
     def self.env_int(name, default, min: nil)
       raw = ::ENV[name]
       return default if raw.nil?
@@ -211,7 +200,7 @@ def self.env_int(name, default, min: nil)
       default
     end
 
-    # Float sibling of env_int, same strict-parse-or-default contract.
+    # Float version of env_int. Same strict-parse-or-default rule.
     def self.env_float(name, default)
       raw = ::ENV[name]
       return default if raw.nil?
@@ -221,38 +210,33 @@ def self.env_float(name, default)
       default
     end
 
-    # Client response timeout (seconds). Single source of truth for the env
-    # var and its default: the client waits this long per request, and the
-    # muxer stretches its token TTL past it (ResponseMuxer#token_ttl_seconds).
+    # Client response timeout, in seconds. The muxer sets its token TTL
+    # longer than this (ResponseMuxer#token_ttl_seconds).
     def self.client_response_timeout
       env_int("PB_NATS_CLIENT_RESPONSE_TIMEOUT", 60)
     end
 
-    # How long a consumer loop parks when its queue pops nil (a closed queue
-    # returns nil immediately forever). Shared by the muxer dispatch loop and
-    # the server intake handlers so the two mirrored loops can't drift.
+    # How long a consumer loop waits after its queue pops nil (a closed
+    # queue always returns nil). Shared by the muxer and server intake
+    # loops, so they cannot drift apart.
     CLOSED_QUEUE_PARK_SECONDS = 0.05
 
-    # nats-pure increments a subscription's pending_size (bytes) for every
-    # inbound message and only decrements it in its own consumption paths
-    # (next_msg / the sub's message thread). The server intake pops
-    # pending_queue directly and never runs those paths, so pending_size grows
-    # monotonically and the byte-based slow-consumer limit would eventually trip
-    # on *cumulative* traffic -- silently dropping every later message on that
-    # subscription. Disable the byte limit; the message-count limit
-    # (pending_queue depth, tracked accurately for free) still bounds a genuinely
-    # slow consumer. Guarded so a non-standard/faked subscription is a no-op.
+    # nats-pure only decrements pending_size (bytes) in its own consumption
+    # paths. Server intake pops pending_queue directly and skips those
+    # paths, so pending_size only grows and would eventually trip the byte
+    # limit on cumulative traffic, dropping every later message. Disable
+    # the byte limit; the message-count limit (pending_queue depth) still
+    # catches a slow consumer. No-op on a non-standard subscription.
     #
-    # NOTE: only the server uses this now. The client muxer instead decrements
-    # pending_size itself after each pop (ResponseMuxer#run_dispatch_loop), which
-    # keeps the counter accurate and lets it enforce a finite byte ceiling.
+    # NOTE: only the server uses this. The client muxer decrements
+    # pending_size itself after each pop (ResponseMuxer#run_dispatch_loop)
+    # and enforces its own byte ceiling.
     def self.disable_subscription_byte_limit!(sub)
       sub.pending_bytes_limit = ::Float::INFINITY if sub.respond_to?(:pending_bytes_limit=)
     end
 
-    # Exponential backoff (seconds) for self-healing worker threads after a fatal
-    # crash, capped. Shared by the ResponseMuxer dispatcher pool and the server
-    # SuperSubscriptionManager handler pool so the formula can't drift between them.
+    # Exponential backoff, in seconds, for a worker thread after a fatal
+    # crash. Shared by the ResponseMuxer and SuperSubscriptionManager pools.
     def self.crash_backoff_seconds(crash_count, cap = 60)
       [(crash_count**2), cap].min
     end

lib/protobuf/nats/byte_bounded_queue.rb

@@ -2,21 +2,18 @@
 
 module Protobuf
   module Nats
-    # A SizedQueue that additionally bounds the total *bytes* of its contents,
-    # not just the message count. The server funnels every subscription into one
-    # shared intake queue, so a per-subscription byte limit (nats-pure's
-    # pending_bytes_limit) can't bound the aggregate heap -- this shared counter
-    # can. Count is still bounded by the SizedQueue capacity it inherits.
+    # A `SizedQueue` that also bounds total bytes, not just message count.
+    # One shared intake queue serves all subscriptions, so a per-subscription
+    # limit (nats-pure's `pending_bytes_limit`) can't bound the combined
+    # heap; this shared counter can.
     #
-    # When a push would exceed the byte ceiling we DROP the message rather than
-    # block: pushes happen on nats-pure's read thread (Subscription#dispatch), and
-    # blocking it would stall PING/PONG and every other subject. A drop mirrors
-    # nats-pure's own SlowConsumer behaviour. Non-message items (the :shutdown
-    # poison pill) carry zero bytes, so they are never dropped by the byte gate.
+    # A push over the byte limit drops the message instead of blocking, like
+    # nats-pure's SlowConsumer: pushes run on nats-pure's read thread
+    # (`Subscription#dispatch`), and blocking it would stall PING/PONG. The
+    # `:shutdown` poison pill counts as zero bytes and is never dropped.
     #
-    # A drop invokes the optional +on_drop+ callback with the dropped byte count,
-    # so the caller owns any (context-specific) instrumentation rather than this
-    # generic queue class hard-coding it.
+    # A drop calls the optional `on_drop` callback with the byte count, so
+    # the caller owns instrumentation.
     class ByteBoundedQueue < ::SizedQueue
       def initialize(max_msgs, max_bytes, on_drop: nil)
         super(max_msgs)
@@ -25,27 +22,45 @@ def initialize(max_msgs, max_bytes, on_drop: nil)
         @bytes = ::Concurrent::AtomicFixnum.new(0)
       end
 
-      # Enqueue unless it would exceed the byte ceiling. The check-then-add races
-      # only concurrent pops (which lower @bytes), so the ceiling can be exceeded
-      # by at most one in-flight message -- a soft limit, like nats-pure's own
-      # byte accounting. Returns self (SizedQueue#push contract). Raises
-      # ThreadError from super on a non_block push into a count-full queue, before
-      # any bytes are counted.
+      # Enqueue unless it would exceed the byte limit. The check-then-add
+      # races only concurrent pops, so this soft limit (like nats-pure's own
+      # byte accounting) can be exceeded by at most one in-flight message.
+      # Returns `self`, the `SizedQueue#push` contract.
       def push(obj, non_block = false)
         bytes = byte_size(obj)
         if bytes > 0 && (@bytes.value + bytes) > @max_bytes
           @on_drop&.call(bytes)
           return self
         end
-        super(obj, non_block)
-        @bytes.increment(bytes)
+
+        # Count bytes before the enqueue, and roll back on failure. Counting
+        # after `super` would let a pop subtract first; `#pop`'s zero-clamp
+        # then swallows it, and the later increment permanently overcounts a
+        # message already gone. That drift only climbs, jamming every push
+        # once it hits `@max_bytes` -- same bug class as nats-pure's
+        # `pending_size` bug.
+        #
+        # A blocking push on a full queue counts bytes while it waits: a
+        # deliberate short overcount, resolved once the push completes.
+        @bytes.increment(bytes) if bytes > 0
+        pushed = false
+        begin
+          super(obj, non_block)
+          pushed = true
+        ensure
+          # `ensure`, not `rescue`: also rolls back on an async
+          # `Thread#raise` or non-`StandardError` unwind.
+          @bytes.decrement(bytes) if bytes > 0 && !pushed
+        end
         self
       end
       alias_method :<<, :push
 
       def pop(non_block = false)
         obj = super
-        # nil == closed/empty non_block; nothing dequeued, nothing to subtract.
+        # nil means closed or empty (non-blocking): nothing to subtract. The
+        # zero-clamp only guards `#clear` racing an in-flight pop; `#push`
+        # counts bytes first, so a normal pop always finds them present.
         @bytes.update { |value| [value - byte_size(obj), 0].max } if obj
         obj
       end
@@ -55,15 +70,15 @@ def clear
         @bytes.value = 0
       end
 
-      # Current resident byte total (gauge for observability).
+      # Current resident byte total (an observability gauge).
       def bytesize
         @bytes.value
       end
 
       private
 
-      # Bytes attributable to a queued item. NATS::Msg carries #data; the
-      # :shutdown poison pill (and any other non-message sentinel) counts as 0.
+      # Bytes charged to a queued item. `NATS::Msg` carries `#data`; the
+      # `:shutdown` poison pill and other sentinels count as 0.
       def byte_size(obj)
         obj.respond_to?(:data) && obj.data ? obj.data.bytesize : 0
       end