< Summary - AsyncResponse (Release / net8.0+net10.0 / unit+integration)

Information
Class: AsyncResponse.Channels.NATS.NatsAsyncResponseChannel
Assembly: AsyncResponse.Channels.NATS
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Channels/AsyncResponse.Channels.NATS/NatsAsyncResponseChannel.cs
Line coverage
100%
Covered lines: 381
Uncovered lines: 0
Coverable lines: 381
Total lines: 812
Line coverage: 100%
Branch coverage
96%
Covered branches: 141
Total branches: 146
Branch coverage: 96.5%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

File(s)

/home/runner/work/AsyncResponse/AsyncResponse/src/Channels/AsyncResponse.Channels.NATS/NatsAsyncResponseChannel.cs

#LineLine coverage
 1using Microsoft.Extensions.DependencyInjection;
 2using Microsoft.Extensions.Logging;
 3using Microsoft.Extensions.Options;
 4using System.Buffers;
 5using System.Diagnostics;
 6using System.Text;
 7using System.Text.Json;
 8
 9namespace AsyncResponse.Channels.NATS;
 10
 11/// <summary>
 12/// NATS-backed response channel:
 13/// <list type="bullet">
 14/// <item><description>Delivers responses over NATS Core request/reply on a subject keyed by
 15/// correlation id: a waiter subscribes and acks each message, and the publisher requests so the NATS
 16/// "no responders" signal reports precisely when nobody is listening.</description></item>
 17/// <item><description>Persists <see cref="RecoveryState"/> in a JetStream Key-Value bucket so a
 18/// response arriving after the waiter died (e.g. a redeploy) is routed through the lost-subscriber
 19/// dispatcher, which asks the payload's ShouldResumeOnRecovery and invokes the resume or failure
 20/// callback.</description></item>
 21/// </list>
 22/// </summary>
 23internal sealed class NatsAsyncResponseChannel : IAsyncResponsePublisher, IRawAsyncResponsePublisher, IRecoverableAsyncR
 24{
 25    private readonly INatsResponseChannelClient _client;
 26    private readonly IRecoveryStateStore _recoveryStateStore;
 27    private readonly AsyncResponseContextPropagation _propagation;
 28    private readonly LostSubscriberCallbackDispatcher _lostSubscriberDispatcher;
 29    private readonly NatsSubjectSchema _subjects;
 30    private readonly NatsAsyncResponseChannelOptions _options;
 31    private readonly ILogger<NatsAsyncResponseChannel> _logger;
 32
 33    /// <summary>Creates a NATS-backed async-response channel.</summary>
 334    public NatsAsyncResponseChannel(
 335        IServiceScopeFactory scopeFactory,
 336        INatsResponseChannelClient client,
 337        IRecoveryStateStore recoveryStateStore,
 338        IOptions<NatsAsyncResponseChannelOptions> options,
 339        AsyncResponseContextPropagation propagation,
 340        ILogger<NatsAsyncResponseChannel> logger)
 41    {
 342        _options = options.Value;
 343        _options.Validate();
 344        _client = client;
 345        _recoveryStateStore = recoveryStateStore;
 346        _propagation = propagation;
 347        _subjects = new NatsSubjectSchema(_options.SubjectPrefix);
 348        _logger = logger;
 349        _lostSubscriberDispatcher = new LostSubscriberCallbackDispatcher(scopeFactory, propagation, logger);
 350    }
 51
 52    // ---------------------------------------------------------------------------------------
 53    // IAsyncResponseSubscriber / IRecoverableAsyncResponseSubscriber
 54
 55    /// <inheritdoc/>
 56    public Task<IAsyncResponseWaiter<T>> CreateResponseWaiter<T>(
 57        string correlationId,
 58        Func<T, ValueTask<bool>>? completionPredicate = null,
 59        TimeSpan? timeout = null) where T : IAsyncResponsePayload
 360        => CreateResponseWaiterCore(correlationId, resumeCallback: null, failureCallback: null, completionPredicate, tim
 61
 62    /// <inheritdoc/>
 63    public Task<IAsyncResponseWaiter<T>> CreateRecoverableResponseWaiter<T>(
 64        string correlationId,
 65        ReflectionCallDto? resumeCallback = null,
 66        ReflectionCallDto? failureCallback = null,
 67        Func<T, ValueTask<bool>>? completionPredicate = null,
 68        TimeSpan? timeout = null) where T : IAsyncResponsePayload
 369        => CreateResponseWaiterCore(correlationId, resumeCallback, failureCallback, completionPredicate, timeout);
 70
 71    private async Task<IAsyncResponseWaiter<T>> CreateResponseWaiterCore<T>(
 72        string correlationId,
 73        ReflectionCallDto? resumeCallback,
 74        ReflectionCallDto? failureCallback,
 75        Func<T, ValueTask<bool>>? completionPredicate,
 76        TimeSpan? timeout) where T : IAsyncResponsePayload
 77    {
 378        if (string.IsNullOrWhiteSpace(correlationId))
 379            throw new ArgumentNullException(nameof(correlationId), "CorrelationId must not be empty or whitespace.");
 80
 81        // Recovery callbacks only make sense if the payload can say whether a late response should
 82        // resume or fail the flow. On this durable channel that decision is real (it survives a
 83        // redeploy), so require the override rather than letting the conservative default silently
 84        // route every recovered response to the failure callback.
 385        if ((resumeCallback is not null || failureCallback is not null)
 386            && !AsyncResponsePayloadReflection.OverridesShouldResumeOnRecovery(typeof(T)))
 87        {
 388            throw new InvalidOperationException(
 389                $"Payload type '{typeof(T)}' registers lost-subscriber recovery callbacks on the NATS channel " +
 390                $"but does not override {nameof(IAsyncResponsePayload)}.{nameof(IAsyncResponsePayload.ShouldResumeOnReco
 391                "Override it to declare which responses resume the flow (return true) versus fail it (return false); " +
 392                "the durable channel needs this to route a response that arrives after the waiter was lost.");
 93        }
 94
 95        // default: first envelope completes the wait
 396        completionPredicate ??= _ => new ValueTask<bool>(true);
 97
 98        // Default timeout aligned with the recovery-state expiry: an infinite wait is never
 99        // meaningful, because once the recovery state expires the correlation id has no recovery
 100        // anyway. Timing out routes the flow through its normal failure handling instead of
 101        // leaving it stuck forever.
 3102        timeout ??= _options.DefaultTimeout ?? _options.RecoveryStateExpiry;
 103        // BEFORE any side effect: an unsupported resolved timeout (non-positive, or past the
 104        // ~49.7-day BCL timer ceiling) used to throw only at timer arming — after the
 105        // subscription and recovery state existed, leaking both — and zero used to slip through
 106        // on some channels entirely, insta-timing-out a fully registered waiter.
 3107        AsyncResponseChannelOptions.EnsureWaiterTimeoutSupported(timeout.Value);
 108
 3109        var storedCorrelationId = correlationId;
 110        // Capture the subscribe-time ExecutionContext so app AsyncLocals (trace, principal, logging
 111        // scope) flow into the message handler, which runs on a background consume-loop thread.
 3112        var capturedContext = ExecutionContext.Capture();
 3113        var subject = _subjects.ResponseSubject(correlationId);
 114
 3115        var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.wait", correlationId: correlationId);
 3116        activity?.SetTag("asyncresponse.channel", "nats");
 3117        AsyncResponseDiagnostics.SetPayloadType(activity, typeof(T));
 3118        activity?.SetTag("asyncresponse.timeout_seconds", timeout.Value.TotalSeconds);
 119
 3120        if (_logger.IsEnabled(LogLevel.Debug))
 3121            _logger.LogDebug("Waiting for response on correlationId {CorrelationId} with timeout {Timeout}.", correlatio
 122
 3123        var tcs = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);
 3124        var registrationId = Guid.NewGuid();
 125
 126        // Single-use cancellation token implementing the timeout. Armed only after subscribe + recovery
 127        // save succeed, but its callback is registered first so a very fast terminal message cleans up safely.
 3128        var cancellationTokenSource = new CancellationTokenSource();
 3129        CancellationTokenRegistration timeoutRegistration = default;
 3130        INatsChannelSubscription? subscription = null;
 131
 132        // -------------------------------------------------------------------------
 133        // Local: CleanupOnceAsync — ends the stream (which ends the consume loop), deletes
 134        // recovery state, and tears down the timeout, exactly once.
 3135        int cleanupStarted = 0;
 3136        int teardownBudgetSpent = 0;
 3137        var subscriptionTornDown = false;
 3138        var cleanupGate = new object();
 3139        Task? cleanupTask = null;
 3140        var streamEndGate = new object();
 3141        Task? streamEndTask = null;
 3142        var consumeLoop = Task.CompletedTask;
 143
 144        // The ONE place the server-side subscription is disposed — the drain and the latched
 145        // cleanup both need the stream ended (whichever runs first), and having each dispose it
 146        // independently doubled the teardown for no benefit. TASK-latched and NEVER-faulting:
 147        // its failure is logged here exactly once, no matter how many latched callers observe
 148        // the task — and a caller that abandoned its bounded wait still gets the late outcome
 149        // recorded instead of it dying as a TaskScheduler.UnobservedTaskException. Callers read
 150        // "completed with subscriptionTornDown false" as teardown failure and backstop-cancel
 151        // (the cleanup core's finally, safe only after the timeout registration is gone).
 152        Task EndStreamOnce()
 153        {
 3154            lock (streamEndGate)
 155            {
 3156                return streamEndTask ??= EndStreamCoreAsync();
 157            }
 3158        }
 159
 160        async Task EndStreamCoreAsync()
 161        {
 3162            if (subscription is null)
 1163                return;
 164
 165            try
 166            {
 3167                await subscription.DisposeAsync().ConfigureAwait(false);
 3168                subscriptionTornDown = true;
 3169                _logger.LogDebug("Unsubscribed from subject {Subject}.", subject);
 3170            }
 3171            catch (Exception teardownEx)
 172            {
 3173                _logger.LogError(teardownEx, "Error during cleanup for subject {Subject}.", subject);
 3174            }
 3175        }
 176
 177        // Task-latched so EVERY caller completes only when the one real cleanup has finished —
 178        // the previous fire-once int latch let a second caller (a disposing waiter racing the
 179        // timeout) return before the task was settled. The core itself never waits on the consume
 180        // loop: draining happens BEFORE the latch (DrainThenCleanupAsync), because the loop's own
 181        // finally also enters this latch — a join inside the core would make the loop await a core
 182        // that is joining the loop.
 183        ValueTask CleanupOnceAsync()
 184        {
 185            Task task;
 3186            lock (cleanupGate)
 187            {
 3188                task = cleanupTask ??= CleanupCoreAsync();
 3189            }
 190
 3191            return task.IsCompletedSuccessfully ? ValueTask.CompletedTask : new ValueTask(task);
 192        }
 193
 194        // Dispose-path cleanup: DRAINS the in-flight delivery before settling. The consume loop
 195        // may be mid Until-predicate holding a claimed terminal message; ending the stream and
 196        // joining the loop guarantees that by the time the latched core cancels, the task is
 197        // either settled by that delivery or genuinely undelivered. Never called from the loop
 198        // itself — loop-invoked cleanup uses CleanupOnceAsync directly, its task already settled
 199        // by the terminal dispatch.
 200        //
 201        // One DisposalDrainTimeout budget covers BOTH steps — a wedged client library can hang
 202        // the subscription dispose just as a wedged Until predicate can hang the loop join. The
 203        // core's cancel is only truthful once the JOIN below has proven the loop ended; any
 204        // drain outcome short of that — budget lapse, anything unforeseen — leaves a delivery
 205        // possibly mid-predicate holding a message already consumed from the stream, and
 206        // "canceled" would tell a re-attaching caller nothing was delivered. Those paths fault
 207        // the task with the explicit indeterminate contract instead (routing durable flows to a
 208        // fresh idempotent restart) and cancel the subscription token so the loop still ends
 209        // once the predicate returns. (A teardown FAILURE no longer throws — the latched
 210        // teardown logs it and leaves subscriptionTornDown false — so it backstop-cancels and
 211        // still proves settlement through the join.)
 212        async ValueTask DrainThenCleanupAsync()
 213        {
 3214            if (Volatile.Read(ref cleanupStarted) == 0)
 215            {
 3216                var drainTimeout = _options.DisposalDrainTimeout;
 217                // ONE budget for the whole disposal: the latched cleanup below skips its own
 218                // teardown wait when a drain already spent this budget on the same latched
 219                // teardown task — a second full wait there made disposal cost double the
 220                // configured DisposalDrainTimeout.
 3221                Volatile.Write(ref teardownBudgetSpent, 1);
 222                try
 223                {
 3224                    using var budget = new CancellationTokenSource(drainTimeout);
 3225                    await EndStreamOnce().WaitAsync(budget.Token).ConfigureAwait(false);
 226
 227                    // A failed teardown surfaces as "completed, subscriptionTornDown false" (the
 228                    // latched core logged it): backstop-cancel so the loop still ends, then FALL
 229                    // THROUGH to the join — a failed teardown proves nothing about a delivery
 230                    // mid-predicate, and skipping the join here let cleanup cancel a response the
 231                    // stream had already handed over.
 3232                    if (!subscriptionTornDown && subscription is not null)
 2233                        await DisarmThenCancelSubscriptionTokenAsync().ConfigureAwait(false);
 234
 235                    // Settlement is PROVEN only by the loop having ended within the remaining
 236                    // budget — either it settled the task with the in-flight delivery, or it
 237                    // ended with nothing in flight and the cleanup's cancel below is truthful.
 3238                    await consumeLoop.WaitAsync(budget.Token).ConfigureAwait(false);
 3239                }
 3240                catch (Exception drainEx)
 241                {
 3242                    _logger.LogWarning(
 3243                        "Disposal drain for correlationId {CorrelationId} did not prove settlement within {DrainTimeout}
 3244                        correlationId, drainTimeout);
 3245                    AsyncResponseDiagnostics.SetError(activity, "indeterminate_delivery", "Disposal drain did not prove 
 246                    // A TrySetResult from the late-finishing delivery loses against this and is
 247                    // dropped; the loop's own cleanup call is a no-op behind the latch. The
 248                    // non-cancellation exception case is unforeseen infrastructure failure —
 249                    // settlement is equally unproven there, so it must not fall back to cancel.
 3250                    if (drainEx is not OperationCanceledException)
 1251                        _logger.LogDebug(drainEx, "Disposal drain failed for subject {Subject}.", subject);
 3252                    tcs.TrySetException(new AsyncResponseIndeterminateDeliveryException(correlationId, drainTimeout));
 3253                    await DisarmThenCancelSubscriptionTokenAsync().ConfigureAwait(false);
 254                }
 255            }
 256
 3257            await CleanupOnceAsync().ConfigureAwait(false);
 258        }
 259
 260        async ValueTask DisarmThenCancelSubscriptionTokenAsync()
 261        {
 262            // Disarm the waiter-timeout registration BEFORE the backstop cancel — the cancel
 263            // would otherwise fire it and stamp a spurious TimeoutException plus a waiter-timeout
 264            // metric onto a disposal that is not a timeout. Idempotent with the cleanup core's
 265            // own registration disposal.
 3266            await timeoutRegistration.DisposeAsync().ConfigureAwait(false);
 267            try
 268            {
 3269                cancellationTokenSource.Cancel();
 3270            }
 1271            catch (ObjectDisposedException)
 272            {
 273                // Cleanup already ran and disposed the source; the loop is ending regardless.
 1274            }
 3275        }
 276
 277        async Task CleanupCoreAsync()
 278        {
 3279            Interlocked.Exchange(ref cleanupStarted, 1);
 280
 281            try
 282            {
 283                try
 284                {
 285                    // Delete the recovery state BEFORE disposing the subscription. In the reverse
 286                    // order a publish landing in the window sees "no responders, state present" and
 287                    // fires a spurious recovery callback for a wait that already reached a terminal
 288                    // state. In this order the window shows a subscriber that drops the message — a
 289                    // late or duplicate terminal message is droppable; a resurrected recovery callback
 290                    // is not.
 3291                    await _recoveryStateStore.TryDeleteAsync(correlationId, registrationId).ConfigureAwait(false);
 3292                }
 3293                catch (Exception ex)
 294                {
 295                    // Best-effort: the KV entry expires on its own, and a transient store failure
 296                    // must not skip the subscription teardown below.
 3297                    _logger.LogError(ex, "Failed to delete recovery state for correlationId {CorrelationId}.", correlati
 3298                }
 299
 300                // End the stream if the drain has not already — dispatch-triggered cleanup (a
 301                // terminal delivery, a loop fault) reaches here without a drain. The task-latch
 302                // keeps the teardown single no matter which path got here first. Bounded like the
 303                // drain: this latched core is what a disposing waiter awaits when terminal
 304                // delivery started cleanup first (the drain skips itself on cleanupStarted), so
 305                // an unbudgeted teardown here let a wedged client library hold DisposeAsync
 306                // hostage past DisposalDrainTimeout. On the bound lapsing, the catch below logs
 307                // and the finally's backstop cancel still ends the consume loop; the abandoned
 308                // teardown task never faults (it logs its own late outcome). When a DRAIN
 309                // preceded this core, it already spent that budget on this same latched task —
 310                // waiting a second one here made disposal cost double the configured bound, so
 311                // the wait is skipped (the teardown is running and self-logging regardless).
 3312                if (Volatile.Read(ref teardownBudgetSpent) == 0)
 3313                    await EndStreamOnce().WaitAsync(_options.DisposalDrainTimeout).ConfigureAwait(false);
 3314                if (subscription is null)
 1315                    subscriptionTornDown = true;
 3316            }
 3317            catch (Exception ex)
 318            {
 3319                _logger.LogError(ex, "Error during cleanup for subject {Subject}.", subject);
 3320            }
 321            finally
 322            {
 3323                await timeoutRegistration.DisposeAsync().ConfigureAwait(false);
 3324                if (!subscriptionTornDown)
 325                {
 326                    // DisposeAsync did not complete, so the server-side subscription may still be
 327                    // pumping messages. Its lifetime is bound to this token (SubscribeAsync received
 328                    // it), and disposing a CTS never cancels — an explicit cancel is the backstop
 329                    // that ends the consume loop. Safe only after the timeout registration above is
 330                    // gone, or the cancel would fire a spurious waiter timeout.
 3331                    cancellationTokenSource.Cancel();
 332                }
 333
 334                // A waiter disposed before any terminal signal must not leave ResponseTask pending
 335                // forever for callers that hold it directly — the timeout died above, so nothing
 336                // else could ever complete the task. A no-op after a normal completion, timeout,
 337                // fault, or a delivery drained by DrainThenCleanupAsync.
 3338                tcs.TrySetCanceled();
 339
 3340                cancellationTokenSource.Dispose();
 3341                activity?.Dispose();
 342            }
 343        }
 344
 345        // -------------------------------------------------------------------------
 346        // Local: ProcessResponseAsync — deserializes and handles a single envelope, completes the TCS when terminal.
 347        async Task ProcessResponseAsync(string? payload)
 348        {
 3349            bool finished = false;
 350            try
 351            {
 3352                if (string.IsNullOrEmpty(payload))
 353                {
 354                    // A non-probe message with no body cannot be a response; ignore it rather than fault.
 3355                    _logger.LogWarning("Received empty response message for correlationId {CorrelationId}; ignoring.", c
 3356                    return;
 357                }
 358
 3359                var envelope = JsonSerializer.Deserialize(payload, AsyncResponseEnvelopeJson.TypeInfo<T>());
 360
 3361                if (envelope == null)
 362                {
 3363                    _logger.LogError("Failed to deserialize envelope for correlationId {CorrelationId}.", correlationId)
 2364                    finished = true;
 2365                    var deserializationError = new JsonException($"Failed to deserialize envelope for correlationId {cor
 2366                    AsyncResponseDiagnostics.SetError(activity, "deserialize_failure", deserializationError.Message);
 2367                    if (!tcs.TrySetException(deserializationError))
 2368                        _logger.LogWarning(deserializationError, "TaskCompletionSource already completed for correlation
 369                }
 3370                else if (!AsyncResponseEnvelopeSchema.IsReadable(envelope.SchemaVersion))
 371                {
 3372                    finished = true;
 3373                    var schemaError = new InvalidOperationException(
 3374                        $"Response envelope for correlationId {correlationId} has schema version {envelope.SchemaVersion
 3375                        $"which this build does not support (current: {AsyncResponseEnvelopeSchema.Current}).");
 2376                    AsyncResponseDiagnostics.SetError(activity, "schema_mismatch", schemaError.Message);
 2377                    if (!tcs.TrySetException(schemaError))
 2378                        _logger.LogWarning(schemaError, "TaskCompletionSource already completed for correlationId {Corre
 379                }
 3380                else if (!envelope.Success)
 381                {
 3382                    finished = true;
 3383                    var remoteFailure = new Exception(envelope.ExceptionMessage ?? "Unknown error during asynchronous pr
 3384                    if (!string.IsNullOrEmpty(envelope.ExceptionStackTrace))
 385                        // Cap on receive too: the publish-side cap only bounds traces we emit, not what
 386                        // a remote we do not control can push at us.
 3387                        remoteFailure.Data["RemoteStackTrace"] = RemoteStackTrace.Cap(envelope.ExceptionStackTrace, _opt
 388
 3389                    _logger.LogWarning("Received error response for correlationId {CorrelationId}: {ErrorMessage}", corr
 3390                    AsyncResponseDiagnostics.SetError(activity, "remote_failure", remoteFailure.Message);
 3391                    if (!tcs.TrySetException(remoteFailure))
 3392                        _logger.LogWarning(remoteFailure, "TaskCompletionSource already completed for correlationId {Cor
 393                }
 394                else
 395                {
 3396                    if (_logger.IsEnabled(LogLevel.Debug))
 3397                        _logger.LogDebug("Received response for correlationId {CorrelationId}.", correlationId);
 3398                    finished = await completionPredicate(envelope.Payload!).ConfigureAwait(false);
 3399                    if (finished && !tcs.TrySetResult(envelope.Payload!))
 3400                        _logger.LogWarning("TaskCompletionSource already completed for correlationId {CorrelationId}.", 
 401                }
 3402            }
 3403            catch (Exception ex)
 404            {
 3405                _logger.LogError(ex, "Error processing message on subject {Subject} for correlationId {CorrelationId}.",
 2406                finished = true;
 2407                AsyncResponseDiagnostics.SetError(activity, ex);
 2408                if (!tcs.TrySetException(ex))
 2409                    _logger.LogWarning(ex, "TaskCompletionSource already completed for correlationId {CorrelationId}.", 
 3410            }
 411            finally
 412            {
 3413                if (finished)
 3414                    await CleanupOnceAsync().ConfigureAwait(false);
 415            }
 416        }
 417
 418        // -------------------------------------------------------------------------
 419        // Local: ProcessUnderCapturedContextAsync — restores the waiter's subscribe-time
 420        // ExecutionContext (app AsyncLocals) plus the correlation id before processing, since the
 421        // consume loop runs on a background thread that never had them.
 422        Task ProcessUnderCapturedContextAsync(string? payload)
 423        {
 424            async Task Process()
 425            {
 3426                using var correlationScope = AsyncResponseContext.PushCorrelationId(storedCorrelationId);
 3427                await ProcessResponseAsync(payload).ConfigureAwait(false);
 3428            }
 429
 3430            if (capturedContext is null)
 3431                return Process();
 432
 3433            Task? task = null;
 3434            ExecutionContext.Run(capturedContext, _ => task = Process(), null);
 3435            return task!;
 436        }
 437
 438        // -------------------------------------------------------------------------
 439        // Local: ConsumeLoopAsync — reads messages serially from the subscription until it is disposed.
 440        async Task ConsumeLoopAsync(INatsChannelSubscription sub)
 441        {
 442            try
 443            {
 3444                await foreach (var message in sub.ReadAsync(CancellationToken.None).ConfigureAwait(false))
 445                {
 446                    // Ack first so the publisher's request resolves quickly (delivery/liveness confirmed)
 447                    // even if processing the payload is slow. A failed ack must not abort the wait.
 448                    try
 449                    {
 3450                        await message.ReplyAsync().ConfigureAwait(false);
 3451                    }
 3452                    catch (Exception replyEx)
 453                    {
 3454                        _logger.LogDebug(replyEx, "Failed to acknowledge response on subject {Subject}.", subject);
 3455                    }
 456
 3457                    if (message.IsProbe)
 458                        continue;
 459
 3460                    await ProcessUnderCapturedContextAsync(message.Payload).ConfigureAwait(false);
 3461                }
 3462            }
 3463            catch (Exception ex)
 464            {
 3465                _logger.LogError(ex, "Response subscription loop failed for subject {Subject}.", subject);
 2466                AsyncResponseDiagnostics.SetError(activity, ex);
 2467                if (!tcs.TrySetException(ex))
 2468                    _logger.LogWarning(ex, "TaskCompletionSource already completed for correlationId {CorrelationId}.", 
 2469                await CleanupOnceAsync().ConfigureAwait(false);
 470            }
 471        }
 472
 3473        timeoutRegistration = cancellationTokenSource.Token.Register(() =>
 3474        {
 3475            _ = Task.Run(async () =>
 3476            {
 3477                _logger.LogWarning("Timed out waiting for response for correlationId {CorrelationId}.", correlationId);
 3478                AsyncResponseDiagnostics.SetError(activity, "timeout", $"Timed out waiting for response for correlationI
 3479                AsyncResponseDiagnostics.RecordWaiterTimeout("nats");
 3480                tcs.TrySetException(new TimeoutException($"Timed out waiting for response for correlationId {correlation
 3481                await DrainThenCleanupAsync().ConfigureAwait(false);
 3482            });
 3483        });
 484
 485        try
 486        {
 3487            subscription = await _client.SubscribeAsync(subject, cancellationTokenSource.Token).ConfigureAwait(false);
 3488            consumeLoop = Task.Run(() => ConsumeLoopAsync(subscription));
 489
 3490            var recoveryState = new RecoveryState
 3491            {
 3492                RegistrationId = registrationId,
 3493                ResumeCallback = resumeCallback,
 3494                FailureCallback = failureCallback,
 3495                CorrelationId = correlationId,
 3496                PayloadTypeFullName = typeof(T).FullName,
 3497                RegisteredAtUtc = DateTime.UtcNow,
 3498                Context = _propagation.Capture()
 3499            };
 3500            await _recoveryStateStore.SaveAsync(correlationId, recoveryState, _options.RecoveryStateExpiry).ConfigureAwa
 501
 502            // Round-trip to the server so the subscription is guaranteed registered before the caller's
 503            // trigger publishes the remote request — closing the subscribe/trigger race.
 3504            await _client.FlushAsync(cancellationTokenSource.Token).ConfigureAwait(false);
 505
 3506            _logger.LogDebug("Subscribed to subject {Subject} for correlationId {CorrelationId}.", subject, correlationI
 3507        }
 3508        catch (Exception ex)
 509        {
 3510            _logger.LogError(ex, "Failed to subscribe to subject {Subject} for correlationId {CorrelationId}.", subject,
 2511            AsyncResponseDiagnostics.SetError(activity, "subscribe_failure", ex.Message);
 2512            await DrainThenCleanupAsync().ConfigureAwait(false);
 513
 514            // Rethrow instead of returning a pre-faulted waiter: the builder's contract is that
 515            // the trigger runs only once the subscription AND recovery state exist. A returned
 516            // waiter would still let the trigger fire the remote operation with no registration
 517            // left to receive (or recover) its response. Cleanup leaves nothing behind and cancels
 518            // the response task rather than faulting it, so no unobserved fault lingers.
 3519            throw;
 520        }
 521
 522        try
 523        {
 3524            if (Volatile.Read(ref cleanupStarted) == 0)
 3525                cancellationTokenSource.CancelAfter(timeout.Value);
 3526        }
 1527        catch (ObjectDisposedException)
 528        {
 529            // A response completed and cleaned up between the check and CancelAfter.
 1530        }
 531
 3532        return new NatsAsyncResponseWaiter<T>(tcs.Task, DrainThenCleanupAsync);
 3533    }
 534
 535    // ---------------------------------------------------------------------------------------
 536    // IAsyncResponsePublisher
 537
 538    /// <inheritdoc/>
 539    public Task SetResponse<T>(T response, string correlationId, CancellationToken cancellationToken = default) where T 
 3540        => SetResponseCore(response, correlationId, cancellationToken);
 541
 542    Task IRawAsyncResponsePublisher.SetRawResponse(object? response, string correlationId, CancellationToken cancellatio
 3543        => SetResponseCore(response, correlationId, cancellationToken);
 544
 545    Task IRawAsyncResponsePublisher.SetRawResponseJson(string responseJson, string correlationId, CancellationToken canc
 3546        => SetRawResponseJsonCore(responseJson, correlationId, cancellationToken);
 547
 548    private async Task SetResponseCore<T>(T response, string correlationId, CancellationToken cancellationToken)
 549    {
 3550        using var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.set_response", ActivityKind.Producer)
 3551        activity?.SetTag("asyncresponse.channel", "nats");
 3552        AsyncResponseDiagnostics.SetPayloadType(activity, typeof(T));
 553
 3554        AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 555
 3556        if (string.IsNullOrWhiteSpace(correlationId))
 557        {
 3558            _logger.LogWarning("CorrelationId is null; cannot publish the response.");
 2559            AsyncResponseDiagnostics.SetError(activity, "correlation_id_null", "CorrelationId is null; cannot publish th
 3560            return;
 561        }
 562
 3563        var subject = _subjects.ResponseSubject(correlationId);
 564        try
 565        {
 3566            var envelope = new AsyncResponseEnvelope<T> { Success = true, Payload = response };
 3567            var json = AsyncResponseEnvelopeJson.Serialize(envelope);
 3568            var outcome = await _client.RequestAsync(subject, json, probe: false, _options.DeliveryConfirmationTimeout, 
 3569            activity?.SetTag("asyncresponse.delivery", outcome.ToString());
 570
 3571            if (outcome == NatsDeliveryOutcome.NoResponders)
 572            {
 573                // Nobody was listening (the waiter died, e.g. with a redeploy): hand the response over
 574                // to the lost-subscriber dispatcher, which asks the payload whether to resume or fail.
 3575                var dispatchResult = await _lostSubscriberDispatcher
 3576                    .DispatchLostResponses(
 3577                        _recoveryStateStore,
 3578                        correlationId,
 3579                        response,
 3580                        subject,
 3581                        cancellationToken,
 3582                        hasLiveSubscriber: () => HasLiveSubscriberAsync(correlationId, cancellationToken))
 3583                    .ConfigureAwait(false);
 3584                if (dispatchResult.RetryLive)
 585                {
 586                    // A waiter subscribed between the request and the recovery-state read —
 587                    // re-attempt the live publish instead of consuming its registration; only a
 588                    // second no-responders consumes it.
 3589                    outcome = await _client.RequestAsync(subject, json, probe: false, _options.DeliveryConfirmationTimeo
 2590                    activity?.SetTag("asyncresponse.delivery", outcome.ToString());
 2591                    if (outcome != NatsDeliveryOutcome.NoResponders)
 2592                        return;
 593
 2594                    dispatchResult = await _lostSubscriberDispatcher
 2595                        .DispatchLostResponses(_recoveryStateStore, correlationId, response, subject, cancellationToken)
 2596                        .ConfigureAwait(false);
 597                }
 598
 3599                AsyncResponseDiagnostics.SetLostSubscriberRoute(activity, dispatchResult.ShouldResume);
 3600                AsyncResponseDiagnostics.RecordLostSubscriber("response", dispatchResult.ShouldResume, dispatchResult.Ca
 3601                activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 602            }
 3603            else if (_logger.IsEnabled(LogLevel.Debug))
 604            {
 3605                _logger.LogDebug("Published response for correlationId {CorrelationId} on subject {Subject}. PayloadType
 606            }
 3607        }
 3608        catch (Exception ex)
 609        {
 3610            _logger.LogError(ex, "Failed to publish response for correlationId {CorrelationId} on subject {Subject}.", c
 2611            AsyncResponseDiagnostics.SetError(activity, ex);
 3612            throw;
 613        }
 3614    }
 615
 616    private async Task SetRawResponseJsonCore(string responseJson, string correlationId, CancellationToken cancellationT
 617    {
 3618        using var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.ingress.raw_response", ActivityKind.P
 3619        activity?.SetTag("asyncresponse.channel", "nats");
 620
 3621        AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 622
 3623        if (string.IsNullOrWhiteSpace(correlationId))
 624        {
 3625            _logger.LogWarning("CorrelationId is null; cannot publish the raw response.");
 2626            AsyncResponseDiagnostics.SetError(activity, "correlation_id_null", "CorrelationId is null; cannot publish th
 3627            return;
 628        }
 629
 3630        var subject = _subjects.ResponseSubject(correlationId);
 631        try
 632        {
 3633            var json = SerializeRawSuccessEnvelope(responseJson);
 3634            var outcome = await _client.RequestAsync(subject, json, probe: false, _options.DeliveryConfirmationTimeout, 
 3635            activity?.SetTag("asyncresponse.delivery", outcome.ToString());
 636
 3637            if (outcome == NatsDeliveryOutcome.NoResponders)
 638            {
 3639                var response = new RawJsonResponse(responseJson).DeserializeUntyped();
 640
 2641                var dispatchResult = await _lostSubscriberDispatcher
 2642                    .DispatchLostResponses(
 2643                        _recoveryStateStore,
 2644                        correlationId,
 2645                        response,
 2646                        subject,
 2647                        cancellationToken,
 3648                        hasLiveSubscriber: () => HasLiveSubscriberAsync(correlationId, cancellationToken))
 2649                    .ConfigureAwait(false);
 2650                if (dispatchResult.RetryLive)
 651                {
 652                    // A waiter subscribed between the request and the recovery-state read —
 653                    // re-attempt the live publish instead of consuming its registration; only a
 654                    // second no-responders consumes it.
 2655                    outcome = await _client.RequestAsync(subject, json, probe: false, _options.DeliveryConfirmationTimeo
 2656                    activity?.SetTag("asyncresponse.delivery", outcome.ToString());
 2657                    if (outcome != NatsDeliveryOutcome.NoResponders)
 2658                        return;
 659
 2660                    dispatchResult = await _lostSubscriberDispatcher
 2661                        .DispatchLostResponses(_recoveryStateStore, correlationId, response, subject, cancellationToken)
 2662                        .ConfigureAwait(false);
 663                }
 664
 2665                AsyncResponseDiagnostics.SetLostSubscriberRoute(activity, dispatchResult.ShouldResume);
 2666                AsyncResponseDiagnostics.RecordLostSubscriber("response", dispatchResult.ShouldResume, dispatchResult.Ca
 2667                activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 3668            }
 3669            else if (_logger.IsEnabled(LogLevel.Debug))
 670            {
 3671                _logger.LogDebug("Published raw response for correlationId {CorrelationId} on subject {Subject}. Outcome
 672            }
 3673        }
 3674        catch (Exception ex)
 675        {
 3676            _logger.LogError(ex, "Failed to publish raw response for correlationId {CorrelationId} on subject {Subject}.
 2677            AsyncResponseDiagnostics.SetError(activity, ex);
 3678            throw;
 679        }
 3680    }
 681
 682    /// <inheritdoc/>
 683    public async Task SetException(Exception exception, string correlationId, CancellationToken cancellationToken = defa
 684    {
 3685        ArgumentNullException.ThrowIfNull(exception);
 686
 3687        using var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.set_exception", ActivityKind.Producer
 3688        activity?.SetTag("asyncresponse.channel", "nats");
 3689        activity?.SetTag("asyncresponse.exception_type", exception.GetType().FullName ?? exception.GetType().Name);
 690
 3691        AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 692
 3693        if (string.IsNullOrWhiteSpace(correlationId))
 694        {
 3695            _logger.LogWarning("CorrelationId is null; cannot publish the exception. Exception: {ExceptionMessage}", exc
 2696            AsyncResponseDiagnostics.SetError(activity, "correlation_id_null", "CorrelationId is null; cannot publish th
 2697            return;
 698        }
 699
 3700        var subject = _subjects.ResponseSubject(correlationId);
 701        try
 702        {
 3703            var envelope = new AsyncResponseEnvelope<object>
 3704            {
 3705                Success = false,
 3706                ExceptionMessage = exception.Message,
 3707                ExceptionStackTrace = RemoteStackTrace.ForWire(exception.StackTrace, _options.IncludeRemoteStackTrace, _
 3708                Payload = null
 3709            };
 3710            var json = AsyncResponseEnvelopeJson.Serialize(envelope);
 3711            var outcome = await _client.RequestAsync(subject, json, probe: false, _options.DeliveryConfirmationTimeout, 
 3712            activity?.SetTag("asyncresponse.delivery", outcome.ToString());
 713
 3714            if (outcome == NatsDeliveryOutcome.NoResponders)
 715            {
 716                // Nobody was listening: exception envelopes always go to the failure callback.
 3717                var dispatchResult = await _lostSubscriberDispatcher
 3718                    .DispatchLostExceptions(
 3719                        _recoveryStateStore,
 3720                        correlationId,
 3721                        exception,
 3722                        subject,
 3723                        cancellationToken,
 3724                        hasLiveSubscriber: () => HasLiveSubscriberAsync(correlationId, cancellationToken))
 3725                    .ConfigureAwait(false);
 3726                if (dispatchResult.RetryLive)
 727                {
 728                    // A waiter subscribed between the request and the recovery-state read —
 729                    // re-attempt the live publish instead of consuming its registration; only a
 730                    // second no-responders consumes it.
 3731                    outcome = await _client.RequestAsync(subject, json, probe: false, _options.DeliveryConfirmationTimeo
 2732                    activity?.SetTag("asyncresponse.delivery", outcome.ToString());
 2733                    if (outcome != NatsDeliveryOutcome.NoResponders)
 2734                        return;
 735
 2736                    dispatchResult = await _lostSubscriberDispatcher
 2737                        .DispatchLostExceptions(_recoveryStateStore, correlationId, exception, subject, cancellationToke
 2738                        .ConfigureAwait(false);
 739                }
 740
 3741                activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 3742                AsyncResponseDiagnostics.RecordLostSubscriber("exception", shouldResume: false, dispatchResult.CallbackI
 743            }
 3744            else if (_logger.IsEnabled(LogLevel.Debug))
 745            {
 3746                _logger.LogDebug("Published exception response for correlationId {CorrelationId} on subject {Subject}. O
 747            }
 3748        }
 3749        catch (Exception ex)
 750        {
 3751            _logger.LogError(ex, "Failed to publish exception response for correlationId {CorrelationId} on subject {Sub
 2752            AsyncResponseDiagnostics.SetError(activity, ex);
 2753            throw;
 754        }
 3755    }
 756
 757    // ---------------------------------------------------------------------------------------
 758    // IActiveSubscriberProbe
 759
 760    /// <inheritdoc/>
 761    public async ValueTask<long> CountActiveSubscribersAsync(string correlationId, CancellationToken cancellationToken =
 762    {
 3763        if (string.IsNullOrWhiteSpace(correlationId))
 3764            return 0L;
 765
 3766        var subject = _subjects.ResponseSubject(correlationId);
 767        try
 768        {
 769            // NATS Core does not expose exact subscriber counts to clients, so the probe reports
 770            // presence: a live waiter answers the ping (1), no-responders or no timely answer means
 771            // none (0). The watchdog only needs "is anyone listening".
 3772            var outcome = await _client.RequestAsync(subject, payload: null, probe: true, _options.PresenceProbeTimeout,
 3773            return outcome == NatsDeliveryOutcome.Replied ? 1L : 0L;
 774        }
 2775        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
 776        {
 2777            throw;
 778        }
 2779        catch (Exception ex)
 780        {
 2781            _logger.LogDebug(ex, "Failed to probe active subscribers for subject {Subject}.", subject);
 2782            return 0L;
 783        }
 3784    }
 785
 786    /// <summary>
 787    /// Re-probes waiter liveness for the lost-subscriber dispatcher's snapshot-race re-check,
 788    /// using the same presence probe the watchdog uses.
 789    /// </summary>
 790    private async ValueTask<bool> HasLiveSubscriberAsync(string correlationId, CancellationToken cancellationToken)
 3791        => await CountActiveSubscribersAsync(correlationId, cancellationToken).ConfigureAwait(false) > 0;
 792
 793    private static string SerializeRawSuccessEnvelope(string payloadJson)
 794    {
 3795        JsonSafety.ThrowIfClearlyNotJson(payloadJson);
 796
 3797        var buffer = new ArrayBufferWriter<byte>();
 3798        using (var writer = new Utf8JsonWriter(buffer))
 799        {
 3800            writer.WriteStartObject();
 3801            writer.WriteNumber("SchemaVersion", AsyncResponseEnvelopeSchema.Current);
 3802            writer.WriteBoolean("Success", true);
 3803            writer.WritePropertyName("Payload");
 3804            writer.WriteRawValue(payloadJson);
 3805            writer.WriteNull("ExceptionMessage");
 3806            writer.WriteNull("ExceptionStackTrace");
 3807            writer.WriteEndObject();
 3808        }
 809
 3810        return Encoding.UTF8.GetString(buffer.WrittenSpan);
 811    }
 812}