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

Information
Class: AsyncResponse.Channels.DbAsyncResponseChannelBase.DbSubscription<T>
Assembly: AsyncResponse.Channels.SqlServer
File(s): /home/runner/work/AsyncResponse/AsyncResponse/src/Channels/Shared/DbChannelShared.cs
Line coverage
100%
Covered lines: 96
Uncovered lines: 0
Coverable lines: 96
Total lines: 1386
Line coverage: 100%
Branch coverage
100%
Covered branches: 32
Total branches: 32
Branch coverage: 100%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
ProcessAsync()100%1616100%
DrainThenCleanupAsync()100%66100%
CleanupCoreAsync()100%1010100%
<CleanupCoreAsync()100%11100%
DropLocalAsync()100%11100%

File(s)

/home/runner/work/AsyncResponse/AsyncResponse/src/Channels/Shared/DbChannelShared.cs

#LineLine coverage
 1using Microsoft.Extensions.DependencyInjection;
 2using Microsoft.Extensions.Logging;
 3using System.Buffers;
 4using System.Collections.Concurrent;
 5using System.Diagnostics;
 6using System.Text;
 7using System.Text.Json;
 8using System.Threading.Channels;
 9
 10namespace AsyncResponse.Channels;
 11
 12// Shared source for the database-backed response channels (PostgreSQL, SQL Server, MongoDB),
 13// mirroring the DurableFlows shared-store pattern: each channel csproj pulls this file in via
 14// <Compile Include="..\Shared\DbChannelShared.cs" />, so the base class compiles INTO each
 15// provider assembly against that provider's concrete seam types. The seam is bound per project
 16// with three global using aliases (declared at the top of the provider's channel file):
 17//
 18//   DbChannelStore   -> the provider's store/SQL adapter (e.g. PostgreSqlChannelSql)
 19//   DbChannelMessage -> the provider's channel-message record (e.g. PostgreSqlChannelMessage)
 20//   DbChannelOptions -> the provider's options class (e.g. PostgreSqlAsyncResponseChannelOptions)
 21//
 22// Because the aliases resolve to concrete sealed types at compile time, store calls on the
 23// per-message paths stay direct (no interface dispatch, no delegate indirection) — see the
 24// benchmark note in RedisAsyncResponseChannel.SetResponseCore for why that matters. The only
 25// virtual seams are the four hooks below, which cover exactly what the three providers genuinely
 26// do differently: the channel-name format, the sweep cadence, the optional wake listener, and the
 27// provider waiter type.
 28
 29/// <summary>
 30/// Provider-agnostic machinery for the database-backed response channels: waiter registration and
 31/// recovery-state bookkeeping, publish with delivery confirmation, the signal-driven dispatch
 32/// sweep, the subscriber heartbeat, and subscription lifecycle/cleanup. Derived channels supply
 33/// the wake mechanism (LISTEN/NOTIFY, adaptive polling, change streams), the channel-name format,
 34/// and the provider waiter type via the protected hooks.
 35/// </summary>
 36internal abstract class DbAsyncResponseChannelBase :
 37    IAsyncResponsePublisher,
 38    IRawAsyncResponsePublisher,
 39    IRecoverableAsyncResponseSubscriber,
 40    IActiveSubscriberProbe,
 41    IAsyncDisposable
 42{
 43    private protected readonly ConcurrentDictionary<string, ConcurrentDictionary<Guid, IDbSubscription>> _subscriptions 
 44
 45    // A signal carries the correlation id to scan (targeted), or null to scan every subscribed
 46    // correlation id (the periodic sweep that is the missed-wake safety net).
 47    private readonly Channel<string?> _signals = Channel.CreateBounded<string?>(new BoundedChannelOptions(1024)
 48    {
 49        SingleReader = true,
 50        SingleWriter = false,
 51        FullMode = BoundedChannelFullMode.DropOldest
 52    });
 53
 54    // Maps a just-published message id to a completion the local dispatch loop trips the instant it
 55    // delivers the message to a live waiter. Same-process delivery (the overwhelmingly common case)
 56    // is confirmed without polling the database; cross-process delivery falls back to polling acked_at.
 57    private readonly ConcurrentDictionary<Guid, TaskCompletionSource<bool>> _pendingConfirmations = new();
 58
 59    private protected readonly DbChannelStore _store;
 60    private readonly IRecoveryStateStore _recoveryStateStore;
 61    private readonly AsyncResponseContextPropagation _propagation;
 62    private readonly LostSubscriberCallbackDispatcher _lostSubscriberDispatcher;
 63    private protected readonly DbChannelOptions _options;
 64    private protected readonly ILogger _logger;
 65    private readonly SerialExecutorRegistry _executors;
 66    private readonly string _instanceId = $"{Environment.MachineName}-{Environment.ProcessId}-{Guid.NewGuid():N}";
 67
 68    // Provider text used in diagnostics. The emitted strings must stay byte-identical to the
 69    // pre-consolidation per-provider channels — tests and dashboards match on them.
 70    private readonly string _channelTypeName;
 71    private readonly string _providerName;
 72    private readonly string _activityTag;
 73    private readonly string _subscriberRecordNoun;
 74    private readonly string _localDispatchRetryHint;
 75
 76    private readonly object _listenerGate = new();
 77    private protected CancellationTokenSource? _listenerCts;
 78    private protected Task? _listenTask;
 79    private protected Task? _dispatchTask;
 80    private protected Task? _heartbeatTask;
 81    private bool _disposed;
 82
 83    /// <summary>Creates the shared machinery for a database-backed async-response channel.</summary>
 84    protected DbAsyncResponseChannelBase(
 85        IServiceScopeFactory scopeFactory,
 86        DbChannelStore store,
 87        IRecoveryStateStore recoveryStateStore,
 88        DbChannelOptions options,
 89        AsyncResponseContextPropagation propagation,
 90        ILogger logger,
 91        string channelTypeName,
 92        string providerName,
 93        string activityTag,
 94        string subscriberRecordNoun,
 95        string localDispatchRetryHint)
 96    {
 97        _store = store;
 98        _recoveryStateStore = recoveryStateStore;
 99        _propagation = propagation;
 100        _options = options;
 101        _options.Validate();
 102        _logger = logger;
 103        _channelTypeName = channelTypeName;
 104        _providerName = providerName;
 105        _activityTag = activityTag;
 106        _subscriberRecordNoun = subscriberRecordNoun;
 107        _localDispatchRetryHint = localDispatchRetryHint;
 108        _lostSubscriberDispatcher = new LostSubscriberCallbackDispatcher(scopeFactory, propagation, logger);
 109        _executors = new SerialExecutorRegistry(logger);
 110    }
 111
 112    /// <summary>
 113    /// The per-correlation channel name used as the serial-executor key and the lost-subscriber
 114    /// channel label. Formats differ per provider (notification channel, schema.table, collection).
 115    /// </summary>
 116    protected abstract string ChannelName(string correlationId);
 117
 118    /// <summary>
 119    /// The dispatch sweep cadence. Fixed (<c>ListenerPollInterval</c>) for the providers with a push
 120    /// wake; adaptive (active/idle) for SQL Server where the sweep IS the cross-process wake.
 121    /// </summary>
 122    protected abstract TimeSpan CurrentPollInterval();
 123
 124    /// <summary>
 125    /// Starts the provider's wake listener loop (LISTEN/NOTIFY, change stream), or returns
 126    /// <c>null</c> when the provider has none and relies on the dispatch sweep alone.
 127    /// </summary>
 128    protected virtual Task? StartWakeListener(CancellationToken cancellationToken) => null;
 129
 130    /// <summary>Wraps the response task in the provider's waiter type.</summary>
 131    protected abstract IAsyncResponseWaiter<T> CreateWaiter<T>(Task<T> responseTask, Func<ValueTask> cleanupAsync)
 132        where T : IAsyncResponsePayload;
 133
 134    /// <inheritdoc />
 135    public Task<IAsyncResponseWaiter<T>> CreateResponseWaiter<T>(
 136        string correlationId,
 137        Func<T, ValueTask<bool>>? completionPredicate = null,
 138        TimeSpan? timeout = null) where T : IAsyncResponsePayload
 139        => CreateResponseWaiterCore(correlationId, null, null, completionPredicate, timeout);
 140
 141    /// <inheritdoc />
 142    public Task<IAsyncResponseWaiter<T>> CreateRecoverableResponseWaiter<T>(
 143        string correlationId,
 144        ReflectionCallDto? resumeCallback = null,
 145        ReflectionCallDto? failureCallback = null,
 146        Func<T, ValueTask<bool>>? completionPredicate = null,
 147        TimeSpan? timeout = null) where T : IAsyncResponsePayload
 148        => CreateResponseWaiterCore(correlationId, resumeCallback, failureCallback, completionPredicate, timeout);
 149
 150    private async Task<IAsyncResponseWaiter<T>> CreateResponseWaiterCore<T>(
 151        string correlationId,
 152        ReflectionCallDto? resumeCallback,
 153        ReflectionCallDto? failureCallback,
 154        Func<T, ValueTask<bool>>? completionPredicate,
 155        TimeSpan? timeout) where T : IAsyncResponsePayload
 156    {
 157        if (string.IsNullOrWhiteSpace(correlationId))
 158            throw new ArgumentNullException(nameof(correlationId), "CorrelationId must not be empty or whitespace.");
 159
 160        if ((resumeCallback is not null || failureCallback is not null)
 161            && !AsyncResponsePayloadReflection.OverridesShouldResumeOnRecovery(typeof(T)))
 162        {
 163            throw new InvalidOperationException(
 164                $"Payload type '{typeof(T)}' registers lost-subscriber recovery callbacks on the {_providerName} channel
 165                $"but does not override {nameof(IAsyncResponsePayload)}.{nameof(IAsyncResponsePayload.ShouldResumeOnReco
 166                "Override it to declare which responses resume the flow (return true) versus fail it (return false); " +
 167                "the durable channel needs this to route a response that arrives after the waiter was lost.");
 168        }
 169
 170        completionPredicate ??= _ => new ValueTask<bool>(true);
 171        timeout ??= _options.DefaultTimeout ?? _options.RecoveryStateExpiry;
 172        // BEFORE any side effect: an unsupported resolved timeout (non-positive, or past the
 173        // ~49.7-day BCL timer ceiling) used to throw only at timer arming — after the
 174        // subscription and recovery state existed, leaking both — and zero used to slip through
 175        // on some channels entirely, insta-timing-out a fully registered waiter.
 176        AsyncResponseChannelOptions.EnsureWaiterTimeoutSupported(timeout.Value);
 177
 178        await _store.EnsureCreatedAsync().ConfigureAwait(false);
 179        EnsureListenerStarted();
 180
 181        // Watermark from the database server's clock, not the app clock: the dispatch loop filters
 182        // pending messages with created_at >= started, and mixing an app-side timestamp with the
 183        // server-stamped created_at would silently drop live deliveries under clock skew.
 184        var startedAtUtc = await _store.GetServerTimeUtcAsync(CancellationToken.None).ConfigureAwait(false);
 185
 186        var storedCorrelationId = correlationId;
 187        var capturedContext = ExecutionContext.Capture();
 188
 189        var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.wait", correlationId: correlationId);
 190        activity?.SetTag("asyncresponse.channel", _activityTag);
 191        AsyncResponseDiagnostics.SetPayloadType(activity, typeof(T));
 192        activity?.SetTag("asyncresponse.timeout_seconds", timeout.Value.TotalSeconds);
 193
 194        var tcs = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);
 195        var registrationId = Guid.NewGuid();
 196        var subscription = new DbSubscription<T>(
 197            this,
 198            correlationId,
 199            registrationId,
 200            startedAtUtc,
 201            completionPredicate,
 202            tcs,
 203            activity);
 204
 205        var timeoutCts = new CancellationTokenSource();
 206        CancellationTokenRegistration timeoutRegistration = default;
 207        subscription.TimeoutRegistration = () => timeoutRegistration.DisposeAsync();
 208        subscription.TimeoutCancellation = timeoutCts;
 209
 210        timeoutRegistration = timeoutCts.Token.Register(
 211            OnWaiterTimeout,
 212            new WaiterTimeoutState<T>(this, subscription, activity, correlationId, tcs));
 213
 214        // Wire the captured-context delegate before the subscription becomes discoverable, so a
 215        // response already stored for this correlation id is processed with the caller's context.
 216        Task ProcessUnderCapturedContextAsync(DbChannelMessage message)
 217        {
 218            async Task Process()
 219            {
 220                using var correlationScope = AsyncResponseContext.PushCorrelationId(storedCorrelationId);
 221                await subscription.ProcessAsync(message).ConfigureAwait(false);
 222            }
 223
 224            if (capturedContext is null)
 225                return Process();
 226
 227            Task? task = null;
 228            ExecutionContext.Run(capturedContext, _ => task = Process(), null);
 229            return task!;
 230        }
 231
 232        subscription.ProcessUnderContextAsync = ProcessUnderCapturedContextAsync;
 233
 234        try
 235        {
 236            var recoveryState = new RecoveryState
 237            {
 238                RegistrationId = registrationId,
 239                ResumeCallback = resumeCallback,
 240                FailureCallback = failureCallback,
 241                CorrelationId = correlationId,
 242                PayloadTypeFullName = typeof(T).FullName,
 243                RegisteredAtUtc = DateTime.UtcNow,
 244                Context = _propagation.Capture()
 245            };
 246            // Subscriber record BEFORE recovery state: "recovery state visible ⇒ subscription
 247            // visible" is the invariant the lost-subscriber dispatcher's live re-check relies on.
 248            // In the reverse order a publisher could see the state, see no subscriber, and consume
 249            // the registration while this waiter is milliseconds from being live.
 250            await _store.UpsertSubscriberAsync(correlationId, registrationId, _instanceId, _options.SubscriberHeartbeatT
 251            await _recoveryStateStore.SaveAsync(correlationId, recoveryState, _options.RecoveryStateExpiry).ConfigureAwa
 252
 253            timeoutCts.CancelAfter(timeout.Value);
 254
 255            if (_logger.IsEnabled(LogLevel.Debug))
 256                _logger.LogDebug("Waiting for {Provider} response on correlationId {CorrelationId} with timeout {Timeout
 257        }
 258        catch (Exception ex)
 259        {
 260            _logger.LogError(ex, "Failed to create {Provider} waiter for correlationId {CorrelationId}.", _providerName,
 261            AsyncResponseDiagnostics.SetError(activity, "subscribe_failure", ex.Message);
 262            await subscription.DrainThenCleanupAsync(deleteRecoveryState: true).ConfigureAwait(false);
 263
 264            // Rethrow instead of returning a pre-faulted waiter: the builder's contract is that
 265            // the trigger runs only once the subscription AND recovery state exist. A returned
 266            // waiter would still let the trigger fire the remote operation with no registration
 267            // left to receive (or recover) its response. Cleanup cancels the response task, so no
 268            // pending task is left behind.
 269            throw;
 270        }
 271
 272        // Publish the subscription only once it is fully armed (heartbeat + timeout + context
 273        // delegate), then signal a scan targeted at this correlation id so any already-stored
 274        // response is delivered promptly without a full sweep.
 275        AddSubscription(correlationId, subscription);
 276        SignalDispatcher(correlationId);
 277
 278        return CreateWaiter<T>(tcs.Task, () => subscription.DrainThenCleanupAsync(deleteRecoveryState: true));
 279    }
 280
 281    /// <inheritdoc />
 282    public Task SetResponse<T>(T response, string correlationId, CancellationToken cancellationToken = default) where T 
 283        => SetResponseCore(response, correlationId, cancellationToken);
 284
 285    Task IRawAsyncResponsePublisher.SetRawResponse(object? response, string correlationId, CancellationToken cancellatio
 286        => SetResponseCore(response, correlationId, cancellationToken);
 287
 288    Task IRawAsyncResponsePublisher.SetRawResponseJson(string responseJson, string correlationId, CancellationToken canc
 289        => SetRawResponseJsonCore(responseJson, correlationId, cancellationToken);
 290
 291    private async Task SetResponseCore<T>(T response, string correlationId, CancellationToken cancellationToken)
 292    {
 293        using var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.set_response", ActivityKind.Producer)
 294        activity?.SetTag("asyncresponse.channel", _activityTag);
 295        AsyncResponseDiagnostics.SetPayloadType(activity, typeof(T));
 296        AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 297
 298        if (string.IsNullOrWhiteSpace(correlationId))
 299        {
 300            _logger.LogWarning("CorrelationId is null; cannot publish the response.");
 301            AsyncResponseDiagnostics.SetError(activity, "correlation_id_null", "CorrelationId is null; cannot publish th
 302            return;
 303        }
 304
 305        try
 306        {
 307            var subscribers = await _store.CountActiveSubscribersAsync(correlationId, cancellationToken).ConfigureAwait(
 308            activity?.SetTag("asyncresponse.subscribers", subscribers);
 309            if (subscribers <= 0)
 310            {
 311                var dispatchResult = await _lostSubscriberDispatcher
 312                    .DispatchLostResponses(
 313                        _recoveryStateStore,
 314                        correlationId,
 315                        response,
 316                        ChannelName(correlationId),
 317                        cancellationToken,
 318                        hasLiveSubscriber: () => HasLiveSubscriberAsync(correlationId, cancellationToken))
 319                    .ConfigureAwait(false);
 320                if (!dispatchResult.RetryLive)
 321                {
 322                    AsyncResponseDiagnostics.SetLostSubscriberRoute(activity, dispatchResult.ShouldResume);
 323                    AsyncResponseDiagnostics.RecordLostSubscriber("response", dispatchResult.ShouldResume, dispatchResul
 324                    activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 325                    return;
 326                }
 327
 328                // A waiter registered between the count and the recovery-state read — publish live
 329                // instead of consuming its registration.
 330            }
 331
 332            var envelope = new AsyncResponseEnvelope<T> { Success = true, Payload = response };
 333            var json = AsyncResponseEnvelopeJson.Serialize(envelope);
 334            var messageId = Guid.NewGuid();
 335            using var confirmation = BeginConfirmation(messageId);
 336            await PublishMessageAsync(messageId, correlationId, json, cancellationToken).ConfigureAwait(false);
 337
 338            if (!await TryConfirmDeliveryAsync(confirmation, cancellationToken).ConfigureAwait(false))
 339            {
 340                var dispatchResult = await _lostSubscriberDispatcher
 341                    .DispatchLostResponses(_recoveryStateStore, correlationId, response, ChannelName(correlationId), can
 342                    .ConfigureAwait(false);
 343                AsyncResponseDiagnostics.SetLostSubscriberRoute(activity, dispatchResult.ShouldResume);
 344                AsyncResponseDiagnostics.RecordLostSubscriber("response", dispatchResult.ShouldResume, dispatchResult.Ca
 345                activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 346            }
 347        }
 348        catch (Exception ex)
 349        {
 350            _logger.LogError(ex, "Failed to publish {Provider} response for correlationId {CorrelationId}.", _providerNa
 351            AsyncResponseDiagnostics.SetError(activity, ex);
 352            throw;
 353        }
 354    }
 355
 356    private async Task SetRawResponseJsonCore(string responseJson, string correlationId, CancellationToken cancellationT
 357    {
 358        using var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.ingress.raw_response", ActivityKind.P
 359        activity?.SetTag("asyncresponse.channel", _activityTag);
 360        AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 361
 362        if (string.IsNullOrWhiteSpace(correlationId))
 363        {
 364            _logger.LogWarning("CorrelationId is null; cannot publish the raw response.");
 365            AsyncResponseDiagnostics.SetError(activity, "correlation_id_null", "CorrelationId is null; cannot publish th
 366            return;
 367        }
 368
 369        try
 370        {
 371            var subscribers = await _store.CountActiveSubscribersAsync(correlationId, cancellationToken).ConfigureAwait(
 372            activity?.SetTag("asyncresponse.subscribers", subscribers);
 373            if (subscribers <= 0)
 374            {
 375                var response = new RawJsonResponse(responseJson).DeserializeUntyped();
 376                var dispatchResult = await _lostSubscriberDispatcher
 377                    .DispatchLostResponses(
 378                        _recoveryStateStore,
 379                        correlationId,
 380                        response,
 381                        ChannelName(correlationId),
 382                        cancellationToken,
 383                        hasLiveSubscriber: () => HasLiveSubscriberAsync(correlationId, cancellationToken))
 384                    .ConfigureAwait(false);
 385                if (!dispatchResult.RetryLive)
 386                {
 387                    AsyncResponseDiagnostics.SetLostSubscriberRoute(activity, dispatchResult.ShouldResume);
 388                    AsyncResponseDiagnostics.RecordLostSubscriber("response", dispatchResult.ShouldResume, dispatchResul
 389                    activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 390                    return;
 391                }
 392
 393                // A waiter registered between the count and the recovery-state read — publish live
 394                // instead of consuming its registration.
 395            }
 396
 397            var messageId = Guid.NewGuid();
 398            using var confirmation = BeginConfirmation(messageId);
 399            await PublishMessageAsync(messageId, correlationId, SerializeRawSuccessEnvelope(responseJson), cancellationT
 400
 401            if (!await TryConfirmDeliveryAsync(confirmation, cancellationToken).ConfigureAwait(false))
 402            {
 403                var response = new RawJsonResponse(responseJson).DeserializeUntyped();
 404                var dispatchResult = await _lostSubscriberDispatcher
 405                    .DispatchLostResponses(_recoveryStateStore, correlationId, response, ChannelName(correlationId), can
 406                    .ConfigureAwait(false);
 407                AsyncResponseDiagnostics.SetLostSubscriberRoute(activity, dispatchResult.ShouldResume);
 408                AsyncResponseDiagnostics.RecordLostSubscriber("response", dispatchResult.ShouldResume, dispatchResult.Ca
 409                activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 410            }
 411        }
 412        catch (Exception ex)
 413        {
 414            _logger.LogError(ex, "Failed to publish {Provider} raw response for correlationId {CorrelationId}.", _provid
 415            AsyncResponseDiagnostics.SetError(activity, ex);
 416            throw;
 417        }
 418    }
 419
 420    /// <inheritdoc />
 421    public async Task SetException(Exception exception, string correlationId, CancellationToken cancellationToken = defa
 422    {
 423        ArgumentNullException.ThrowIfNull(exception);
 424
 425        using var activity = AsyncResponseDiagnostics.StartActivity("asyncresponse.set_exception", ActivityKind.Producer
 426        activity?.SetTag("asyncresponse.channel", _activityTag);
 427        activity?.SetTag("asyncresponse.exception_type", exception.GetType().FullName ?? exception.GetType().Name);
 428        AsyncResponseDiagnostics.SetCorrelationId(activity, correlationId);
 429
 430        if (string.IsNullOrWhiteSpace(correlationId))
 431        {
 432            _logger.LogWarning("CorrelationId is null; cannot publish the exception. Exception: {ExceptionMessage}", exc
 433            AsyncResponseDiagnostics.SetError(activity, "correlation_id_null", "CorrelationId is null; cannot publish th
 434            return;
 435        }
 436
 437        try
 438        {
 439            var subscribers = await _store.CountActiveSubscribersAsync(correlationId, cancellationToken).ConfigureAwait(
 440            activity?.SetTag("asyncresponse.subscribers", subscribers);
 441            if (subscribers <= 0)
 442            {
 443                var dispatchResult = await _lostSubscriberDispatcher
 444                    .DispatchLostExceptions(
 445                        _recoveryStateStore,
 446                        correlationId,
 447                        exception,
 448                        ChannelName(correlationId),
 449                        cancellationToken,
 450                        hasLiveSubscriber: () => HasLiveSubscriberAsync(correlationId, cancellationToken))
 451                    .ConfigureAwait(false);
 452                if (!dispatchResult.RetryLive)
 453                {
 454                    activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 455                    AsyncResponseDiagnostics.RecordLostSubscriber("exception", shouldResume: false, dispatchResult.Callb
 456                    return;
 457                }
 458
 459                // A waiter registered between the count and the recovery-state read — publish live
 460                // instead of consuming its registration.
 461            }
 462
 463            var envelope = new AsyncResponseEnvelope<object>
 464            {
 465                Success = false,
 466                ExceptionMessage = exception.Message,
 467                ExceptionStackTrace = RemoteStackTrace.ForWire(exception.StackTrace, _options.IncludeRemoteStackTrace, _
 468                Payload = null
 469            };
 470            var json = AsyncResponseEnvelopeJson.Serialize(envelope);
 471            var messageId = Guid.NewGuid();
 472            using var confirmation = BeginConfirmation(messageId);
 473            await PublishMessageAsync(messageId, correlationId, json, cancellationToken).ConfigureAwait(false);
 474
 475            if (!await TryConfirmDeliveryAsync(confirmation, cancellationToken).ConfigureAwait(false))
 476            {
 477                // No live re-check here: TryClaimForRecoveryAsync already won the message for the
 478                // recovery path, so live delivery of it is no longer possible.
 479                var dispatchResult = await _lostSubscriberDispatcher
 480                    .DispatchLostExceptions(_recoveryStateStore, correlationId, exception, ChannelName(correlationId), c
 481                    .ConfigureAwait(false);
 482                activity?.SetTag("asyncresponse.recovery.callback_invoked", dispatchResult.CallbackInvoked);
 483                AsyncResponseDiagnostics.RecordLostSubscriber("exception", shouldResume: false, dispatchResult.CallbackI
 484            }
 485        }
 486        catch (Exception ex)
 487        {
 488            _logger.LogError(ex, "Failed to publish {Provider} exception response for correlationId {CorrelationId}.", _
 489            AsyncResponseDiagnostics.SetError(activity, ex);
 490            throw;
 491        }
 492    }
 493
 494    /// <inheritdoc />
 495    public async ValueTask<long> CountActiveSubscribersAsync(string correlationId, CancellationToken cancellationToken =
 496    {
 497        if (string.IsNullOrWhiteSpace(correlationId))
 498            return 0L;
 499
 500        try
 501        {
 502            return await _store.CountActiveSubscribersAsync(correlationId, cancellationToken).ConfigureAwait(false);
 503        }
 504        catch (Exception ex) when (ex is not OperationCanceledException)
 505        {
 506            _logger.LogDebug(ex, "Failed to count {Provider} subscribers for correlationId {CorrelationId}.", _providerN
 507            return 0L;
 508        }
 509    }
 510
 511    /// <summary>
 512    /// Drops local subscriptions while leaving recovery state intact. Used by the sample app to
 513    /// simulate a redeploy for lost-subscriber integration tests.
 514    /// </summary>
 515    internal async Task DropLocalSubscriptionsAsync(CancellationToken cancellationToken = default)
 516    {
 517        foreach (var (correlationId, group) in _subscriptions.ToArray())
 518        {
 519            foreach (var subscription in group.Values.ToArray())
 520            {
 521                await subscription.DropLocalAsync(cancellationToken).ConfigureAwait(false);
 522                group.TryRemove(subscription.Id, out _);
 523            }
 524
 525            if (group.IsEmpty)
 526                _subscriptions.TryRemove(correlationId, out _);
 527
 528            await _executors.RemoveAsync(ChannelName(correlationId)).ConfigureAwait(false);
 529        }
 530    }
 531
 532    /// <summary>
 533    /// Re-probes waiter liveness for the lost-subscriber dispatcher's snapshot-race re-check,
 534    /// using the same active-subscriber count the publish path consulted.
 535    /// </summary>
 536    private async ValueTask<bool> HasLiveSubscriberAsync(string correlationId, CancellationToken cancellationToken)
 537        => await _store.CountActiveSubscribersAsync(correlationId, cancellationToken).ConfigureAwait(false) > 0;
 538
 539    private protected void AddSubscription(string correlationId, IDbSubscription subscription)
 540    {
 541        // Register with the executor registry BEFORE publishing into the subscription map: every
 542        // dispatch path consults the map and then enqueues, so a delivery racing a visible-but-
 543        // unregistered subscription on a correlation id reused within the tombstone lifetime would
 544        // be silently dropped. In the reversed window (registered, not yet visible) the delivery
 545        // just waits for the next sweep or falls back to lost-subscriber recovery.
 546        _executors.OnSubscriptionRegistered(ChannelName(correlationId));
 547        var group = _subscriptions.GetOrAdd(correlationId, _ => new ConcurrentDictionary<Guid, IDbSubscription>());
 548        group[subscription.Id] = subscription;
 549    }
 550
 551    private void RemoveSubscription(string correlationId, Guid registrationId)
 552    {
 553        if (!_subscriptions.TryGetValue(correlationId, out var group))
 554            return;
 555
 556        if (group.TryRemove(registrationId, out _))
 557            _executors.OnSubscriptionRetired(ChannelName(correlationId));
 558        if (group.IsEmpty)
 559            _subscriptions.TryRemove(correlationId, out _);
 560    }
 561
 562    private protected void EnsureListenerStarted()
 563    {
 564        lock (_listenerGate)
 565        {
 566            // Checked under the same gate DisposeAsync sets it under: a racing registration must
 567            // never recreate the CTS and loops after disposal tore them down.
 568            if (_disposed)
 569                throw new ObjectDisposedException(_channelTypeName);
 570
 571            if (_listenerCts is not null)
 572                return;
 573
 574            var listenerCts = new CancellationTokenSource();
 575            _listenerCts = listenerCts;
 576            _listenTask = StartWakeListener(listenerCts.Token);
 577            _dispatchTask = Task.Run(() => DispatchLoopAsync(listenerCts.Token));
 578            _heartbeatTask = Task.Run(() => HeartbeatLoopAsync(listenerCts.Token));
 579        }
 580    }
 581
 582    private async Task HeartbeatLoopAsync(CancellationToken cancellationToken)
 583    {
 584        while (!cancellationToken.IsCancellationRequested)
 585        {
 586            try
 587            {
 588                await Task.Delay(_options.SubscriberHeartbeatInterval, cancellationToken).ConfigureAwait(false);
 589                var registrations = SnapshotActiveRegistrations();
 590                if (registrations.Count > 0)
 591                {
 592                    await _store.HeartbeatSubscribersAsync(
 593                        _instanceId,
 594                        registrations,
 595                        _options.SubscriberHeartbeatTimeout,
 596                        cancellationToken).ConfigureAwait(false);
 597                }
 598            }
 599            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
 600            {
 601                return;
 602            }
 603            catch (Exception ex)
 604            {
 605                _logger.LogWarning(ex, "{Provider} subscriber heartbeat failed; retrying for all local waiters.", _provi
 606            }
 607        }
 608    }
 609
 610    private List<(string CorrelationId, Guid RegistrationId)> SnapshotActiveRegistrations()
 611    {
 612        // Full (correlation id, registration id) pairs: the heartbeat UPSERTs the subscriber
 613        // records, so it needs everything required to re-create one the store's expiry pruning
 614        // (relational pruner / TTL reaper) has already deleted.
 615        var registrations = new List<(string CorrelationId, Guid RegistrationId)>();
 616        foreach (var (correlationId, group) in _subscriptions)
 617        {
 618            foreach (var subscription in group.Values)
 619            {
 620                if (!subscription.Dropped)
 621                    registrations.Add((correlationId, subscription.Id));
 622            }
 623        }
 624
 625        return registrations;
 626    }
 627
 628    private async Task DispatchLoopAsync(CancellationToken cancellationToken)
 629    {
 630        while (!cancellationToken.IsCancellationRequested)
 631        {
 632            try
 633            {
 634                var scope = await CollectDispatchScopeAsync(cancellationToken).ConfigureAwait(false);
 635                await DispatchPendingMessagesAsync(scope, cancellationToken).ConfigureAwait(false);
 636            }
 637            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
 638            {
 639                return;
 640            }
 641            catch (Exception ex)
 642            {
 643                _logger.LogWarning(ex, "{Provider} response dispatch loop failed; retrying after poll delay.", _provider
 644                await Task.Delay(CurrentPollInterval(), cancellationToken).ConfigureAwait(false);
 645            }
 646        }
 647    }
 648
 649    /// <summary>
 650    /// Waits for the next dispatch trigger and returns its scope. <c>null</c> means scan every
 651    /// subscribed correlation id — a full sweep requested explicitly (a null signal) or by the
 652    /// periodic poll that is the missed-wake / cross-process-delivery safety net. A non-null set
 653    /// scans only the signaled correlation ids, so a flood of wake signals never forces a scan of
 654    /// every waiter.
 655    /// </summary>
 656    private protected async Task<HashSet<string>?> CollectDispatchScopeAsync(CancellationToken cancellationToken)
 657    {
 658        // The WhenAny loser is cancelled via the per-iteration linked source: an abandoned
 659        // WaitToReadAsync would otherwise stay parked in the channel's blocked-reader list until
 660        // the next signal — one per poll interval, accumulating without bound on an idle channel.
 661        using var iteration = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
 662        var delay = Task.Delay(CurrentPollInterval(), iteration.Token);
 663        var signal = _signals.Reader.WaitToReadAsync(iteration.Token).AsTask();
 664        var completed = await Task.WhenAny(delay, signal).ConfigureAwait(false);
 665        iteration.Cancel();
 666        if (completed == delay)
 667            return null;
 668
 669        await signal.ConfigureAwait(false);
 670
 671        var scope = new HashSet<string>(StringComparer.Ordinal);
 672        var fullSweep = false;
 673        while (_signals.Reader.TryRead(out var correlationId))
 674        {
 675            if (string.IsNullOrEmpty(correlationId))
 676                fullSweep = true;
 677            else
 678                scope.Add(correlationId);
 679        }
 680
 681        return fullSweep || scope.Count == 0 ? null : scope;
 682    }
 683
 684    private protected async Task DispatchPendingMessagesAsync(HashSet<string>? scope, CancellationToken cancellationToke
 685    {
 686        foreach (var (correlationId, group) in _subscriptions)
 687        {
 688            if (scope is not null && !scope.Contains(correlationId))
 689                continue;
 690
 691            var subscriptions = new List<IDbSubscription>(group.Count);
 692            foreach (var subscription in group.Values)
 693            {
 694                if (!subscription.Dropped)
 695                    subscriptions.Add(subscription);
 696            }
 697            if (subscriptions.Count == 0)
 698                continue;
 699
 700            var since = subscriptions.Min(static s => s.StartedAtUtc).AddSeconds(-1);
 701            var seenCutoff = DateTimeOffset.UtcNow - _options.MessageRetention - TimeSpan.FromMinutes(1);
 702            foreach (var subscription in subscriptions)
 703                subscription.PruneSeen(seenCutoff);
 704
 705            DateTimeOffset? afterCreatedAtUtc = null;
 706            Guid? afterId = null;
 707            while (true)
 708            {
 709                var messages = await _store.LoadMessagesAsync(
 710                    correlationId,
 711                    since,
 712                    _options.PendingMessageBatchSize,
 713                    afterCreatedAtUtc,
 714                    afterId,
 715                    cancellationToken).ConfigureAwait(false);
 716                foreach (var message in messages)
 717                {
 718                    await _executors.EnqueueAsync(
 719                        ChannelName(correlationId),
 720                        () => DispatchMessageToSubscribersAsync(message, subscriptions, cancellationToken),
 721                        cancellationToken).ConfigureAwait(false);
 722                }
 723
 724                if (messages.Count < _options.PendingMessageBatchSize)
 725                    break;
 726
 727                var last = messages[^1];
 728                afterCreatedAtUtc = last.CreatedAtUtc;
 729                afterId = last.Id;
 730            }
 731        }
 732    }
 733
 734    private async Task PublishMessageAsync(
 735        Guid messageId,
 736        string correlationId,
 737        string envelopeJson,
 738        CancellationToken cancellationToken)
 739    {
 740        // The insert itself carries the remote wake where the provider has one (a NOTIFY rides the
 741        // PostgreSQL insert; MongoDB change streams observe it) and the SQL Server sweep polls it
 742        // up. Only the local fast path and a targeted local signal are needed on top. The store
 743        // returns the SERVER-stamped created_at for the local fast-path message: subscription
 744        // watermarks are server-clock, and an app-clock timestamp here silently disabled the fast
 745        // path whenever the app clock ran more than the 1s tolerance behind the database — delivery
 746        // then quietly degraded to sweep latency on every publish.
 747        var createdAtUtc = await _store.InsertMessageAsync(messageId, correlationId, envelopeJson, _options.MessageReten
 748            .ConfigureAwait(false);
 749        await TryDispatchLocalSubscribersAsync(
 750            new DbChannelMessage(messageId, correlationId, envelopeJson, createdAtUtc),
 751            cancellationToken).ConfigureAwait(false);
 752        SignalDispatcher(correlationId);
 753    }
 754
 755    private protected async Task DispatchMessageToSubscribersAsync(
 756        DbChannelMessage message,
 757        IReadOnlyList<IDbSubscription> subscriptions,
 758        CancellationToken cancellationToken)
 759    {
 760        // Only subscriptions that are still live, inside their delivery watermark, and have not
 761        // already processed this message. Skipping when there is nothing to deliver also avoids a
 762        // redundant claim on every re-sweep.
 763        var hasTargets = false;
 764        foreach (var subscription in subscriptions)
 765        {
 766            if (!subscription.Dropped && IsWithinWatermark(subscription, message) && !subscription.HasSeen(message.Id))
 767            {
 768                hasTargets = true;
 769                break;
 770            }
 771        }
 772        if (!hasTargets)
 773            return;
 774
 775        // Take the message for live delivery. The claim sets acked_at unless the publisher already
 776        // routed it to recovery (recovery_claimed); losing the claim means recovery owns it, so it is
 777        // not delivered to the waiter and handled a second time.
 778        //
 779        // Claim-then-dispatch is deliberate — keep this ordering. The in-process handoff is
 780        // at-most-once by design: a crash between the claim and the waiter's continuation can only
 781        // lose delivery to waiters in THIS dying process, which no ordering could save (their
 782        // continuations die with it), while pre-registered fan-out waiters in other processes
 783        // still receive the acked message (IsWithinWatermark admits acked_at > started_at).
 784        // Dispatch-then-ack behind an expiring claim would re-open the stale-redelivery wrong-data
 785        // bug the strict acked exclusion in IsWithinWatermark closes. Durability across process
 786        // death belongs to the layer above: flow re-execution, publish-time recovery routing, and
 787        // the step timeout.
 788        if (!await _store.TryClaimForDeliveryAsync(message.Id, cancellationToken).ConfigureAwait(false))
 789        {
 790            foreach (var subscription in subscriptions)
 791            {
 792                if (!subscription.Dropped)
 793                    subscription.MarkSeen(message.Id);
 794            }
 795            return;
 796        }
 797
 798        // Wake the publisher immediately if it is waiting in this process — no acked_at polling needed.
 799        if (_pendingConfirmations.TryGetValue(message.Id, out var confirmation))
 800            confirmation.TrySetResult(true);
 801
 802        foreach (var subscription in subscriptions)
 803        {
 804            if (subscription.Dropped || !IsWithinWatermark(subscription, message) || !subscription.MarkSeen(message.Id))
 805                continue;
 806
 807            await subscription.ProcessUnderContextAsync(message).ConfigureAwait(false);
 808        }
 809    }
 810
 811    /// <summary>
 812    /// Per-subscription delivery watermark. The sweep queries with the OLDEST waiter's watermark on
 813    /// a shared correlation id, so without this filter a late-joining waiter would receive retained
 814    /// messages created before it registered. Same 1s tolerance as the query watermark.
 815    /// <para>
 816    /// The creation-time tolerance alone re-admits history: a message created inside the 1s skew
 817    /// window may have already been delivered and acked for a PREVIOUS waiter that reused the
 818    /// correlation id, and per-subscription seen-tracking cannot dedupe what a different
 819    /// subscription processed. A message acked before this subscription existed is history, not
 820    /// delivery — waiters that legitimately participate in a delivery (including cross-process
 821    /// fan-out) were registered before its claim stamped <c>acked_at</c>. The acked comparison is
 822    /// deliberately strict, with no skew tolerance: under skew, strictness can only make a waiter
 823    /// whose registration raced another process's in-flight ack keep waiting for its own response,
 824    /// whereas a tolerance would re-open the stale-redelivery window this check closes.
 825    /// </para>
 826    /// <para>
 827    /// The comparison must be STRICTLY greater, and that is load-bearing rather than stylistic: a
 828    /// server clock's resolution is far coarser than its column precision, so equal timestamps are
 829    /// routine, not a measure-zero tie. SQL Server stamps <c>datetime2(7)</c> from
 830    /// <c>SYSUTCDATETIME()</c> — 100ns precision, but the clock behind it advances in ~5ms ticks
 831    /// (measured: 30,344 samples over 300ms yielded 61 distinct values, mean gap 4.9ms), and
 832    /// MongoDB's <c>$$NOW</c> is millisecond-resolution. A waiter that reuses a correlation id
 833    /// within one tick of the previous waiter's ack therefore registers at exactly
 834    /// <c>acked_at</c>, and a non-strict comparison hands it the response its predecessor already
 835    /// consumed. Registration is ordered strictly after that ack in real time and the clock is
 836    /// non-decreasing, so <c>acked_at &lt;= started_at</c> always holds for history and the strict
 837    /// form excludes it deterministically — not probabilistically.
 838    /// </para>
 839    /// <para>
 840    /// The tie is symmetric, and its resolution is deliberate: the same equality can also be a
 841    /// genuine cross-process fan-out delivery (this waiter registered and another process's claim
 842    /// stamped <c>acked_at</c> inside one clock tick), and the strict form then excludes the
 843    /// waiter from its own response — it recovers through its step timeout and the
 844    /// idempotent-restart contract. That at-most-once cost (a rare missed delivery that
 845    /// self-heals) is chosen over the wrong-data redelivery a tolerant comparison re-opens. No
 846    /// timestamp can separate the two same-tick cases; only an identity carried on the claim (the
 847    /// claiming registration id, or a monotonic sequence) could — a possible store-schema
 848    /// evolution if the trade ever bites in practice. The recovery asymmetry is worth naming for
 849    /// the eventual triager: excluded HISTORY re-subscribes and proceeds immediately, while an
 850    /// excluded fan-out waiter stalls for its full timeout first — a durable-flow step's default
 851    /// is 7 days, and a plain waiter surfaces a TimeoutException to its caller. "Bites in
 852    /// practice" looks like a long stall, not a quick retry.
 853    /// </para>
 854    /// </summary>
 855    private static bool IsWithinWatermark(IDbSubscription subscription, DbChannelMessage message)
 856        => message.CreatedAtUtc >= subscription.StartedAtUtc.AddSeconds(-1)
 857           && (message.AckedAtUtc is null || message.AckedAtUtc > subscription.StartedAtUtc);
 858
 859    private async Task TryDispatchLocalSubscribersAsync(DbChannelMessage message, CancellationToken cancellationToken)
 860    {
 861        if (!_subscriptions.TryGetValue(message.CorrelationId, out var group))
 862            return;
 863
 864        var subscriptions = new List<IDbSubscription>(group.Count);
 865        foreach (var subscription in group.Values)
 866        {
 867            if (!subscription.Dropped)
 868                subscriptions.Add(subscription);
 869        }
 870        if (subscriptions.Count == 0)
 871            return;
 872
 873        // Same-process fast path: skips the wake round trip / sweep latency but still runs on the
 874        // per-correlation serial executor — completion predicates are guaranteed serial, in-order
 875        // invocation on every channel, and a direct dispatch here could otherwise run concurrently
 876        // with a sweep-enqueued dispatch of a different message for the same subscription. MarkSeen
 877        // keeps the sweep from double-processing this message.
 878        await _executors.EnqueueAsync(
 879            ChannelName(message.CorrelationId),
 880            new LocalDispatchWorkItem(this, message, subscriptions, cancellationToken).InvokeAsync,
 881            cancellationToken).ConfigureAwait(false);
 882    }
 883
 884    /// <summary>
 885    /// Registers an in-process delivery completion for a message id. Disposing it removes the entry,
 886    /// so a publish that throws or completes never leaks the registration.
 887    /// </summary>
 888    private protected PendingConfirmation BeginConfirmation(Guid messageId)
 889    {
 890        var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
 891        _pendingConfirmations[messageId] = tcs;
 892        return new PendingConfirmation(this, messageId, tcs);
 893    }
 894
 895    /// <summary>
 896    /// Confirms a published response reached a live waiter. Returns <c>true</c> once a waiter has
 897    /// acknowledged it; on confirmation timeout, atomically claims the message for the lost-subscriber
 898    /// path and returns <c>false</c> only if that claim wins — so the recovery callback and a
 899    /// slow-but-live waiter are mutually exclusive.
 900    /// </summary>
 901    private protected async Task<bool> TryConfirmDeliveryAsync(PendingConfirmation confirmation, CancellationToken cance
 902    {
 903        if (await WaitForAcknowledgementAsync(confirmation, cancellationToken).ConfigureAwait(false))
 904            return true;
 905
 906        return !await _store.TryClaimForRecoveryAsync(confirmation.MessageId, cancellationToken).ConfigureAwait(false);
 907    }
 908
 909    private protected void SignalDispatcher(string? correlationId = null) => _signals.Writer.TryWrite(correlationId);
 910
 911    private async Task<bool> WaitForAcknowledgementAsync(PendingConfirmation confirmation, CancellationToken cancellatio
 912    {
 913        var deadline = DateTimeOffset.UtcNow + _options.DeliveryConfirmationTimeout;
 914        while (DateTimeOffset.UtcNow < deadline)
 915        {
 916            var remaining = deadline - DateTimeOffset.UtcNow;
 917            if (remaining <= TimeSpan.Zero)
 918                break;
 919
 920            var pollDelay = remaining < _options.DeliveryConfirmationPollInterval
 921                ? remaining
 922                : _options.DeliveryConfirmationPollInterval;
 923
 924            // Fast path: an in-process delivery trips the completion and we return without a query.
 925            await Task.WhenAny(confirmation.Delivered, Task.Delay(pollDelay, cancellationToken)).ConfigureAwait(false);
 926            if (confirmation.Delivered.IsCompletedSuccessfully)
 927                return true;
 928
 929            // Slow path: a delivery in another process only set acked_at, so poll for it.
 930            if (await _store.IsMessageAcknowledgedAsync(confirmation.MessageId, cancellationToken).ConfigureAwait(false)
 931                return true;
 932        }
 933
 934        return confirmation.Delivered.IsCompletedSuccessfully
 935            || await _store.IsMessageAcknowledgedAsync(confirmation.MessageId, cancellationToken).ConfigureAwait(false);
 936    }
 937
 938    private static string SerializeRawSuccessEnvelope(string payloadJson)
 939    {
 940        JsonSafety.ThrowIfClearlyNotJson(payloadJson);
 941
 942        var buffer = new ArrayBufferWriter<byte>();
 943        using (var writer = new Utf8JsonWriter(buffer))
 944        {
 945            writer.WriteStartObject();
 946            writer.WriteNumber("SchemaVersion", AsyncResponseEnvelopeSchema.Current);
 947            writer.WriteBoolean("Success", true);
 948            writer.WritePropertyName("Payload");
 949            writer.WriteRawValue(payloadJson);
 950            writer.WriteNull("ExceptionMessage");
 951            writer.WriteNull("ExceptionStackTrace");
 952            writer.WriteEndObject();
 953        }
 954
 955        return Encoding.UTF8.GetString(buffer.WrittenSpan);
 956    }
 957
 958    [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage]
 959    private static void OnWaiterTimeout(object? state)
 960        => ((IWaiterTimeoutState)state!).Schedule();
 961
 962    private async Task HandleWaiterTimeoutAsync<T>(
 963        DbSubscription<T> subscription,
 964        Activity? activity,
 965        string correlationId,
 966        TaskCompletionSource<T> tcs) where T : IAsyncResponsePayload
 967    {
 968        _logger.LogWarning("Timed out waiting for {Provider} response for correlationId {CorrelationId}.", _providerName
 969        AsyncResponseDiagnostics.SetError(activity, "timeout", $"Timed out waiting for response for correlationId {corre
 970        AsyncResponseDiagnostics.RecordWaiterTimeout(_activityTag);
 971        tcs.TrySetException(new TimeoutException($"Timed out waiting for response for correlationId {correlationId}."));
 972        await subscription.DrainThenCleanupAsync(deleteRecoveryState: true).ConfigureAwait(false);
 973    }
 974
 975    private interface IWaiterTimeoutState
 976    {
 977        void Schedule();
 978    }
 979
 980    [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage]
 981    private sealed class WaiterTimeoutState<T>(
 982        DbAsyncResponseChannelBase owner,
 983        DbSubscription<T> subscription,
 984        Activity? activity,
 985        string correlationId,
 986        TaskCompletionSource<T> tcs) : IWaiterTimeoutState where T : IAsyncResponsePayload
 987    {
 988        public void Schedule()
 989            => _ = Task.Run(async () =>
 990            {
 991                try
 992                {
 993                    await owner.HandleWaiterTimeoutAsync(subscription, activity, correlationId, tcs).ConfigureAwait(fals
 994                }
 995                catch (Exception ex)
 996                {
 997                    // Fire-and-forget: nothing awaits this task, so an escaped fault would vanish.
 998                    owner._logger.LogError(ex, "Error handling {Provider} waiter timeout for correlationId {CorrelationI
 999                }
 1000            });
 1001    }
 1002
 1003    [System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage]
 1004    private sealed class LocalDispatchWorkItem(
 1005        DbAsyncResponseChannelBase owner,
 1006        DbChannelMessage message,
 1007        IReadOnlyList<IDbSubscription> subscriptions,
 1008        CancellationToken cancellationToken)
 1009    {
 1010        public async Task InvokeAsync()
 1011        {
 1012            try
 1013            {
 1014                await owner.DispatchMessageToSubscribersAsync(message, subscriptions, cancellationToken).ConfigureAwait(
 1015            }
 1016            catch (Exception ex) when (ex is not OperationCanceledException || !cancellationToken.IsCancellationRequeste
 1017            {
 1018                owner._logger.LogDebug(
 1019                    ex,
 1020                    "Local {Provider} response dispatch failed for correlationId {CorrelationId}; {RetryHint}.",
 1021                    owner._providerName,
 1022                    message.CorrelationId,
 1023                    owner._localDispatchRetryHint);
 1024            }
 1025        }
 1026    }
 1027
 1028    /// <inheritdoc />
 1029    public async ValueTask DisposeAsync()
 1030    {
 1031        CancellationTokenSource? cts;
 1032        Task? listenTask;
 1033        Task? dispatchTask;
 1034        Task? heartbeatTask;
 1035        lock (_listenerGate)
 1036        {
 1037            // Set under the gate so EnsureListenerStarted can never observe "not disposed" and
 1038            // then recreate the CTS/loops this teardown is about to stop.
 1039            _disposed = true;
 1040            cts = _listenerCts;
 1041            listenTask = _listenTask;
 1042            dispatchTask = _dispatchTask;
 1043            heartbeatTask = _heartbeatTask;
 1044            _listenerCts = null;
 1045            _listenTask = null;
 1046            _dispatchTask = null;
 1047            _heartbeatTask = null;
 1048        }
 1049
 1050        if (cts is not null)
 1051        {
 1052            await cts.CancelAsync().ConfigureAwait(false);
 1053            try
 1054            {
 1055                await Task.WhenAll(new[] { listenTask, dispatchTask, heartbeatTask }.OfType<Task>()).ConfigureAwait(fals
 1056            }
 1057            catch (OperationCanceledException)
 1058            {
 1059            }
 1060            cts.Dispose();
 1061        }
 1062
 1063        foreach (var (correlationId, group) in _subscriptions.ToArray())
 1064        {
 1065            foreach (var subscription in group.Values.ToArray())
 1066                await subscription.DrainThenCleanupAsync(deleteRecoveryState: false).ConfigureAwait(false);
 1067            await _executors.RemoveAsync(ChannelName(correlationId)).ConfigureAwait(false);
 1068        }
 1069    }
 1070
 1071    /// <summary>Scopes an in-process delivery completion; <see cref="Dispose"/> unregisters it.</summary>
 1072    private protected readonly struct PendingConfirmation(
 1073        DbAsyncResponseChannelBase owner,
 1074        Guid messageId,
 1075        TaskCompletionSource<bool> tcs) : IDisposable
 1076    {
 1077        public Guid MessageId => messageId;
 1078        public Task<bool> Delivered => tcs.Task;
 1079        public void Dispose() => owner._pendingConfirmations.TryRemove(messageId, out _);
 1080    }
 1081
 1082    private protected interface IDbSubscription
 1083    {
 1084        Guid Id { get; }
 1085        DateTimeOffset StartedAtUtc { get; }
 1086        bool Dropped { get; }
 1087        Func<DbChannelMessage, Task> ProcessUnderContextAsync { get; set; }
 1088        bool HasSeen(Guid messageId);
 1089        bool MarkSeen(Guid messageId);
 1090        void PruneSeen(DateTimeOffset cutoffUtc);
 1091        Task ProcessAsync(DbChannelMessage message);
 1092        ValueTask CleanupOnceAsync(bool deleteRecoveryState);
 1093        ValueTask DrainThenCleanupAsync(bool deleteRecoveryState);
 1094        ValueTask DropLocalAsync(CancellationToken cancellationToken);
 1095    }
 1096
 1097    private sealed class DbSubscription<T> : IDbSubscription where T : IAsyncResponsePayload
 1098    {
 1099        private readonly DbAsyncResponseChannelBase _owner;
 1100        private readonly string _correlationId;
 1101        private readonly Func<T, ValueTask<bool>> _completionPredicate;
 1102        private readonly TaskCompletionSource<T> _tcs;
 1103        private readonly Activity? _activity;
 1104        private readonly HashSet<Guid> _seen = [];
 1105        private readonly Queue<(Guid Id, DateTimeOffset SeenAtUtc)> _seenOrder = [];
 1106        private readonly object _seenGate = new();
 1107        private int _cleanupStarted;
 1108        private volatile bool _dropped;
 1109        private readonly object _cleanupGate = new();
 1110        private Task? _cleanupTask;
 1111
 1112        public DbSubscription(
 1113            DbAsyncResponseChannelBase owner,
 1114            string correlationId,
 1115            Guid registrationId,
 1116            DateTimeOffset startedAtUtc,
 1117            Func<T, ValueTask<bool>> completionPredicate,
 1118            TaskCompletionSource<T> tcs,
 1119            Activity? activity)
 1120        {
 1121            _owner = owner;
 1122            _correlationId = correlationId;
 1123            Id = registrationId;
 1124            StartedAtUtc = startedAtUtc;
 1125            _completionPredicate = completionPredicate;
 1126            _tcs = tcs;
 1127            _activity = activity;
 1128            ProcessUnderContextAsync = ProcessAsync;
 1129        }
 1130
 1131        public Guid Id { get; }
 1132        public DateTimeOffset StartedAtUtc { get; }
 1133        public bool Dropped => _dropped;
 1134        public Func<ValueTask>? TimeoutRegistration { get; set; }
 1135        public CancellationTokenSource? TimeoutCancellation { get; set; }
 1136        public Func<DbChannelMessage, Task> ProcessUnderContextAsync { get; set; }
 1137
 1138        public bool HasSeen(Guid messageId)
 1139        {
 1140            lock (_seenGate)
 1141            {
 1142                return _seen.Contains(messageId);
 1143            }
 1144        }
 1145
 1146        public bool MarkSeen(Guid messageId)
 1147        {
 1148            lock (_seenGate)
 1149            {
 1150                if (!_seen.Add(messageId))
 1151                    return false;
 1152
 1153                // Use the local observation time, not the database creation time. This keeps the
 1154                // pruning queue monotonic and avoids immediate eviction when app and DB clocks differ.
 1155                _seenOrder.Enqueue((messageId, DateTimeOffset.UtcNow));
 1156                return true;
 1157            }
 1158        }
 1159
 1160        public void PruneSeen(DateTimeOffset cutoffUtc)
 1161        {
 1162            lock (_seenGate)
 1163            {
 1164                while (_seenOrder.TryPeek(out var entry) && entry.SeenAtUtc < cutoffUtc)
 1165                {
 1166                    _seenOrder.Dequeue();
 1167                    _seen.Remove(entry.Id);
 1168                }
 1169            }
 1170        }
 1171
 1172        public async Task ProcessAsync(DbChannelMessage message)
 1173        {
 31174            if (_dropped)
 31175                return;
 1176
 31177            var finished = false;
 1178            try
 1179            {
 31180                var envelope = JsonSerializer.Deserialize(message.EnvelopeJson, AsyncResponseEnvelopeJson.TypeInfo<T>())
 31181                if (envelope is null)
 1182                {
 31183                    finished = true;
 31184                    var error = new JsonException($"Failed to deserialize envelope for correlationId {_correlationId}.")
 31185                    AsyncResponseDiagnostics.SetError(_activity, "deserialize_failure", error.Message);
 31186                    _tcs.TrySetException(error);
 1187                }
 31188                else if (!AsyncResponseEnvelopeSchema.IsReadable(envelope.SchemaVersion))
 1189                {
 31190                    finished = true;
 31191                    var error = new InvalidOperationException(
 31192                        $"Response envelope for correlationId {_correlationId} has schema version {envelope.SchemaVersio
 31193                        $"which this build does not support (current: {AsyncResponseEnvelopeSchema.Current}).");
 31194                    AsyncResponseDiagnostics.SetError(_activity, "schema_mismatch", error.Message);
 31195                    _tcs.TrySetException(error);
 1196                }
 31197                else if (!envelope.Success)
 1198                {
 31199                    finished = true;
 31200                    var remoteFailure = new Exception(envelope.ExceptionMessage ?? "Unknown error during asynchronous pr
 31201                    if (!string.IsNullOrEmpty(envelope.ExceptionStackTrace))
 31202                        remoteFailure.Data["RemoteStackTrace"] = RemoteStackTrace.Cap(envelope.ExceptionStackTrace, _own
 31203                    AsyncResponseDiagnostics.SetError(_activity, "remote_failure", remoteFailure.Message);
 31204                    _tcs.TrySetException(remoteFailure);
 1205                }
 1206                else
 1207                {
 31208                    finished = await _completionPredicate(envelope.Payload!).ConfigureAwait(false);
 31209                    if (finished)
 31210                        _tcs.TrySetResult(envelope.Payload!);
 1211                }
 31212            }
 31213            catch (Exception ex)
 1214            {
 31215                finished = true;
 31216                _owner._logger.LogError(ex, "Error processing {Provider} response for correlationId {CorrelationId}.", _
 31217                AsyncResponseDiagnostics.SetError(_activity, ex);
 31218                _tcs.TrySetException(ex);
 31219            }
 1220            finally
 1221            {
 31222                if (finished)
 31223                    await CleanupOnceAsync(deleteRecoveryState: true).ConfigureAwait(false);
 1224            }
 31225        }
 1226
 1227        /// <summary>
 1228        /// Task-latched so EVERY caller completes only when the one real cleanup has finished —
 1229        /// a fire-once flag alone would let a disposing waiter racing the timeout return before
 1230        /// the response task was settled.
 1231        /// </summary>
 1232        public ValueTask CleanupOnceAsync(bool deleteRecoveryState)
 1233        {
 1234            Task task;
 1235            lock (_cleanupGate)
 1236            {
 1237                task = _cleanupTask ??= CleanupCoreAsync(deleteRecoveryState);
 1238            }
 1239
 1240            return task.IsCompletedSuccessfully ? ValueTask.CompletedTask : new ValueTask(task);
 1241        }
 1242
 1243        /// <summary>
 1244        /// Dispose-path cleanup: DRAINS the per-correlation serial executor before settling. A
 1245        /// delivery may be mid <c>Until</c>-predicate holding a message the claim already acked;
 1246        /// the marker work item completes only after that in-flight item finished, so by the time
 1247        /// cleanup cancels, the task is either settled by the delivery or genuinely undelivered —
 1248        /// never a cancellation stealing a consumed response. Must NOT be called from dispatch
 1249        /// code (which runs ON the executor): the dispatch-triggered cleanup calls
 1250        /// <see cref="CleanupOnceAsync"/> directly, its task already settled.
 1251        /// <para>
 1252        /// The drain is bounded by <c>DisposalDrainTimeout</c> — a single budget covering marker
 1253        /// ADMISSION too, since a full bounded queue behind a wedged item blocks the enqueue
 1254        /// itself. A lapsed budget must not fall back to the cleanup's cancel: the wedged delivery
 1255        /// holds a message the claim already consumed, and "canceled" would tell a re-attaching
 1256        /// caller nothing was delivered. It faults the task with the explicit indeterminate
 1257        /// contract instead, routing durable flows to a fresh idempotent restart. An enqueue
 1258        /// suppressed by the registry's tombstone is the opposite case — the retired executor
 1259        /// finished everything it ever admitted, so nothing is in flight and the plain cancel
 1260        /// below is truthful.
 1261        /// </para>
 1262        /// </summary>
 1263        public async ValueTask DrainThenCleanupAsync(bool deleteRecoveryState)
 1264        {
 31265            if (Volatile.Read(ref _cleanupStarted) == 0)
 1266            {
 11267                var drainTimeout = _owner._options.DisposalDrainTimeout;
 11268                var drained = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
 1269                try
 1270                {
 11271                    using var budget = new CancellationTokenSource(drainTimeout);
 11272                    var accepted = await _owner._executors.EnqueueAsync(_owner.ChannelName(_correlationId), () =>
 11273                    {
 11274                        drained.TrySetResult();
 11275                        return Task.CompletedTask;
 11276                    }, budget.Token).ConfigureAwait(false);
 11277                    if (accepted)
 11278                        await drained.Task.WaitAsync(budget.Token).ConfigureAwait(false);
 11279                }
 11280                catch (Exception drainEx)
 1281                {
 1282                    // Budget lapse — or an unforeseen drain failure: either way the marker never
 1283                    // ran, so an in-flight delivery cannot be ruled out (only accepted=false
 1284                    // proves the executor finished everything). Settlement unproven means the
 1285                    // cleanup's cancel below would be a false "nothing was delivered" — fault
 1286                    // with the explicit indeterminate contract instead. A TrySetResult from the
 1287                    // late-finishing dispatch loses against this and is dropped; its cleanup
 1288                    // call is a no-op behind the latch.
 11289                    _owner._logger.LogWarning(
 11290                        "Disposal drain for {Provider} correlationId {CorrelationId} did not prove settlement within {Dr
 11291                        _owner._providerName, _correlationId, drainTimeout);
 11292                    AsyncResponseDiagnostics.SetError(_activity, "indeterminate_delivery", "Disposal drain did not prove
 11293                    if (drainEx is not OperationCanceledException)
 11294                        _owner._logger.LogDebug(drainEx, "Dispatch drain failed for correlationId {CorrelationId}.", _co
 11295                    _tcs.TrySetException(new AsyncResponseIndeterminateDeliveryException(_correlationId, drainTimeout));
 11296                }
 11297            }
 1298
 31299            await CleanupOnceAsync(deleteRecoveryState).ConfigureAwait(false);
 31300        }
 1301
 1302        private async Task CleanupCoreAsync(bool deleteRecoveryState)
 1303        {
 1304            // The flag is kept alongside the task latch: dispatch cores and white-box tests gate
 1305            // on it, and a pre-set flag (test isolation) must keep skipping the network cleanup.
 31306            if (Interlocked.Exchange(ref _cleanupStarted, 1) != 0)
 31307                return;
 1308
 1309            // A waiter disposed before any terminal signal must not leave ResponseTask pending
 1310            // forever for callers that hold it directly — the timeout dies with this cleanup, so
 1311            // nothing else could ever complete the task. This also covers channel DisposeAsync at
 1312            // host shutdown, which runs this cleanup over every in-flight subscription and would
 1313            // otherwise hang still-awaiting WaitAsync callers. Cancellation is a no-op after a
 1314            // normal completion, timeout, fault, or a delivery drained by DrainThenCleanupAsync.
 31315            _tcs.TrySetCanceled();
 1316
 1317            try
 1318            {
 31319                _dropped = true;
 1320
 1321                try
 1322                {
 1323                    // Delete the recovery state BEFORE removing the subscription (locally and in the
 1324                    // subscriber store). In the reverse order a publish landing in the window sees
 1325                    // "no subscriber, state present" and fires a spurious recovery callback for a wait
 1326                    // that already reached a terminal state. In this order the window shows a
 1327                    // subscriber that drops the message — a late or duplicate terminal message is
 1328                    // droppable; a resurrected recovery callback is not. (Shutdown/redeploy paths pass
 1329                    // deleteRecoveryState: false and keep the state for lost-subscriber recovery.)
 31330                    if (deleteRecoveryState)
 31331                        await _owner._recoveryStateStore.TryDeleteAsync(_correlationId, Id).ConfigureAwait(false);
 11332                }
 31333                catch (Exception ex)
 1334                {
 1335                    // Best-effort: the state expires on its own, and a transient store failure must
 1336                    // not skip the local teardown below.
 31337                    _owner._logger.LogError(ex, "Failed to delete {Provider} recovery state for correlationId {Correlati
 31338                }
 1339
 1340                try
 1341                {
 31342                    await _owner._store.DeleteSubscriberAsync(_correlationId, Id, CancellationToken.None).ConfigureAwait
 11343                }
 31344                catch (Exception ex)
 1345                {
 1346                    // Best-effort: an orphaned subscriber record ages out via the heartbeat timeout.
 31347                    _owner._logger.LogError(ex, "Failed to delete {Provider} subscriber {SubscriberRecord} for correlati
 31348                }
 1349            }
 1350            finally
 1351            {
 1352                // Purely local teardown runs no matter which network call above failed — the
 1353                // cleanup latch is already set, so a skipped removal would leak the subscription
 1354                // map entry and the executor until process exit.
 31355                _owner.RemoveSubscription(_correlationId, Id);
 1356
 1357                // Schedule the executor retirement on the thread pool; do not await directly —
 1358                // dispatch-loop deliveries run this cleanup ON the executor, and RemoveAsync waits
 1359                // for the executor's drain loop to finish, which would be a circular await.
 31360                var channelName = _owner.ChannelName(_correlationId);
 31361                _ = Task.Run(async () =>
 31362                {
 31363                    try
 31364                    {
 31365                        await _owner._executors.RemoveAsync(channelName).ConfigureAwait(false);
 31366                    }
 11367                    catch (Exception ex)
 31368                    {
 11369                        _owner._logger.LogError(ex, "Failed to retire the executor for channel {Channel}.", channelName)
 11370                    }
 31371                });
 1372
 31373                if (TimeoutRegistration is not null)
 11374                    await TimeoutRegistration().ConfigureAwait(false);
 31375                TimeoutCancellation?.Dispose();
 31376                _activity?.Dispose();
 1377            }
 31378        }
 1379
 1380        public async ValueTask DropLocalAsync(CancellationToken cancellationToken)
 1381        {
 11382            _dropped = true;
 11383            await _owner._store.DeleteSubscriberAsync(_correlationId, Id, cancellationToken).ConfigureAwait(false);
 11384        }
 1385    }
 1386}
Generated by: ReportGenerator 5.5.10.0
08/03/2026 - 08:59:44
GitHub | reportgenerator.io

Methods/Properties

ProcessAsync()
DrainThenCleanupAsync()
CleanupCoreAsync()
<CleanupCoreAsync()
DropLocalAsync()