fixChildLeaseCreationAfterSplit (#46075)

* fix child lease create with null continuation token after split

---------

Co-authored-by: annie-mac <[email protected]>

Author: xinlian12

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/changefeed/epkversion/IncrementalChangeFeedProcessorTest.java

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed an issue where child partition lease is getting created with null continuation token when change feed processor restart after split - See [PR 46075](https://github.com/Azure/azure-sdk-for-java/pull/46075)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/changefeed/pkversion/IncrementalChangeFeedProcessorTest.java

@@ -195,6 +195,10 @@ public Mono<Lease> createLeaseIfNotExist(FeedRangeEpkImpl feedRange, String cont
                 return Mono.error(ex);
             })
             .map(documentResourceResponse -> {
+                logger.info(
+                    "Successfully created lease document for {} with continuation token {}",
+                    leaseToken,
+                    continuationToken);
                 if (documentResourceResponse == null) {
                     return null;
                 }

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,12 +3,9 @@
 package com.azure.cosmos.implementation.changefeed.pkversion;
 
 import com.azure.cosmos.implementation.CosmosSchedulers;
-import com.azure.cosmos.implementation.Strings;
 import com.azure.cosmos.implementation.changefeed.Bootstrapper;
 import com.azure.cosmos.implementation.changefeed.LeaseStore;
 import com.azure.cosmos.implementation.changefeed.LeaseStoreManager;
-import com.azure.cosmos.implementation.changefeed.common.ChangeFeedMode;
-import com.azure.cosmos.implementation.changefeed.common.ChangeFeedState;
 import com.azure.cosmos.implementation.changefeed.common.LeaseVersion;
 import com.azure.cosmos.models.ChangeFeedProcessorOptions;
 import org.slf4j.Logger;

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreManagerImpl.java

@@ -191,6 +191,10 @@ public Mono<Lease> createLeaseIfNotExist(String leaseToken, String continuationT
                 return Mono.error(ex);
             })
             .map(documentResourceResponse -> {
+                logger.info(
+                    "Successfully created lease document for {} with continuation token {}.",
+                    leaseToken,
+                    continuationToken);
                 if (documentResourceResponse == null) {
                     return null;
                 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/BootstrapperImpl.java

@@ -12,7 +12,7 @@
  */
 public interface PartitionSynchronizer {
     /**
-     * Creates missing leases.
+     * Creates missing leases during startup.
      *
      * @return a deferred computation of this operation.
      */

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreManagerImpl.java

@@ -4,7 +4,6 @@
 
 import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.implementation.PartitionKeyRange;
-import com.azure.cosmos.implementation.Resource;
 import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient;
 import com.azure.cosmos.implementation.changefeed.Lease;
 import com.azure.cosmos.implementation.changefeed.LeaseContainer;
@@ -17,8 +16,11 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-import java.util.HashSet;
-import java.util.Set;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 import static com.azure.cosmos.BridgeInternal.extractContainerSelfLink;
 
@@ -55,15 +57,19 @@ public PartitionSynchronizerImpl(
 
     @Override
     public Mono<Void> createMissingLeases() {
+        Map<String, List<String>> leaseTokenMap = new ConcurrentHashMap<>();
+
         return this.enumPartitionKeyRanges()
-            .map(Resource::getId)
+            .map(partitionKeyRange -> {
+                leaseTokenMap.put(partitionKeyRange.getId(), partitionKeyRange.getParents());
+                return partitionKeyRange.getId();
+            })
             .collectList()
             .flatMap( partitionKeyRangeIds -> {
                 logger.info(
                     "Checking whether leases for any partition is missing - partitions - {}",
                     String.join(", ", partitionKeyRangeIds));
-                Set<String> leaseTokens = new HashSet<>(partitionKeyRangeIds);
-                return this.createLeases(leaseTokens).then();
+                return this.createLeases(leaseTokenMap).then();
             })
             .onErrorResume( throwable -> {
                 logger.error("Failed to create missing leases.", throwable);
@@ -143,41 +149,45 @@ private Flux<PartitionKeyRange> enumPartitionKeyRanges() {
     }
 
     /**
-     * Creates leases if they do not exist. This might happen on initial start or if some lease was unexpectedly lost.
+     * Creates leases if they do not exist for the partition or partition's parent partitions.
+     * This might happen on initial start or if some lease was unexpectedly lost.
      * <p>
      * Leases are created without the continuation token. It means partitions will be read according to
      *   'From Beginning' or 'From current time'.
      * Same applies also to split partitions. We do not search for parent lease and take continuation token since this
      *   might end up of reprocessing all the events since the split.
      *
-     * @param leaseTokens a hash set of all the lease tokens.
+     * @param leaseTokenMap a map of all the lease tokens and their mapping parent lease tokens.
      * @return a deferred computation of this call.
      */
-    private Flux<Lease> createLeases(Set<String> leaseTokens)
+    private Flux<Lease> createLeases(Map<String, List<String>> leaseTokenMap)
     {
-        Set<String> addedLeaseTokens = new HashSet<>(leaseTokens);
-
+        List<String> leaseTokensToBeAdded = new ArrayList<>();
         return this.leaseContainer.getAllLeases()
+            .map(lease -> lease.getLeaseToken())
+            .collectList()
+            .flatMapMany(existingLeaseTokens -> {
+                // only create lease documents if there is no existing lease document matching the partition or its parent partitions
+                leaseTokensToBeAdded.addAll(
+                    leaseTokenMap.entrySet().stream()
+                        .filter(entry -> !existingLeaseTokens.contains(entry.getKey()))
+                        .filter(entry -> entry.getValue() == null ||
+                            entry.getValue().isEmpty() ||
+                            entry.getValue().stream().noneMatch(existingLeaseTokens::contains))
+                        .map(Map.Entry::getKey)
+                        .collect(Collectors.toList())
+                );
+
+                logger.info("Missing lease documents for partitions: [{}]", String.join(", ", leaseTokensToBeAdded));
+                return Flux.fromIterable(leaseTokensToBeAdded);
+            })
+            .flatMap(leaseTokenToBeAdded -> {
+                logger.debug("Adding a new lease document for partition {}", leaseTokenToBeAdded);
+                return this.leaseManager.createLeaseIfNotExist(leaseTokenToBeAdded, null);
+            }, this.degreeOfParallelism)
             .map(lease -> {
-                if (lease != null) {
-                    logger.debug("Found an existing lease document for partition {}", lease.getLeaseToken());
-                    // Get leases after getting ranges, to make sure that no other hosts checked in continuation for
-                    //   split partition after we got leases.
-                    addedLeaseTokens.remove(lease.getLeaseToken());
-                }
-
+                logger.info("Added new lease document for partition {}", lease.getLeaseToken());
                 return lease;
-            })
-            .thenMany(Flux.fromIterable(addedLeaseTokens)
-                .flatMap( addedRangeId -> {
-                    logger.debug("Adding a new lease document for partition {}", addedRangeId);
-
-                    return this.leaseManager.createLeaseIfNotExist(addedRangeId, null);
-                }, this.degreeOfParallelism)
-                .map( lease -> {
-                    logger.info("Added new lease document for partition {}", lease.getLeaseToken());
-                    return lease;
-                })
-            );
+            });
     }
 }