[Internal] Upgrade Resiliency: Fixes Duplicate Channel and Task Creation.

[kundadebdatta]

Pull Request Template

Description

Scope: Applicable when the "Advanced Replica Selection" feature is enabled using the environment variable AZURE_COSMOS_REPLICA_VALIDATION_ENABLED = True

Problem: One or our recent conversations with the IC3 team has helped us unveiling a potential issue with the open connection async flow. The team migrated to the latest preview sdk version 3.35.2-preview and experienced a sudden and consistent spike in number of tcp connections since the cosmos client initialization. To understand this in more detail, please take a look at the number of open tcp connections graph below.

It appears that during the time of impact, the RPS was around 15000 Reqs/Sec, and the number of tcp connections has doubled compared to their historical trends.

Analysis : The reason, this is happening, is because it is possible that while validating both the Unknown and Unhealthy replicas, during the OpenChannelAsync() stage, in the LoadBalancingPartition, today there is no way to validate if there is a connection already made to the given endpoint of the BE replica, thus if there are N number of parallel requests land to this method, they all ended up opening N tcp connections for the given endpoint. This is what is causing the sudden spike on the number of open tcp connections.

        internal Task OpenChannelAsync(Guid activityId)
        {
            IChannel channel = null;
            this.capacityLock.EnterWriteLock();
            try
            {
                if (this.capacity < this.maxCapacity)
                {
                    channel = this.OpenChannelAndIncrementCapacity(
                        activityId: activityId);
                }
                else
                {
                   ....
                }
            }
         }

Solution: The code changes in this PR is doing the following -

• It uses the latest version of the Cosmos.Direct version 3.31.5 package, that refactors the behavior in the LoadBalancingPartition, and adds an extra check to validate if the openChannels dictionary in the LoadBalancingPartiton contains a healthy LbChannelState for a specific endpoint. If there exists a healthy channel established to that endpoint, then we skip the open connection creation for the given endpoint. This basically means, if there is a healthy connection present for a given endpoint, no new connection will required to be opened.

• Additionally, this PR adds an additional RefreshAsync() method in the AsyncLazyWithRefreshTask<T> in the AsyncCacheNonBlocking<K, V> that maintains a reference to the background task. If the task is already running for the cache key, then any future Refresh requests will be a no-op. This dramatically reduces the background task creation, thus reducing the number of duplicate replica validations for the same partition key range.

• This PR also reduces the replica validation scope, only to the Unhealthy replicas by default. The Unknown replicas will be validated only when the CosmosClient is initialized with CreateAndInitializeAsync() flow.

Detailed Sequence Diagram: Below is the sequence diagram for validating the channel health before opening the connections.

sequenceDiagram
    participant X as GatewayAddressCache <br> [v3 Code]
    participant A as RntbdOpenConnectionHandler <br> [Direct Code]
    participant B as Rntbd.TransportClient <br> [Direct Code]
    participant C as ChannelDictionary <br> [Direct Code]
    participant D as LoadBalancingChannel <br> [Direct Code]
    participant E as LoadBalancingPartition <br> [Direct Code]
    participant F as Channel <br> [Direct Code]
    Note over X: The open connection process is <br> triggered as a background task.    
    X->>A: 1. TryOpenRntbdChannelsAsync <br> (transport addresses)
    Note over A: The incoming requests <br> acquires the semaphore <br> and opens channel.    
    A->>B: 2. OpenConnectionAsync()
    Note over B: Rntbd.TransportClient <br> implements abstract class <br> Microsoft.Azure.Documents.TransportClient
    B->>C: 3. GetChannel <br> (validationRequired: true)
    C-->>B: 4. Gets a <br> LoadBalancingChannel
    B->>D: 5. OpenChannelAsync() 
    D->>E: 6. OpenChannelAsync()
    E->>E: 7. Are there any healthy LbChannelState <br> present in the openChannels dictionary ?
    E-->>D: 8. Yes - Healthy LbChannelState present. <br> Skip new channel creation.
    E->>F: 9. No - healthy LbChannelState present. <br> Create Channel using OpenChannelAsync().
    F->>E: 10. Returns Channel Init Task.   
    E-->>D: 11. Returns Channel Init Task.
    D-->>C: 12. Returns Channel Init <br> Task Or No-Op.
    C-->>B: 13. Returns Channel Init <br> Task Or No-Op.
    B-->>A: 14. Returns Channel Init <br> Task Or No-Op.
    A-->>A: 15. a) If successful, mark transport address to Connected <br> b) If exception is caught, mark transport address to Unhealthy.

