Avoid NullPointerException when x-region calls are made under high concurrency when using region scoped session container. (#46758)

* Make session container and bloom filter insertions atomic.

* Make session container and bloom filter insertions atomic.

* Make session container and bloom filter insertions atomic.

* Fixing tests.

* Ensure once a global session is used to not go back to using region-scoped session tokens.

* Ensure once a global session is used to not go back to using region-scoped session tokens.

* Ensure bloom filter is updated atomically with session container update.

* Add a test which concurrently updates and reads from the RegionScopedSessionContainer which "occasional" cross region calls.

* Fixing tests.

* Clean up PR.

* Add user-agent suffix and try-catch block in PartitionScopedRegionLevelProgress.

* Track globalLsn progress when hub region is seen in bloom filter.

* Centralize fallback to global session token flows and introduce use hub flow.

* Fixing baseSessionToken resolution flow.

* Enhancing RegionScopedSessionContainerConcurrencyTest to make lsns decrease from service perspective yet session token resolutions client-side should be monotonically increasing.

* Enhancing RegionScopedSessionContainerConcurrencyTest to make lsns decrease from service perspective yet session token resolutions client-side should be monotonically increasing.

* Enhancing RegionScopedSessionContainerConcurrencyTest to make lsns decrease from service perspective yet session token resolutions client-side should be monotonically increasing.

* Add license header.

* Fixing compilation errors.

* Addressing review comments.

* Addressing review comments.

* Addressing review comments.

* Updating CHANGELOG.md

Author: jeet1995

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RegionScopedSessionContainerConcurrencyTest.java

@@ -480,8 +480,8 @@ public void setSessionToken_NoSessionTokenForPartitionKeyRangeId() throws Except
 
         assertThat(globalProgress).isNotNull();
         assertThat(globalProgress.getSessionToken()).isNotNull();
-        assertThat(globalProgress.getSessionToken().convertToString()).isNotNull();
-        assertThat(globalProgress.getSessionToken().convertToString()).isEqualTo(sessionToken);
+        assertThat(globalProgress.getSessionToken().get().convertToString()).isNotNull();
+        assertThat(globalProgress.getSessionToken().get().convertToString()).isEqualTo(sessionToken);
 
         RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document,
             collectionName + "/docs",  Utils.getUTF8Bytes(""), new HashMap<>());

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RegionScopedSessionContainerTest.java

