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

Information
Class: Microsoft.Extensions.DependencyInjection.CosmosDurableFlowServiceCollectionExtensions
Assembly: AsyncResponse.DurableFlows.Cosmos
File(s): /_/src/DurableFlows/AsyncResponse.DurableFlows.Cosmos/CosmosDurableFlows.cs
Line coverage
100%
Covered lines: 13
Uncovered lines: 0
Coverable lines: 13
Total lines: 939
Line coverage: 100%
Branch coverage
100%
Covered branches: 4
Total branches: 4
Branch coverage: 100%
Method coverage

Feature is only available for sponsors

Upgrade to PRO version

Metrics

MethodBranch coverage Crap Score Cyclomatic complexity Line coverage
WithCosmosDurableFlows(...)100%44100%

File(s)

/_/src/DurableFlows/AsyncResponse.DurableFlows.Cosmos/CosmosDurableFlows.cs

#LineLine coverage
 1using AsyncResponse;
 2using AsyncResponse.DurableFlows.Internal;
 3using AsyncResponse.DurableFlows.Cosmos;
 4using Microsoft.Azure.Cosmos;
 5using Microsoft.Extensions.DependencyInjection.Extensions;
 6using Microsoft.Extensions.Options;
 7using Newtonsoft.Json;
 8using System.Buffers;
 9using System.Net;
 10using System.Text;
 11
 12namespace Microsoft.Extensions.DependencyInjection
 13{
 14    /// <summary>DI registration for the Azure Cosmos DB durable-flow state store.</summary>
 15    public static class CosmosDurableFlowServiceCollectionExtensions
 16    {
 17        /// <summary>
 18        /// Stores durable-flow state in Azure Cosmos DB. Hosts may either register a
 19        /// <see cref="CosmosClient"/> singleton or set connection options here.
 20        /// </summary>
 21        public static AsyncResponseRegistrationBuilder WithCosmosDurableFlows(
 22            this AsyncResponseRegistrationBuilder builder,
 23            Action<CosmosDurableFlowOptions>? configure = null)
 24        {
 25            // Singleton on purpose: database/container provisioning is cached per store instance
 26            // and Cosmos metadata operations are RU-charged and rate-limited — a scoped store would
 27            // re-issue them on every flow execution. A host-registered CosmosClient is reused when
 28            // present; otherwise the store creates and owns one from ConnectionString. Nothing is
 29            // registered as a bare CosmosClient service, so unrelated resolutions of that type are
 30            // never answered — or broken — by this package.
 20431            builder.Services.TryAddSingleton(provider =>
 20432            {
 20433                var options = provider.GetRequiredService<IOptions<CosmosDurableFlowOptions>>();
 20434
 20435                var shared = provider.GetService<CosmosClient>();
 20436                if (shared is not null)
 20037                    return new CosmosFlowStateStore(shared, options);
 20438
 439                if (string.IsNullOrWhiteSpace(options.Value.ConnectionString))
 240                    throw new InvalidOperationException($"{nameof(CosmosDurableFlowOptions)}.{nameof(CosmosDurableFlowOp
 241                return new CosmosFlowStateStore(new CosmosClient(options.Value.ConnectionString), options, ownsClient: t
 20442            });
 20443            return builder.WithDurableFlows<CosmosFlowStateStore, CosmosDurableFlowOptions>(configure);
 44        }
 45    }
 46}
 47
 48namespace AsyncResponse.DurableFlows.Cosmos
 49{
 50/// <summary>Options for the Azure Cosmos DB durable-flow state store.</summary>
 51public sealed class CosmosDurableFlowOptions : DurableFlowOptions
 52{
 53    /// <summary>Optional Cosmos DB connection string used when no <see cref="CosmosClient"/> is registered.</summary>
 54    public string? ConnectionString { get; set; }
 55
 56    /// <summary>Cosmos database name. Required.</summary>
 57    public string? DatabaseName { get; set; }
 58
 59    /// <summary>Container storing one durable-flow ledger document per flow id.</summary>
 60    public string ContainerName { get; set; } = "asyncresponse_flow_state";
 61
 62    /// <summary>Partition-key path for the container. Default: <c>/flowId</c>.</summary>
 63    public string PartitionKeyPath { get; set; } = "/flowId";
 64
 65    /// <summary>Creates the database and container on first use.</summary>
 66    public bool AutoCreateContainer { get; set; } = true;
 67
 68    /// <summary>Optional throughput used when auto-creating the container.</summary>
 69    public int? Throughput { get; set; }
 70
 71    /// <summary>
 72    /// Maximum size in bytes of the COMPLETE ledger document accepted by writes — the item as it
 73    /// is serialized for Cosmos, with the ledger JSON embedded (and therefore escaped a second
 74    /// time) as its <c>stateJson</c> string and the sibling fields beside it. Cosmos caps the item
 75    /// as a whole at 2 MB, not the ledger inside it: a ledger whose own JSON is well under the
 76    /// budget can escape into a document over it, so the budget is enforced on what is actually
 77    /// sent. Oversized ledgers fail fast with an actionable error instead of the raw Cosmos 413
 78    /// the executor would retry into the dead-letter queue. Default: 1.9 MB (headroom under the
 79    /// item cap); <c>null</c> disables the guard.
 80    /// </summary>
 81    public long? MaxStateBytes { get; set; } = 1_900_000;
 82
 83    /// <summary>Validates option values and throws on misconfiguration.</summary>
 84    public void Validate()
 85    {
 86        if (string.IsNullOrWhiteSpace(DatabaseName))
 87            throw new InvalidOperationException($"{nameof(CosmosDurableFlowOptions)}.{nameof(DatabaseName)} must be conf
 88        if (string.IsNullOrWhiteSpace(ContainerName))
 89            throw new InvalidOperationException($"{nameof(CosmosDurableFlowOptions)}.{nameof(ContainerName)} must be con
 90        if (string.IsNullOrWhiteSpace(PartitionKeyPath) || !PartitionKeyPath.StartsWith("/", StringComparison.Ordinal))
 91            throw new InvalidOperationException($"{nameof(CosmosDurableFlowOptions)}.{nameof(PartitionKeyPath)} must sta
 92        // Every store operation addresses documents with new PartitionKey(flowId), and the ledger
 93        // document only carries the flow id under 'id' and 'flowId' — any other partition-key path
 94        // would make every write fail with a Cosmos partition-key-mismatch error at runtime, so
 95        // reject it up front instead of letting validation pass on a doomed configuration.
 96        if (!string.Equals(PartitionKeyPath, "/flowId", StringComparison.Ordinal)
 97            && !string.Equals(PartitionKeyPath, "/id", StringComparison.Ordinal))
 98        {
 99            throw new InvalidOperationException(
 100                $"{nameof(CosmosDurableFlowOptions)}.{nameof(PartitionKeyPath)} must be '/flowId' or '/id': the store " 
 101                "partitions every operation by the flow id, and the ledger document carries no other property to " +
 102                $"satisfy a container partitioned on '{PartitionKeyPath}'.");
 103        }
 104        if (Throughput is <= 0)
 105            throw new InvalidOperationException($"{nameof(CosmosDurableFlowOptions)}.{nameof(Throughput)} must be positi
 106        DurableFlowStoreShared.ValidateMaxStateBytes(MaxStateBytes, nameof(CosmosDurableFlowOptions));
 107    }
 108}
 109
 110/// <summary>Azure Cosmos DB implementation of <see cref="IFlowStateStore"/>.</summary>
 111public sealed class CosmosFlowStateStore : IFlowStateStore, IDisposable
 112{
 113    // Time authority: this store keeps the app clock (DateTime.UtcNow) for expiry and lease
 114    // comparisons. Cosmos conditional writes (ETag preconditions) evaluate client-supplied
 115    // values only — there is no server-clock expression usable inside a point write — so the
 116    // read-check-replace cycles below compare against the app clock and rely on the ETag fence
 117    // for atomicity. Multi-node deployments should keep worker clocks synchronized well inside
 118    // the lease window. (The server's own TTL sweep, by contrast, runs on the service clock.)
 119    private readonly CosmosClient _client;
 120    private readonly CosmosDurableFlowOptions _options;
 121    private readonly SemaphoreSlim _ensureGate = new(1, 1);
 122    private readonly bool _ownsClient;
 123    private volatile bool _created;
 124
 125    public CosmosFlowStateStore(CosmosClient client, IOptions<CosmosDurableFlowOptions> options, bool ownsClient = false
 126    {
 127        _client = client;
 128        _options = options.Value;
 129        _options.Validate();
 130        _ownsClient = ownsClient;
 131    }
 132
 133    public async Task<FlowState?> LoadAsync(string flowId, CancellationToken cancellationToken = default)
 134    {
 135        ArgumentException.ThrowIfNullOrWhiteSpace(flowId);
 136        var container = await GetContainerAsync(cancellationToken).ConfigureAwait(false);
 137
 138        // Point reads have no predicate, so the expiry check happens client-side on the app
 139        // clock — see the time-authority note on this class.
 140        var document = await ReadDocumentAsync(container, flowId, cancellationToken).ConfigureAwait(false);
 141        if (document is null || document.ExpiresAtUtc <= DateTime.UtcNow)
 142        {
 143            // Callers acknowledge the wake-up on null, so "absent" must not come from a read
 144            // alone. Session consistency is read-your-writes for the client that WROTE: a
 145            // different process never received the writer's session token, so its read may be
 146            // served by a replica that has not applied the create yet (a plain 404, sub-status 0
 147            // — 1002 only answers a token the replica cannot satisfy) or by one still holding an
 148            // older version whose expiry has since been extended. Inside that replication lag a
 149            // live run's only wake-up was acknowledged as "no state". The write path answers
 150            // authoritatively, and its 412 brings this client's session token up to the write
 151            // region's, so the re-read below is current.
 152            for (var attempt = 0; ; attempt++)
 153            {
 154                if (!await ExistsOnWritePathAsync(container, flowId, cancellationToken).ConfigureAwait(false))
 155                    return null;
 156
 157                document = await ReadDocumentAsync(container, flowId, cancellationToken).ConfigureAwait(false);
 158                if (document is not null)
 159                    break;
 160
 161                // Present for writes and absent for reads, repeatedly: a delete can win that race
 162                // once, not every time. This client's reads are not session-consistent with its
 163                // writes (an Eventual or Consistent Prefix account or client), so nothing it
 164                // reads can prove the run is gone.
 165                if (attempt == MaxAbsenceConfirmations - 1)
 166                {
 167                    throw new FlowStateUnreadableException(
 168                        flowId,
 169                        "the container's write path reports its document present while reads keep answering 404; the " +
 170                        "CosmosClient's reads are not session-consistent with its writes (the store needs Session consis
 171                }
 172            }
 173
 174            if (document.ExpiresAtUtc <= DateTime.UtcNow)
 175                return null;
 176        }
 177
 178        // The document exists, so a missing required field is an unreadable ledger, not an
 179        // absent one. Reporting it as absent let the executor ack the only wake-up of a run that
 180        // is still sitting in the container.
 181        if (document.Revision is not { } revision)
 182            throw new FlowStateUnreadableException(flowId, "its stored document has no revision");
 183
 184        if (string.IsNullOrEmpty(document.StateJson))
 185            throw new FlowStateUnreadableException(flowId, "its stored document has no state JSON");
 186
 187        return DurableFlowStoreShared.ReadState(flowId, document.StateJson, revision);
 188    }
 189
 190    private static async Task<CosmosFlowStateDocument?> ReadDocumentAsync(Container container, string flowId, Cancellati
 191    {
 192        try
 193        {
 194            var response = await container.ReadItemAsync<CosmosFlowStateDocument>(
 195                flowId,
 196                new PartitionKey(flowId),
 197                cancellationToken: cancellationToken).ConfigureAwait(false);
 198            return response.Resource;
 199        }
 200        catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == 0)
 201        {
 202            // Sub-status 0 is the only 404 that CAN mean "no such item". Cosmos also answers 404
 203            // for conditions where the ledger still exists — 1002 ReadSessionNotAvailable (the
 204            // replicas in reach are behind the session token this client already holds, surfaced
 205            // once the SDK's session retries exhaust) and 1003/1004 (container/database
 206            // recreated). Mapping those to null silently dropped a live run's only wake-up;
 207            // letting them throw routes the delivery through retry/dead-letter instead. Sub-status
 208            // 0 is still only what one replica says, which is why LoadAsync confirms it on the
 209            // write path. (DynamoDB pins this contract point with ConsistentRead and MongoDB with
 210            // primary reads; Cosmos cannot strengthen a read per request.)
 211            return null;
 212        }
 213    }
 214
 215    /// <inheritdoc />
 216    public void ValidateCreate(string flowId, FlowState state, TimeSpan ttl)
 217    {
 218        DurableFlowStoreShared.ValidateCreate(flowId, state, ttl);
 219        var stateJson = DurableFlowStoreShared.SerializeBounded(flowId, state, _options.MaxStateBytes, "Cosmos DB");
 220        ThrowIfDocumentTooLarge(flowId, CreateDocument(flowId, stateJson, state.Revision, ttl, DateTime.UtcNow));
 221    }
 222
 223    public async Task<bool> TryCreateAsync(string flowId, FlowState state, TimeSpan ttl, CancellationToken cancellationT
 224    {
 225        DurableFlowStoreShared.ValidateCreate(flowId, state, ttl);
 226        var stateJson = DurableFlowStoreShared.SerializeBounded(flowId, state, _options.MaxStateBytes, "Cosmos DB");
 227        var container = await GetContainerAsync(cancellationToken).ConfigureAwait(false);
 228
 229        // Bounded retry: documents carry native Cosmos TTL alongside the logical ExpiresAtUtc, so
 230        // the server's TTL sweep can physically purge an expired document between our create's
 231        // 409 and the follow-up read (or the conditional replace). A vanished conflicting
 232        // document means the slot is free — create again — not that a live competitor won.
 233        // The short growing delay lets one purge cycle finish instead of burning all attempts
 234        // inside the same inconsistency window (the Linux emulator is markedly worse here than
 235        // the service, but the race itself is real on both).
 236        for (var attempt = 0; attempt < 4; attempt++)
 237        {
 238            if (attempt > 0)
 239                await Task.Delay(TimeSpan.FromMilliseconds(50 * attempt), cancellationToken).ConfigureAwait(false);
 240
 241            var now = DateTime.UtcNow;
 242            var document = CreateDocument(flowId, stateJson, state.Revision, ttl, now);
 243            if (attempt == 0)
 244                ThrowIfDocumentTooLarge(flowId, document);
 245            try
 246            {
 247                await container.CreateItemAsync(document, new PartitionKey(flowId), cancellationToken: cancellationToken
 248                return true;
 249            }
 250            catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.Conflict)
 251            {
 252                try
 253                {
 254                    var current = await container.ReadItemAsync<CosmosFlowStateDocument>(
 255                        flowId,
 256                        new PartitionKey(flowId),
 257                        cancellationToken: cancellationToken).ConfigureAwait(false);
 258                    if (current.Resource.ExpiresAtUtc > now)
 259                        return false;
 260
 261                    await container.ReplaceItemAsync(
 262                        document,
 263                        flowId,
 264                        new PartitionKey(flowId),
 265                        new ItemRequestOptions { IfMatchEtag = current.ETag },
 266                        cancellationToken).ConfigureAwait(false);
 267                    return true;
 268                }
 269                catch (CosmosException retryEx) when (retryEx.StatusCode == HttpStatusCode.PreconditionFailed)
 270                {
 271                    // The ETag moved under us. That is either a live writer winning the id or the
 272                    // TTL purge touching the expired document; only the next read can tell them
 273                    // apart, so retry instead of conceding — a live occupant surfaces as an
 274                    // unexpired read (false) on the next attempt, a purged one as a clean create.
 275                    continue;
 276                }
 277                catch (CosmosException retryEx) when (retryEx.StatusCode == HttpStatusCode.NotFound)
 278                {
 279                    // The expired document was TTL-purged after our 409 — the id is free again.
 280                    continue;
 281                }
 282            }
 283        }
 284
 285        return false;
 286    }
 287
 288    public async Task<bool> TryUpdateAsync(
 289        string flowId,
 290        FlowState state,
 291        long expectedRevision,
 292        TimeSpan ttl,
 293        string? leaseId = null,
 294        CancellationToken cancellationToken = default)
 295    {
 296        DurableFlowStoreShared.ValidateUpdate(flowId, state, expectedRevision, ttl);
 297        var stateJson = DurableFlowStoreShared.SerializeBounded(flowId, state, _options.MaxStateBytes, "Cosmos DB");
 298        var container = await GetContainerAsync(cancellationToken).ConfigureAwait(false);
 299        for (var attempt = 0; attempt < 4; attempt++)
 300        {
 301            var now = DateTime.UtcNow;
 302            try
 303            {
 304                var current = await container.ReadItemAsync<CosmosFlowStateDocument>(
 305                    flowId,
 306                    new PartitionKey(flowId),
 307                    cancellationToken: cancellationToken).ConfigureAwait(false);
 308                var document = current.Resource;
 309                if (document.ExpiresAtUtc <= now || document.Revision != expectedRevision)
 310                    return false;
 311                // Positive form (SQL-sibling parity: `lease_id = @lease_id AND lease_expires_at_utc > now()`
 312                // is false for NULL). The negated `LeaseExpiresAtUtc <= now` was ALSO false for a
 313                // null deadline, so a document with a lease id and no expiry passed the fence.
 314                if (leaseId is not null && !(document.LeaseId == leaseId && document.LeaseExpiresAtUtc > now))
 315                    return false;
 316
 317                document.StateJson = stateJson;
 318                document.ExpiresAtUtc = DurableFlowStoreShared.AddSaturating(now, ttl);
 319                document.UpdatedAtUtc = now;
 320                document.Revision = state.Revision;
 321                document.Ttl = CosmosTtlSeconds(ttl);
 322                ThrowIfDocumentTooLarge(flowId, document);
 323                await container.ReplaceItemAsync(
 324                    document,
 325                    flowId,
 326                    new PartitionKey(flowId),
 327                    new ItemRequestOptions { IfMatchEtag = current.ETag },
 328                    cancellationToken).ConfigureAwait(false);
 329                return true;
 330            }
 331            catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == 0)
 332            {
 333                // Sub-status 0 only (LoadAsync's discriminator, for the same reason): a 404/1002
 334                // ReadSessionNotAvailable from a lagging replica names a ledger that still exists.
 335                // Reporting it as "gone" made the lease mark itself lost, the delivery redeliver,
 336                // and the step's already-performed side effect run a second time. Letting the
 337                // other sub-statuses throw routes the delivery through retry instead.
 338                return false;
 339            }
 340            catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.PreconditionFailed)
 341            {
 342                // Lease acquire/renew/release patch the document's lease fields and change its
 343                // ETag without changing the ledger revision. Re-read and retry so that benign
 344                // race is not reported as a lost execution lease; a real state race fails the
 345                // revision check above.
 346            }
 347        }
 348
 349        return false;
 350    }
 351
 352    public Task<bool> TryAcquireLeaseAsync(string flowId, string leaseId, TimeSpan leaseDuration, CancellationToken canc
 353        => UpdateLeaseAsync(flowId, leaseId, leaseDuration, acquire: true, cancellationToken);
 354
 355    public Task<bool> TryRenewLeaseAsync(string flowId, string leaseId, TimeSpan leaseDuration, CancellationToken cancel
 356        => UpdateLeaseAsync(flowId, leaseId, leaseDuration, acquire: false, cancellationToken);
 357
 358    public async Task ReleaseLeaseAsync(string flowId, string leaseId, CancellationToken cancellationToken = default)
 359    {
 360        var container = await GetContainerAsync(cancellationToken).ConfigureAwait(false);
 361        for (var attempt = 0; attempt < 4; attempt++)
 362        {
 363            try
 364            {
 365                var current = await ReadLeaseAsync(container, flowId, cancellationToken).ConfigureAwait(false);
 366                if (current is null || current.LeaseId != leaseId)
 367                    return;
 368
 369                // Every write refreshes _ts — the anchor the server-side TTL counts from — so
 370                // re-persisting the stored full-window ttl would restart the physical-retention
 371                // countdown and decouple it from the logical ExpiresAtUtc. Rewrite it from the
 372                // remaining logical window instead. (Checkpoints recompute both together in
 373                // TryUpdateAsync; only the lease paths write without moving ExpiresAtUtc.)
 374                await PatchLeaseAsync(
 375                    container,
 376                    flowId,
 377                    current.ETag,
 378                    [
 379                        PatchOperation.Set<string?>(LeaseIdPath, null),
 380                        PatchOperation.Set<DateTime?>(LeaseExpiresAtPath, null),
 381                        PatchOperation.Set(TtlPath, CosmosTtlSeconds(current.ExpiresAtUtc, DateTime.UtcNow))
 382                    ],
 383                    cancellationToken).ConfigureAwait(false);
 384                return;
 385            }
 386            catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == 0)
 387            {
 388                return;
 389            }
 390            catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.PreconditionFailed)
 391            {
 392            }
 393        }
 394    }
 395
 396    /// <inheritdoc />
 397    public async Task<FlowLeaseObservation?> ObserveLeaseAsync(string flowId, CancellationToken cancellationToken = defa
 398    {
 399        ArgumentException.ThrowIfNullOrWhiteSpace(flowId);
 400        var container = await GetContainerAsync(cancellationToken).ConfigureAwait(false);
 401
 402        // The same projecting point query the lease writes read through, so stateJson never
 403        // crosses the wire — and, unlike them, nothing here is compared with a clock: an expired
 404        // lease nobody has taken over must keep reading as the same lease, because the engine's
 405        // proof of a live holder is that two observations DIFFER. Whether it has lapsed stays
 406        // UpdateLeaseAsync's call.
 407        //
 408        // That proof only holds when each observation is at least as new as the moment it was
 409        // asked for. A waiting delivery runs in a process that never received the holder's
 410        // session token, so a bare query may be served by a lagging replica: a baseline older
 411        // than the wait makes a renewal written BEFORE the delivery arrived look like one written
 412        // while it waited, and the delivery is acknowledged as a duplicate of a holder that may
 413        // already be dead. The write-path round trip first pins the query to the write region's
 414        // progress as of now (and answers absence authoritatively, without the query).
 415        if (!await ExistsOnWritePathAsync(container, flowId, cancellationToken).ConfigureAwait(false))
 416            return FlowLeaseObservation.Unheld;
 417
 418        try
 419        {
 420            var current = await ReadLeaseAsync(container, flowId, cancellationToken).ConfigureAwait(false);
 421            return DurableFlowStoreShared.LeaseObservation(current?.LeaseId, current?.LeaseExpiresAtUtc);
 422        }
 423        catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == 0)
 424        {
 425            return FlowLeaseObservation.Unheld;
 426        }
 427    }
 428
 429    /// <summary>
 430    /// How often <see cref="LoadAsync"/> re-asks the write path when it reports the ledger present
 431    /// and the follow-up read still answers 404, before it gives up proving absence.
 432    /// </summary>
 433    private const int MaxAbsenceConfirmations = 3;
 434
 435    /// <summary>
 436    /// An <c>If-Match</c> value no document can carry (service ETags are quoted GUIDs, and the
 437    /// wildcard is <c>*</c>), so the conditional patch below can never apply.
 438    /// </summary>
 439    internal const string NeverMatchingEtag = "\"asyncresponse-consistency-barrier\"";
 440
 441    // Inert even if it could apply: no reader or writer of the ledger document knows this path.
 442    private const string BarrierPath = "/consistencyBarrier";
 443
 444    /// <summary>
 445    /// Asks the container's WRITE path whether the ledger document exists, with a conditional
 446    /// patch whose precondition cannot hold. Writes are served by the partition's write-region
 447    /// primary, never by a lagging read replica: 404 (sub-status 0) there is an authoritative
 448    /// "no such item", and 412 means the document exists. Nothing is ever written.
 449    /// <para>
 450    /// The 412 has a second effect the callers depend on. The SDK records the session token of a
 451    /// 412/409/404 response exactly as it does a successful one (StoreClient and
 452    /// GatewayStoreModel both capture it), so after this call the client's session token for the
 453    /// partition is at least the write region's progress as of now, and under Session consistency
 454    /// the NEXT read from this client cannot be served by a replica behind it — it is current, or
 455    /// it fails with 404/1002, which every path here lets throw.
 456    /// </para>
 457    /// <para>
 458    /// What that guarantees, by the client's effective consistency level: Strong reads, and
 459    /// Bounded Staleness reads served from the write region, were already current; Session (the
 460    /// account default) is made current by the token; Bounded Staleness read from another
 461    /// region, Consistent Prefix and Eventual send no session token on reads, so only the absence
 462    /// answer is authoritative there and the following read can still lag. An account with
 463    /// multiple WRITE regions has no single authoritative write path (it already lets two regions
 464    /// win the same ETag-fenced lease write and resolves them last-writer-wins), so none of this
 465    /// holds on one. The cost is one extra bodiless request, on the two decision paths only —
 466    /// never on a load that found its document.
 467    /// </para>
 468    /// </summary>
 469    private static async Task<bool> ExistsOnWritePathAsync(Container container, string flowId, CancellationToken cancell
 470    {
 471        try
 472        {
 473            await container.PatchItemAsync<CosmosFlowStateDocument>(
 474                flowId,
 475                new PartitionKey(flowId),
 476                [PatchOperation.Set(BarrierPath, 0)],
 477                new PatchItemRequestOptions { IfMatchEtag = NeverMatchingEtag, EnableContentResponseOnWrite = false },
 478                cancellationToken).ConfigureAwait(false);
 479            // Unreachable against the service; a patch that applied still proves the document exists.
 480            return true;
 481        }
 482        catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.PreconditionFailed)
 483        {
 484            return true;
 485        }
 486        catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == 0)
 487        {
 488            return false;
 489        }
 490    }
 491
 492    // JSON-pointer paths of the lease fields, matching CosmosFlowStateDocument's property names.
 493    private const string LeaseIdPath = "/leaseId";
 494    private const string LeaseExpiresAtPath = "/leaseExpiresAtUtc";
 495    private const string TtlPath = "/ttl";
 496
 497    /// <summary>
 498    /// The lease-relevant slice of one ledger document, read with a projecting point query so a
 499    /// lease acquire, heartbeat, or release never transfers <c>stateJson</c>. A point read has no
 500    /// projection — it returned the whole document, StateJson included, and the follow-up
 501    /// ReplaceItemAsync sent it all back — so an idle execution's every renewal (default: each
 502    /// 20 seconds) moved and re-serialized the full ledger twice, proportional to its size.
 503    /// Together with the conditional patches below, lease maintenance now costs O(lease fields)
 504    /// on the wire regardless of ledger size. RU cost still depends on the service's accounting
 505    /// for the loaded document; measure it (docs/durable-flow-state-stores.md).
 506    /// <para>
 507    /// Only the SQL text is shared. A <see cref="QueryDefinition"/> is a mutable parameter bag —
 508    /// <c>WithParameter</c> replaces the named parameter in place and returns the same instance —
 509    /// so one static definition parameterized per call handed concurrent lease operations each
 510    /// other's ids: flow A's query could execute with <c>@id = B</c> under A's partition key,
 511    /// return no document, and fail a healthy renewal (which abandons and replays the run). Every
 512    /// call builds its own definition.
 513    /// </para>
 514    /// </summary>
 515    private const string LeaseProjectionSql =
 516        "SELECT c.id, c._etag, c.expiresAtUtc, c.revision, c.leaseId, c.leaseExpiresAtUtc FROM c WHERE c.id = @id";
 517
 518    private static async Task<CosmosLeaseProjection?> ReadLeaseAsync(Container container, string flowId, CancellationTok
 519    {
 520        using var iterator = container.GetItemQueryIterator<CosmosLeaseProjection>(
 521            new QueryDefinition(LeaseProjectionSql).WithParameter("@id", flowId),
 522            requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(flowId), MaxItemCount = 1 });
 523        while (iterator.HasMoreResults)
 524        {
 525            var page = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false);
 526            foreach (var projection in page)
 527            {
 528                if (string.IsNullOrEmpty(projection.ETag))
 529                {
 530                    // The fence for every lease write. A projection without it cannot be acted
 531                    // on safely, and silently treating it as "not held" would let the executor
 532                    // acknowledge a wake-up as a duplicate against a run nobody holds.
 533                    throw new InvalidOperationException(
 534                        $"The Cosmos DB durable-flow store's lease query for '{flowId}' returned no _etag; the registere
 535                }
 536
 537                return projection;
 538            }
 539        }
 540
 541        return null;
 542    }
 543
 544    /// <summary>
 545    /// A conditional partial update of the lease fields: fenced by the projection's ETag exactly
 546    /// as the replace was, with no document content in the response (there is nothing the
 547    /// caller reads back).
 548    /// </summary>
 549    private static Task PatchLeaseAsync(
 550        Container container,
 551        string flowId,
 552        string etag,
 553        IReadOnlyList<PatchOperation> operations,
 554        CancellationToken cancellationToken)
 555        => container.PatchItemAsync<CosmosFlowStateDocument>(
 556            flowId,
 557            new PartitionKey(flowId),
 558            operations,
 559            new PatchItemRequestOptions { IfMatchEtag = etag, EnableContentResponseOnWrite = false },
 560            cancellationToken);
 561
 562    public async Task<bool> TryDeleteAsync(string flowId, CancellationToken cancellationToken = default)
 563    {
 564        ArgumentException.ThrowIfNullOrWhiteSpace(flowId);
 565        var container = await GetContainerAsync(cancellationToken).ConfigureAwait(false);
 566
 567        try
 568        {
 569            await container.DeleteItemAsync<CosmosFlowStateDocument>(
 570                flowId,
 571                new PartitionKey(flowId),
 572                cancellationToken: cancellationToken).ConfigureAwait(false);
 573            return true;
 574        }
 575        catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == 0)
 576        {
 577            return false;
 578        }
 579    }
 580
 581    private async Task<Container> GetContainerAsync(CancellationToken cancellationToken)
 582    {
 583        if (!_created)
 584            await EnsureCreatedAsync(cancellationToken).ConfigureAwait(false);
 585
 586        return _client.GetContainer(_options.DatabaseName, _options.ContainerName);
 587    }
 588
 589    private async Task EnsureCreatedAsync(CancellationToken cancellationToken)
 590    {
 591        if (_created)
 592            return;
 593
 594        await _ensureGate.WaitAsync(cancellationToken).ConfigureAwait(false);
 595        try
 596        {
 597            if (_created)
 598                return;
 599
 600            ContainerResponse container;
 601            if (_options.AutoCreateContainer)
 602            {
 603                var database = await _client.CreateDatabaseIfNotExistsAsync(
 604                    _options.DatabaseName,
 605                    cancellationToken: cancellationToken).ConfigureAwait(false);
 606
 607                // DefaultTimeToLive = -1 enables per-item TTL without a container-wide default.
 608                var properties = new ContainerProperties(_options.ContainerName, _options.PartitionKeyPath)
 609                {
 610                    DefaultTimeToLive = -1
 611                };
 612                container = await database.Database.CreateContainerIfNotExistsAsync(
 613                    properties,
 614                    _options.Throughput,
 615                    cancellationToken: cancellationToken).ConfigureAwait(false);
 616            }
 617            else
 618            {
 619                container = await _client
 620                    .GetContainer(_options.DatabaseName, _options.ContainerName)
 621                    .ReadContainerAsync(cancellationToken: cancellationToken)
 622                    .ConfigureAwait(false);
 623            }
 624
 625            if (!string.Equals(container.Resource.PartitionKeyPath, _options.PartitionKeyPath, StringComparison.Ordinal)
 626                throw new InvalidOperationException(
 627                    $"Cosmos container '{_options.ContainerName}' uses partition key '{container.Resource.PartitionKeyPa
 628                    $"but '{_options.PartitionKeyPath}' is required.");
 629            if (container.Resource.DefaultTimeToLive is null)
 630                throw new InvalidOperationException(
 631                    $"Cosmos container '{_options.ContainerName}' does not have TTL enabled. " +
 632                    "Enable container TTL (DefaultTimeToLive = -1) before using it for durable flows.");
 633
 634            ValidateHostSerializer();
 635            _created = true;
 636        }
 637        finally
 638        {
 639            _ensureGate.Release();
 640        }
 641    }
 642
 643    /// <summary>
 644    /// Fails provisioning fast when a host-registered <see cref="CosmosClient"/> carries a custom
 645    /// serializer that does not honor the flow-state document's JSON property names.
 646    /// <see cref="CosmosFlowStateDocument"/> is attributed for both Newtonsoft.Json
 647    /// (<c>[JsonProperty]</c>) and System.Text.Json (<c>[JsonPropertyName]</c>), so the SDK's
 648    /// default serializer and STJ-based serializers both map correctly; a serializer honoring
 649    /// neither would write documents whose <c>id</c> Cosmos rejects — or, worse, whose fields
 650    /// silently round-trip as nulls. The probe serializes a sentinel document through the host's
 651    /// serializer and verifies the wire property names survive.
 652    /// </summary>
 653    private void ValidateHostSerializer()
 654    {
 655        if (_client.ClientOptions?.Serializer is not { } serializer)
 656            return; // The SDK default serializer honors [JsonProperty]; nothing to probe.
 657
 658        var now = DateTime.UtcNow;
 659        using var stream = serializer.ToStream(new CosmosFlowStateDocument
 660        {
 661            Id = "asyncresponse-serializer-probe",
 662            FlowId = "asyncresponse-serializer-probe",
 663            StateJson = "{}",
 664            ExpiresAtUtc = now,
 665            UpdatedAtUtc = now,
 666            Revision = 0,
 667            LeaseId = "probe",
 668            LeaseExpiresAtUtc = now,
 669            Ttl = 1
 670        });
 671        using var probe = System.Text.Json.JsonDocument.Parse(stream);
 672        if (!probe.RootElement.TryGetProperty("id", out _) || !probe.RootElement.TryGetProperty("leaseExpiresAtUtc", out
 673        {
 674            throw new InvalidOperationException(
 675                $"The registered CosmosClient's serializer ({serializer.GetType().Name}) does not honor the durable-flow
 676                "JSON property names ('id', 'flowId', 'leaseExpiresAtUtc', ...). Flow-state documents would be written w
 677                "property names and could not be read back. Configure the serializer to honor Newtonsoft.Json [JsonPrope
 678                "System.Text.Json [JsonPropertyName] attributes, or let the SDK use its default serializer.");
 679        }
 680    }
 681
 682    private async Task<bool> UpdateLeaseAsync(
 683        string flowId,
 684        string leaseId,
 685        TimeSpan leaseDuration,
 686        bool acquire,
 687        CancellationToken cancellationToken)
 688    {
 689        DurableFlowStoreShared.ValidateLeaseArgs(flowId, leaseId, leaseDuration);
 690
 691        var container = await GetContainerAsync(cancellationToken).ConfigureAwait(false);
 692        for (var attempt = 0; attempt < 4; attempt++)
 693        {
 694            var now = DateTime.UtcNow;
 695            try
 696            {
 697                var document = await ReadLeaseAsync(container, flowId, cancellationToken).ConfigureAwait(false);
 698                if (document is null || document.ExpiresAtUtc <= now || document.Revision is null)
 699                    return false;
 700                if (acquire)
 701                {
 702                    if (document.LeaseId is not null && document.LeaseId != leaseId && document.LeaseExpiresAtUtc > now)
 703                        return false;
 704                }
 705                else if (!(document.LeaseId == leaseId && document.LeaseExpiresAtUtc > now))
 706                {
 707                    // Positive form: a renewal needs a live deadline, and a null one is not live.
 708                    return false;
 709                }
 710
 711                // Same _ts realignment as ReleaseLeaseAsync: a lease heartbeat writes the
 712                // document without moving ExpiresAtUtc, so it must not restart the server TTL's
 713                // full retention window.
 714                await PatchLeaseAsync(
 715                    container,
 716                    flowId,
 717                    document.ETag,
 718                    [
 719                        PatchOperation.Set(LeaseIdPath, leaseId),
 720                        PatchOperation.Set(LeaseExpiresAtPath, DurableFlowStoreShared.AddSaturating(now, leaseDuration))
 721                        PatchOperation.Set(TtlPath, CosmosTtlSeconds(document.ExpiresAtUtc, now))
 722                    ],
 723                    cancellationToken).ConfigureAwait(false);
 724                return true;
 725            }
 726            catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == 0)
 727            {
 728                return false;
 729            }
 730            catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.PreconditionFailed)
 731            {
 732            }
 733        }
 734
 735        return false;
 736    }
 737
 738    /// <summary>
 739    /// Enforces <see cref="CosmosDurableFlowOptions.MaxStateBytes"/> on the document as Cosmos
 740    /// will receive it. <see cref="DurableFlowStoreShared.SerializeBounded"/> already refused a
 741    /// ledger whose own JSON is over the budget — a cheap first check, since the document can only
 742    /// be larger — but the ledger travels inside the document as a string value, so every quote
 743    /// and backslash in it is escaped again: a 1.2 MB ledger made of escaped quotes is a 2.4 MB
 744    /// document, accepted by the inner check and refused by Cosmos's 2 MB item cap on every
 745    /// retry. Measured through the host's own serializer when one is registered (its escaping
 746    /// and property naming are what go on the wire), else through the SDK default's
 747    /// Newtonsoft-based shape.
 748    /// </summary>
 749    private void ThrowIfDocumentTooLarge(string flowId, CosmosFlowStateDocument document)
 750    {
 751        if (_options.MaxStateBytes is not { } limit)
 752            return;
 753
 754        var size = MeasureDocumentBytes(document);
 755        if (size > limit)
 756            throw new FlowStateTooLargeException(flowId, size, limit, "Cosmos DB");
 757    }
 758
 759    private long MeasureDocumentBytes(CosmosFlowStateDocument document)
 760    {
 761        if (_client.ClientOptions?.Serializer is { } serializer)
 762        {
 763            using var stream = serializer.ToStream(document);
 764            if (stream.CanSeek)
 765                return stream.Length;
 766
 767            long total = 0;
 768            var buffer = ArrayPool<byte>.Shared.Rent(16 * 1024);
 769            try
 770            {
 771                int read;
 772                while ((read = stream.Read(buffer, 0, buffer.Length)) > 0)
 773                    total += read;
 774            }
 775            finally
 776            {
 777                ArrayPool<byte>.Shared.Return(buffer);
 778            }
 779
 780            return total;
 781        }
 782
 783        return MeasureDefaultDocumentBytes(document, _client.ClientOptions?.SerializerOptions);
 784    }
 785
 786    // The SDK's default wire shape, written as scalars rather than by reflecting over the
 787    // document. JsonTextWriter preserves Newtonsoft's escaping/date rules without IL2026/3050.
 788    internal static long MeasureDefaultDocumentBytes(CosmosFlowStateDocument document, CosmosSerializationOptions? optio
 789    {
 790        using var text = new StringWriter(System.Globalization.CultureInfo.InvariantCulture);
 791        using var writer = new JsonTextWriter(text) { Formatting = options?.Indented == true ? Formatting.Indented : For
 792        writer.WriteStartObject();
 793        writer.WritePropertyName("id"); writer.WriteValue(document.Id);
 794        writer.WritePropertyName("flowId"); writer.WriteValue(document.FlowId);
 795        writer.WritePropertyName("stateJson"); writer.WriteValue(document.StateJson);
 796        writer.WritePropertyName("expiresAtUtc"); writer.WriteValue(document.ExpiresAtUtc);
 797        writer.WritePropertyName("updatedAtUtc"); writer.WriteValue(document.UpdatedAtUtc);
 798        if (document.Revision is not null || options?.IgnoreNullValues != true)
 799        {
 800            writer.WritePropertyName("revision"); writer.WriteValue(document.Revision);
 801        }
 802        if (document.LeaseId is not null)
 803        {
 804            writer.WritePropertyName("leaseId"); writer.WriteValue(document.LeaseId);
 805        }
 806        if (document.LeaseExpiresAtUtc is not null)
 807        {
 808            writer.WritePropertyName("leaseExpiresAtUtc"); writer.WriteValue(document.LeaseExpiresAtUtc);
 809        }
 810        if (document.Ttl is not null)
 811        {
 812            writer.WritePropertyName("ttl"); writer.WriteValue(document.Ttl);
 813        }
 814        writer.WriteEndObject();
 815        writer.Flush();
 816        return Encoding.UTF8.GetByteCount(text.ToString());
 817    }
 818
 819    private static CosmosFlowStateDocument CreateDocument(string flowId, string stateJson, long revision, TimeSpan ttl, 
 820        => new()
 821        {
 822            Id = flowId,
 823            FlowId = flowId,
 824            StateJson = stateJson,
 825            ExpiresAtUtc = DurableFlowStoreShared.AddSaturating(now, ttl),
 826            UpdatedAtUtc = now,
 827            Revision = revision,
 828            // Cosmos reaps the item itself once container TTL is enabled. Ceiling keeps the
 829            // server-side TTL from being shorter than the requested duration.
 830            Ttl = CosmosTtlSeconds(ttl)
 831        };
 832
 833    /// <summary>Per-item TTL in whole seconds, rounded up and saturated at int.MaxValue (~68 years) for absurd expiries
 834    private static int CosmosTtlSeconds(TimeSpan ttl)
 835        => (int)Math.Min(Math.Ceiling(ttl.TotalSeconds), int.MaxValue);
 836
 837    /// <summary>
 838    /// Remaining per-item TTL in whole seconds until <paramref name="expiresAtUtc"/>, rounded up
 839    /// and floored at 1 (Cosmos rejects 0). Used by replaces that keep the logical expiry in place:
 840    /// an already-due document collapses to the shortest legal TTL so the next sweep purges it
 841    /// instead of the replace granting it a fresh retention window.
 842    /// </summary>
 843    private static int CosmosTtlSeconds(DateTime expiresAtUtc, DateTime now)
 844        => (int)Math.Min(Math.Max(Math.Ceiling((expiresAtUtc - now).TotalSeconds), 1), int.MaxValue);
 845
 846    /// <summary>Disposes the Cosmos client when the store created (and therefore owns) it.</summary>
 847    public void Dispose()
 848    {
 849        _ensureGate.Dispose();
 850        if (_ownsClient)
 851            _client.Dispose();
 852    }
 853}
 854
 855/// <summary>
 856/// One durable-flow ledger document. Attributed for BOTH Newtonsoft.Json and System.Text.Json:
 857/// the Cosmos SDK's default serializer is Newtonsoft-based, but hosts may register a
 858/// <see cref="CosmosClient"/> with an STJ-based serializer — with single-stack attributes such a
 859/// client would silently write PascalCase property names (breaking <c>id</c> and every read
 860/// back). <see cref="CosmosFlowStateStore"/> additionally probes custom serializers at
 861/// provisioning time and fails fast when neither attribute set is honored.
 862/// </summary>
 863internal sealed class CosmosFlowStateDocument
 864{
 865    [JsonProperty("id")]
 866    [System.Text.Json.Serialization.JsonPropertyName("id")]
 867    public string Id { get; set; } = "";
 868
 869    [JsonProperty("flowId")]
 870    [System.Text.Json.Serialization.JsonPropertyName("flowId")]
 871    public string FlowId { get; set; } = "";
 872
 873    [JsonProperty("stateJson")]
 874    [System.Text.Json.Serialization.JsonPropertyName("stateJson")]
 875    public string StateJson { get; set; } = "";
 876
 877    [JsonProperty("expiresAtUtc")]
 878    [System.Text.Json.Serialization.JsonPropertyName("expiresAtUtc")]
 879    public DateTime ExpiresAtUtc { get; set; }
 880
 881    [JsonProperty("updatedAtUtc")]
 882    [System.Text.Json.Serialization.JsonPropertyName("updatedAtUtc")]
 883    public DateTime UpdatedAtUtc { get; set; }
 884
 885    [JsonProperty("revision")]
 886    [System.Text.Json.Serialization.JsonPropertyName("revision")]
 887    public long? Revision { get; set; }
 888
 889    [JsonProperty("leaseId", NullValueHandling = NullValueHandling.Ignore)]
 890    [System.Text.Json.Serialization.JsonPropertyName("leaseId")]
 891    [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritin
 892    public string? LeaseId { get; set; }
 893
 894    [JsonProperty("leaseExpiresAtUtc", NullValueHandling = NullValueHandling.Ignore)]
 895    [System.Text.Json.Serialization.JsonPropertyName("leaseExpiresAtUtc")]
 896    [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritin
 897    public DateTime? LeaseExpiresAtUtc { get; set; }
 898
 899    /// <summary>Cosmos per-item TTL in seconds; honored once the container enables TTL.</summary>
 900    [JsonProperty("ttl", NullValueHandling = NullValueHandling.Ignore)]
 901    [System.Text.Json.Serialization.JsonPropertyName("ttl")]
 902    [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritin
 903    public int? Ttl { get; set; }
 904}
 905
 906/// <summary>
 907/// The lease slice of <see cref="CosmosFlowStateDocument"/> plus the document's <c>_etag</c>, as
 908/// returned by the store's projecting lease query — everything a lease acquire, renewal, or
 909/// release decides on and fences with, and nothing else (no <c>stateJson</c>). Attributed for
 910/// both serializer stacks for the same reason the document is.
 911/// </summary>
 912internal sealed class CosmosLeaseProjection
 913{
 914    [JsonProperty("id")]
 915    [System.Text.Json.Serialization.JsonPropertyName("id")]
 916    public string Id { get; set; } = "";
 917
 918    /// <summary>The fence for every lease write; the store refuses to act on a projection without one.</summary>
 919    [JsonProperty("_etag")]
 920    [System.Text.Json.Serialization.JsonPropertyName("_etag")]
 921    public string ETag { get; set; } = "";
 922
 923    [JsonProperty("expiresAtUtc")]
 924    [System.Text.Json.Serialization.JsonPropertyName("expiresAtUtc")]
 925    public DateTime ExpiresAtUtc { get; set; }
 926
 927    [JsonProperty("revision")]
 928    [System.Text.Json.Serialization.JsonPropertyName("revision")]
 929    public long? Revision { get; set; }
 930
 931    [JsonProperty("leaseId")]
 932    [System.Text.Json.Serialization.JsonPropertyName("leaseId")]
 933    public string? LeaseId { get; set; }
 934
 935    [JsonProperty("leaseExpiresAtUtc")]
 936    [System.Text.Json.Serialization.JsonPropertyName("leaseExpiresAtUtc")]
 937    public DateTime? LeaseExpiresAtUtc { get; set; }
 938}
 939}