Loading

Benchmark Results: Below are the benchmark results before and after the fix.

VM Configuration: SKU: Standard D16s v3, vCPU: 16, RAM: 64 GB.
VM Location: West US 2.
Cosmos Test Account Location: West US 2.
Cosmos Test Account RU Count: 500K.
Upgrade Duration: 5 hours.
Upgrade Domains: 20
Test Run Time: 6 Hours.
RPS: Approx. 5000.

TCP Connection Count Before the Fix:

TCP Connection Count After the Fix:

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)

Closing issues

To automatically close an issue: closes #4033

[FabianMeiswinkel]

LGTM - Thanks. Very good PR description - thanks for setting an example here!

[kundadebdatta]

[like] Debdatta Kunda reacted to your message:

…

________________________________ From: Fabian Meiswinkel ***@***.***> Sent: Friday, November 3, 2023 10:53:24 PM To: Azure/azure-cosmos-dotnet-v3 ***@***.***> Cc: Debdatta Kunda ***@***.***>; Assign ***@***.***> Subject: Re: [Azure/azure-cosmos-dotnet-v3] [Internal] Upgrade Resiliency: Fixes Duplicate Channel and Task Creation. (PR #4123) @FabianMeiswinkel approved this pull request. LGTM - Thanks. Very good PR description - thanks for setting an example here! — Reply to this email directly, view it on GitHub<#4123 (review)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AU2KHTPTCIWSBM27GTFQL7DYCVYWJAVCNFSM6AAAAAA57T6UFOVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMYTOMJTGYZTCNZXGU>. You are receiving this because you were assigned.Message ID: ***@***.***>

[xinlian12]

should not this semaphore be released only when the refresh completed?

[kundadebdatta]

Discussed this offline. The below code should cover this:

if (addresses
    .Get(Protocol.Tcp)
    .ReplicaTransportAddressUris
    .Any(x => x.ShouldRefreshHealthStatus()))

The semaphore will be released as soon as the task is successfully scheduled. and the above check is good enough to block other parallel threads to create duplicate tasks.

[kirankumarkolli]

nit: try-finally be moved inside IF (code refractoring and fainally clause then don't need explicit check

[kundadebdatta]

Agreed. I will refactor this with the next PR to master. Good catch!

[kirankumarkolli]

Clarification: What will happen if this check not exists?

[kundadebdatta]

If the below condition doesn't exist, then there is no gate to detect if other threads has acquired the semaphore in the past and already scheduled a background refresh task. So, without the check, the try block will get executed.

        if (slimAcquired)
        {
            this.serverPartitionAddressCache.Refresh(
                key: partitionKeyRangeIdentity,
                singleValueInitFunc: (currentCachedValue) => this.GetAddressesForRangeIdAsync(
                    request,
                    cachedAddresses: currentCachedValue,
                    partitionKeyRangeIdentity.CollectionRid,
                    partitionKeyRangeIdentity.PartitionKeyRangeId,
                    forceRefresh: true));
        }

Also, the code snippet bool slimAcquired = await this.semaphore.WaitAsync(0); checks the semaphore hook is acquired at that point of time. If it's already acquired, then the thread immediately returns a false and continue it's execution. This guarantees that no duplicate tasks were created.