@@ -8,6 +8,7 @@
 
 #### Bugs Fixed
 * Fixed a possible memory leak (Netty buffers) in Gateway mode caused by a race condition when timeouts are happening. - [47228](https://github.com/Azure/azure-sdk-for-java/pull/47228) and [47251](https://github.com/Azure/azure-sdk-for-java/pull/47251)
+* Fixed `NullPointerException` in region-scoped session container when false positive regions are reported for session progress tracking. [PR 46758](https://github.com/Azure/azure-sdk-for-java/pull/46758)
 
 #### Other Changes
 * Enabled hostname validation for RNTBD connections to backend - [PR 47111](https://github.com/Azure/azure-sdk-for-java/pull/47111)

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -466,6 +466,7 @@ public static class SubStatusCodes {
         public static final int INVALID_RESULT = 20910;
         public static final int CLOSED_CLIENT = 20912;
         public static final int PPCB_INVALID_STATE = 20913;
+        public static final int REGION_SCOPED_SESSION_CONTAINER_IN_BAD_STATE = 20914;
 
         //SDK Codes (Server)
         // IMPORTANT - whenever possible use consistency substatus codes that .Net SDK also uses

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java

@@ -18,6 +18,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -201,13 +202,6 @@ public ISessionToken resolvePartitionLocalSessionToken(RxDocumentServiceRequest
             this.firstPreferredReadableRegionCached.set(extractFirstEffectivePreferredReadableRegion());
         }
 
-        boolean shouldUseBloomFilter = shouldUseBloomFilter(
-            request,
-            partitionKeyRangeId,
-            partitionKeyInternal,
-            partitionKeyDefinition,
-            partitionScopedRegionLevelProgress);
-
         return SessionTokenHelper.resolvePartitionLocalSessionToken(
             request,
             this.partitionKeyBasedBloomFilter,
@@ -217,7 +211,7 @@ public ISessionToken resolvePartitionLocalSessionToken(RxDocumentServiceRequest
             collectionRid,
             partitionKeyRangeId,
             this.firstPreferredReadableRegionCached.get(),
-            shouldUseBloomFilter);
+            this.globalEndpointManager);
     }
 
     @Override
@@ -339,41 +333,6 @@ private void setSessionToken(RxDocumentServiceRequest request, ResourceId resour
         }
     }
 
-    private void recordPartitionKeyInBloomFilter(
-        RxDocumentServiceRequest request,
-        Long collectionRid,
-        String regionRoutedTo,
-        PartitionKeyInternal partitionKeyInternal,
-        PartitionKeyDefinition partitionKeyDefinition) {
-
-        if (Strings.isNullOrEmpty(this.firstPreferredReadableRegionCached.get())) {
-            this.firstPreferredReadableRegionCached.set(extractFirstEffectivePreferredReadableRegion());
-        }
-
-        this.partitionKeyBasedBloomFilter.tryRecordPartitionKey(
-            request,
-            collectionRid,
-            this.firstPreferredReadableRegionCached.get(),
-            regionRoutedTo,
-            partitionKeyInternal,
-            partitionKeyDefinition);
-    }
-
-    private void recordRegionScopedSessionToken(
-        RxDocumentServiceRequest request,
-        PartitionScopedRegionLevelProgress partitionScopedRegionLevelProgress,
-        ISessionToken parsedSessionToken,
-        String partitionKeyRangeId,
-        String regionRoutedTo) {
-
-        partitionScopedRegionLevelProgress.tryRecordSessionToken(
-            request,
-            parsedSessionToken,
-            partitionKeyRangeId,
-            this.firstPreferredReadableRegionCached.get(),
-            regionRoutedTo);
-    }
-
     private void addSessionTokenAndTryRecordEpkInBloomFilter(RxDocumentServiceRequest request, ResourceId resourceId, String partitionKeyRangeId, ISessionToken parsedSessionToken) {
 
         final Long collectionResourceId = resourceId.getUniqueDocumentCollectionId();
@@ -399,27 +358,14 @@ private void addSessionTokenAndTryRecordEpkInBloomFilter(RxDocumentServiceReques
 
         if (partitionScopedRegionLevelProgress != null) {
 
-            if (shouldUseBloomFilter(
-                request,
-                partitionKeyRangeId,
-                partitionKeyInternal,
-                partitionKeyDefinition,
-                partitionScopedRegionLevelProgress)) {
-
-                this.recordPartitionKeyInBloomFilter(
-                    request,
-                    collectionResourceId,
-                    regionRoutedTo,
-                    partitionKeyInternal.v,
-                    partitionKeyDefinition.v);
-            }
-
             this.recordRegionScopedSessionToken(
                 request,
                 partitionScopedRegionLevelProgress,
                 parsedSessionToken,
+                collectionResourceId,
                 partitionKeyRangeId,
-                regionRoutedTo);
+                regionRoutedTo
+            );
 
         } else {
             this.collectionResourceIdToPartitionScopedRegionLevelProgress.compute(
@@ -444,25 +390,31 @@ private void addSessionTokenAndTryRecordEpkInBloomFilter(RxDocumentServiceReques
                     request,
                     partitionScopedRegionLevelProgress,
                     parsedSessionToken,
+                    collectionResourceId,
                     partitionKeyRangeId,
-                    regionRoutedTo);
+                    regionRoutedTo
+                );
             }
+        }
+    }
 
-            if (shouldUseBloomFilter(
-                request,
-                partitionKeyRangeId,
-                partitionKeyInternal,
-                partitionKeyDefinition,
-                partitionScopedRegionLevelProgress)) {
+    private void recordRegionScopedSessionToken(
+        RxDocumentServiceRequest request,
+        PartitionScopedRegionLevelProgress partitionScopedRegionLevelProgress,
+        ISessionToken parsedSessionToken,
+        Long collectionRid,
+        String partitionKeyRangeId,
+        String regionRoutedTo) {
 
-                this.recordPartitionKeyInBloomFilter(
-                    request,
-                    collectionResourceId,
-                    regionRoutedTo,
-                    partitionKeyInternal.v,
-                    partitionKeyDefinition.v);
-            }
-        }
+        partitionScopedRegionLevelProgress.tryRecordSessionToken(
+            request,
+            parsedSessionToken,
+            collectionRid,
+            partitionKeyRangeId,
+            this.firstPreferredReadableRegionCached.get(),
+            regionRoutedTo,
+            this.partitionKeyBasedBloomFilter,
+            this.globalEndpointManager);
     }
 
     private String getCombinedSessionToken(PartitionScopedRegionLevelProgress partitionScopedRegionLevelProgress) {
@@ -472,53 +424,28 @@ private String getCombinedSessionToken(PartitionScopedRegionLevelProgress partit
         StringBuilder result = new StringBuilder();
         if (tokens != null) {
             for (Iterator<Map.Entry<String, ConcurrentHashMap<String, PartitionScopedRegionLevelProgress.RegionLevelProgress>>> iterator = tokens.entrySet().iterator(); iterator.hasNext(); ) {
-                Map.Entry<String,  ConcurrentHashMap<String, PartitionScopedRegionLevelProgress.RegionLevelProgress>> entry = iterator.next();
+                Map.Entry<String, ConcurrentHashMap<String, PartitionScopedRegionLevelProgress.RegionLevelProgress>> entry = iterator.next();
 
                 String partitionKeyRangeId = entry.getKey();
-                String sessionTokenAsString = entry.getValue().get(PartitionScopedRegionLevelProgress.GLOBAL_PROGRESS_KEY).getSessionToken().convertToString();
-
-                result = result.append(partitionKeyRangeId).append(":").append(sessionTokenAsString);
-                if (iterator.hasNext()) {
-                    result = result.append(",");
-                }
-            }
-        }
 
-        return result.toString();
-    }
-
-    // validate whether the request can be scoped to a logical partition
-    // along with whether multi-write is enabled for the request / account
-    // multi-write setup needs to be verified since multiple regions can make progress
-    // independently as multiple regions can process writes
-    private boolean shouldUseBloomFilter(
-        RxDocumentServiceRequest request,
-        String partitionKeyRangeId,
-        Utils.ValueHolder<PartitionKeyInternal> partitionKeyInternal,
-        Utils.ValueHolder<PartitionKeyDefinition> partitionKeyDefinition,
-        PartitionScopedRegionLevelProgress partitionScopedRegionLevelProgress) {
+                Optional<ISessionToken> sessionToken = entry.getValue().get(PartitionScopedRegionLevelProgress.GLOBAL_PROGRESS_KEY).getSessionToken();
 
-        checkNotNull(request, "request cannot be null!");
-        checkNotNull(this.globalEndpointManager, "globalEndpointManager cannot be nulL!");
-        checkNotNull(partitionScopedRegionLevelProgress, "partitionScopedRegionLevelProgress cannot be null!");
+                if (sessionToken.isPresent()) {
 
-        partitionKeyInternal.v = request.getPartitionKeyInternal();
+                    ISessionToken sessionTokenInner = sessionToken.get();
+                    String sessionTokenAsString = sessionTokenInner.convertToString();
 
-        if (partitionKeyInternal.v == null) {
-            return false;
-        }
-
-        partitionKeyDefinition.v = request.getPartitionKeyDefinition();
-
-        if (partitionKeyDefinition.v == null) {
-            return false;
-        }
-
-        if (partitionScopedRegionLevelProgress.getHasPartitionSeenNonPointRequestsForDocuments(partitionKeyRangeId)) {
-            return false;
+                    result = result.append(partitionKeyRangeId).append(":").append(sessionTokenAsString);
+                    if (iterator.hasNext()) {
+                        result = result.append(",");
+                    }
+                } else {
+                    logger.warn("No session token present for partitionKeyRangeId {} in global progress", partitionKeyRangeId);
+                }
+            }
         }
 
-        return globalEndpointManager.canUseMultipleWriteLocations(request);
+        return result.toString();
     }
 
     private String extractFirstEffectivePreferredReadableRegion() {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/PartitionScopedRegionLevelProgress.java

@@ -1535,6 +1535,10 @@ private void addUserAgentSuffix(UserAgentContainer userAgentContainer, Set<UserA
             }
         }
 
+        if (!(this.sessionContainer instanceof RegionScopedSessionContainer)) {
+            userAgentFeatureFlags.remove(UserAgentFeatureFlags.RegionScopedSessionCapturing);
+        }
+
         userAgentContainer.setFeatureEnabledFlagsAsSuffix(userAgentFeatureFlags);
     }
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RegionScopedSessionContainer.java

@@ -173,50 +173,21 @@ static ISessionToken resolvePartitionLocalSessionToken(RxDocumentServiceRequest
                                                            Long collectionRid,
                                                            String partitionKeyRangeId,
                                                            String firstEffectivePreferredReadableRegion,
-                                                           boolean canUseBloomFilter) {
+                                                           GlobalEndpointManager globalEndpointManager) {
 
         if (partitionScopedRegionLevelProgress != null) {
 
-            Set<String> partitionKeyPossibleRegions = new HashSet<>();
-
             if (partitionScopedRegionLevelProgress.isPartitionKeyRangeIdPresent(partitionKeyRangeId)) {
-
-                if (canUseBloomFilter) {
-                    partitionKeyPossibleRegions = partitionKeyBasedBloomFilter
-                        .tryGetPossibleRegionsLogicalPartitionResolvedTo(
-                            request,
-                            collectionRid,
-                            partitionKey,
-                            partitionKeyDefinition);
-
-                    return partitionScopedRegionLevelProgress
-                        .tryResolveSessionToken(
-                            request,
-                            partitionKeyPossibleRegions,
-                            partitionKeyRangeId,
-                            firstEffectivePreferredReadableRegion,
-                            true);
-
-                }
-
                 return partitionScopedRegionLevelProgress
                     .tryResolveSessionToken(
                         request,
-                        partitionKeyPossibleRegions,
+                        collectionRid,
                         partitionKeyRangeId,
                         firstEffectivePreferredReadableRegion,
-                        false);
+                        partitionKeyBasedBloomFilter,
+                        globalEndpointManager);
 
             } else {
-                if (canUseBloomFilter) {
-                    partitionKeyPossibleRegions = partitionKeyBasedBloomFilter
-                        .tryGetPossibleRegionsLogicalPartitionResolvedTo(
-                            request,
-                            collectionRid,
-                            partitionKey,
-                            partitionKeyDefinition);
-                }
-
                 ISessionToken parentSessionToken = null;
 
                 Collection<String> parents = request.requestContext.resolvedPartitionKeyRange.getParents();
@@ -232,10 +203,11 @@ static ISessionToken resolvePartitionLocalSessionToken(RxDocumentServiceRequest
                             resolvedSessionTokenForParentPkRangeId = partitionScopedRegionLevelProgress
                                 .tryResolveSessionToken(
                                     request,
-                                    partitionKeyPossibleRegions,
+                                    collectionRid,
                                     parentPkRangeId,
                                     firstEffectivePreferredReadableRegion,
-                                    canUseBloomFilter);
+                                    partitionKeyBasedBloomFilter,
+                                    globalEndpointManager);
 
                             if (resolvedSessionTokenForParentPkRangeId != null) {
                                 parentSessionToken = parentSessionToken != null ?

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

@@ -19,7 +19,8 @@ public enum UserAgentFeatureFlags {
     PerPartitionCircuitBreaker(1 << 1),
     ThinClient(1 << 2),
     // BinaryEncoding(1 << 3),
-    Http2(1 << 4);
+    Http2(1 << 4),
+    RegionScopedSessionCapturing(1 << 5);
 
     private final int value;