diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ChangeFeedProcessorBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ChangeFeedProcessorBuilder.java index 00fdd3f2a7f7..1bbcc3f4f2b2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ChangeFeedProcessorBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ChangeFeedProcessorBuilder.java @@ -133,6 +133,33 @@ public ChangeFeedProcessorBuilder handleChanges(Consumer> consume return this; } + /** + * Sets a consumer function which will be called to process changes for LatestVersion change feed mode. + * + * + * 
+     * .handleLatestVersionChanges(docs -> {
+     *     for (JsonNode item : docs) {
+     *         // Implementation for handling and processing of each JsonNode item goes here
+     *     }
+     * })
+     *
+ * + * + * @param consumer the {@link Consumer} to call for handling the feeds. + * @return current Builder. + */ + @Beta(value = Beta.SinceVersion.V4_40_0, warningText = Beta.PREVIEW_SUBJECT_TO_CHANGE_WARNING) + public ChangeFeedProcessorBuilder handleLatestVersionChanges(Consumer> consumer) { + checkNotNull(consumer, "Argument 'consumer' can not be null"); + checkArgument(this.incrementalModeLeaseConsumer == null, "consumer has already been defined"); + + this.incrementalModeLeaseConsumer = consumer; + this.changeFeedMode = ChangeFeedMode.INCREMENTAL; + this.leaseVersion = LeaseVersion.EPK_RANGE_BASED_LEASE; + return this; + } + /** * Sets a consumer function which will be called to process changes for AllVersionsAndDeletes change feed mode. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java index 49a60b6378d1..fb546d25604b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java @@ -275,6 +275,9 @@ public static class HttpHeaders { // Client Encryption Headers public static final String IS_CLIENT_ENCRYPTED_HEADER = "x-ms-cosmos-is-client-encrypted"; public static final String INTENDED_COLLECTION_RID_HEADER = "x-ms-cosmos-intended-collection-rid"; + + // SDK supported capacities headers + public static final String SDK_SUPPORTED_CAPABILITIES = "x-ms-cosmos-sdk-supported-capabilities"; } public static class A_IMHeaderValues { @@ -282,6 +285,18 @@ public static class A_IMHeaderValues { public static final String FULL_FIDELITY_FEED = "Full-Fidelity Feed"; } + public static class SDKSupportedCapabilities { + private static final long None = 0; // 0 + private static final long PartitionMerge = 1; // 1 << 0 + + public static final long SUPPORTED_CAPABILITIES; + public static final long SUPPORTED_CAPABILITIES_EXCLUDE_MERGE; + static { + SUPPORTED_CAPABILITIES = PartitionMerge; + SUPPORTED_CAPABILITIES_EXCLUDE_MERGE = None; + } + } + public static class Versions { public static final String CURRENT_VERSION = "2020-07-15"; public static final String QUERY_VERSION = "1.0"; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 360f7c816890..e6909ffe7cf0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -1648,6 +1648,12 @@ private Mono populateHeadersAsync(RxDocumentServiceReq request.getHeaders().put(HttpConstants.HttpHeaders.API_TYPE, this.apiType.toString()); } + if (!request.getHeaders().containsKey(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES)) { + request.getHeaders().put( + HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES, + String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPORTED_CAPABILITIES)); + } + if ((RequestVerb.POST.equals(httpMethod) || RequestVerb.PUT.equals(httpMethod)) && !request.getHeaders().containsKey(HttpConstants.HttpHeaders.CONTENT_TYPE)) { request.getHeaders().put(HttpConstants.HttpHeaders.CONTENT_TYPE, RuntimeConstants.MediaTypes.JSON); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java index f3d0c3fbfc8e..25a0f44d037c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java @@ -78,6 +78,8 @@ public RxGatewayStoreModel( "no-cache"); this.defaultHeaders.put(HttpConstants.HttpHeaders.VERSION, HttpConstants.Versions.CURRENT_VERSION); + this.defaultHeaders.put(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES, + String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPORTED_CAPABILITIES)); if (apiType != null){ this.defaultHeaders.put(HttpConstants.HttpHeaders.API_TYPE, apiType.toString()); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/ChangeFeedContextClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/ChangeFeedContextClient.java index 9b7098b6da56..878714fdb806 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/ChangeFeedContextClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/ChangeFeedContextClient.java @@ -6,12 +6,14 @@ import com.azure.cosmos.CosmosAsyncDatabase; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.CosmosBulkOperationResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; import com.azure.cosmos.models.CosmosContainerResponse; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; import com.azure.cosmos.models.CosmosDatabaseResponse; +import com.azure.cosmos.models.CosmosItemOperation; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -172,4 +174,11 @@ Mono> readItem(String itemId, PartitionKey partitionKe * @return The list of partition key ranges. */ Mono> getOverlappingRanges(Range range); + + /** + * Executes flux of operations in Bulk. + * @param operations Flux of operation which will be executed by this container. + * @return A Flux of {@link CosmosBulkOperationResponse} which contains operation and it's response or exception. + */ + Flux> executeBulkOperations(Flux operations); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/LeaseStoreManager.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/LeaseStoreManager.java index 6c1edbf95386..8a8bc56354ee 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/LeaseStoreManager.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/LeaseStoreManager.java @@ -8,6 +8,7 @@ import reactor.core.publisher.Mono; import java.time.Duration; +import java.util.List; /** * Defines an interface for operations with {@link Lease}. @@ -68,6 +69,14 @@ interface LeaseStoreManagerBuilderDefinition { */ Mono delete(Lease lease); + /** + * DELETE all the leases passed in. + * + * @param leases the leases to be removed. + * @return a representation of the deferred computation of this call. + */ + Mono delete(List leases); + /** * Acquire ownership of the lease. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedContextClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedContextClientImpl.java index 92a97a12e551..41d0f4c8c09d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedContextClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedContextClientImpl.java @@ -11,12 +11,14 @@ import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.CosmosBulkOperationResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; import com.azure.cosmos.models.CosmosContainerResponse; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; import com.azure.cosmos.models.CosmosDatabaseResponse; +import com.azure.cosmos.models.CosmosItemOperation; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; @@ -108,6 +110,12 @@ public Mono> getOverlappingRanges(Range range) { .publishOn(this.scheduler); } + @Override + public Flux> executeBulkOperations(Flux operations) { + return this.cosmosContainer.executeBulkOperations(operations) + .publishOn(this.scheduler); + } + @Override public Flux> readPartitionKeyRangeFeed(String partitionKeyRangesOrCollectionLink, CosmosQueryRequestOptions cosmosQueryRequestOptions) { return this.documentClient.readPartitionKeyRanges(partitionKeyRangesOrCollectionLink, cosmosQueryRequestOptions) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/BootstrapperImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/BootstrapperImpl.java index 9e0530dad77d..93b7c83608bd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/BootstrapperImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/BootstrapperImpl.java @@ -5,6 +5,7 @@ import com.azure.cosmos.implementation.CosmosSchedulers; import com.azure.cosmos.implementation.changefeed.Bootstrapper; import com.azure.cosmos.implementation.changefeed.LeaseStore; +import com.azure.cosmos.implementation.changefeed.LeaseStoreManager; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -21,13 +22,20 @@ class BootstrapperImpl implements Bootstrapper { private final Logger logger = LoggerFactory.getLogger(BootstrapperImpl.class); private final PartitionSynchronizer synchronizer; private final LeaseStore leaseStore; + private final LeaseStoreManager pkVersionLeaseStoreManager; private final Duration lockTime; private final Duration sleepTime; private volatile boolean isInitialized; private volatile boolean isLockAcquired; - public BootstrapperImpl(PartitionSynchronizer synchronizer, LeaseStore leaseStore, Duration lockTime, Duration sleepTime) { + public BootstrapperImpl( + PartitionSynchronizer synchronizer, + LeaseStore leaseStore, + Duration lockTime, + Duration sleepTime, + LeaseStoreManager pkVersionLeaseStoreManager) { + checkNotNull(synchronizer, "Argument 'synchronizer' can not be null"); checkNotNull(leaseStore, "Argument 'leaseStore' can not be null"); checkArgument(lockTime != null && this.isPositive(lockTime), "lockTime should be non-null and positive"); @@ -35,6 +43,7 @@ public BootstrapperImpl(PartitionSynchronizer synchronizer, LeaseStore leaseStor this.synchronizer = synchronizer; this.leaseStore = leaseStore; + this.pkVersionLeaseStoreManager = pkVersionLeaseStoreManager; this.lockTime = lockTime; this.sleepTime = sleepTime; @@ -66,7 +75,18 @@ public Mono initialize() { logger.info("Another instance is initializing the store"); return Mono.just(isLockAcquired).delayElement(this.sleepTime, CosmosSchedulers.COSMOS_PARALLEL); } else { - return this.synchronizer.createMissingLeases() + return this.shouldBootstrapFromPkVersionLeases() + .flatMap(shouldBootstrapFromPkVersion -> { + if (shouldBootstrapFromPkVersion) { + return this.pkVersionLeaseStoreManager.getAllLeases() + .collectList() + .flatMap(pkVersionLeases -> { + return this.synchronizer.createMissingLeases(pkVersionLeases) + .then(this.pkVersionLeaseStoreManager.delete(pkVersionLeases)); // after bootstrapping, we are going to delete all th pk version leases + }); + } + return this.synchronizer.createMissingLeases(); + }) .then(this.leaseStore.markInitialized()); } }) @@ -85,4 +105,8 @@ public Mono initialize() { .repeat(() -> !this.isInitialized) .then(); } + + private Mono shouldBootstrapFromPkVersionLeases() { + return this.pkVersionLeaseStoreManager == null ? Mono.just(false) : this.pkVersionLeaseStoreManager.isInitialized(); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ChangeFeedProcessorImplBase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ChangeFeedProcessorImplBase.java index 50c23f2d2d73..c83452dde83e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ChangeFeedProcessorImplBase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ChangeFeedProcessorImplBase.java @@ -3,6 +3,7 @@ package com.azure.cosmos.implementation.changefeed.epkversion; +import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.ChangeFeedProcessor; import com.azure.cosmos.ConsistencyLevel; import com.azure.cosmos.CosmosAsyncContainer; @@ -67,9 +68,12 @@ public abstract class ChangeFeedProcessorImplBase implements ChangeFeedProces private final Scheduler scheduler; private volatile String databaseResourceId; + private volatile String databaseId; private volatile String collectionResourceId; + private volatile String collectionId; private PartitionLoadBalancingStrategy loadBalancingStrategy; private LeaseStoreManager leaseStoreManager; + private LeaseStoreManager pkVersionLeaseStoreManager; private HealthMonitor healthMonitor; private volatile PartitionManager partitionManager; @@ -257,12 +261,14 @@ private Mono initializeCollectionPropertiesForBuild() { .readDatabase(this.feedContextClient.getDatabaseClient(), null) .map(databaseResourceResponse -> { this.databaseResourceId = databaseResourceResponse.getProperties().getResourceId(); + this.databaseId = databaseResourceResponse.getProperties().getId(); return this.databaseResourceId; }) .flatMap( id -> this.feedContextClient .readContainer(this.feedContextClient.getContainerClient(), null) .map(documentCollectionResourceResponse -> { this.collectionResourceId = documentCollectionResourceResponse.getProperties().getResourceId(); + this.collectionId = documentCollectionResourceResponse.getProperties().getId(); return this; })); } @@ -287,6 +293,19 @@ private Mono getLeaseStoreManager() { .requestOptionsFactory(requestOptionsFactory) .hostName(this.hostName) .build(); + + if (this.canBootstrapFromPkVersionLeaseStore()) { + String pkVersionLeasePrefix = this.getPkVersionLeasePrefix(); + this.pkVersionLeaseStoreManager = + com.azure.cosmos.implementation.changefeed.pkversion.LeaseStoreManagerImpl.builder() + .leasePrefix(pkVersionLeasePrefix) + .leaseCollectionLink(this.leaseContextClient.getContainerClient()) + .leaseContextClient(this.leaseContextClient) + .requestOptionsFactory(requestOptionsFactory) + .hostName(this.hostName) + .build(); + } + return Mono.just(this.leaseStoreManager); }); } @@ -319,20 +338,45 @@ private String getLeasePrefix() { this.collectionResourceId); } + private String getPkVersionLeasePrefix() { + String optionsPrefix = this.changeFeedProcessorOptions.getLeasePrefix(); + + if (optionsPrefix == null) { + optionsPrefix = ""; + } + + URI uri = this.feedContextClient.getServiceEndpoint(); + + return String.format( + "%s%s_%s_%s", + optionsPrefix, + uri.getHost(), + this.databaseId, + this.collectionId); + } + abstract Class getPartitionProcessorItemType(); + abstract boolean canBootstrapFromPkVersionLeaseStore(); private Mono buildPartitionManager(LeaseStoreManager leaseStoreManager) { CheckpointerObserverFactory factory = new CheckpointerObserverFactory<>(this.observerFactory, new CheckpointFrequency()); PartitionSynchronizerImpl synchronizer = new PartitionSynchronizerImpl( this.feedContextClient, - this.feedContextClient.getContainerClient(), + BridgeInternal.extractContainerSelfLink(this.feedContextClient.getContainerClient()), leaseStoreManager, leaseStoreManager, DEFAULT_DEGREE_OF_PARALLELISM, - DEFAULT_QUERY_PARTITIONS_MAX_BATCH_SIZE); - - Bootstrapper bootstrapper = new BootstrapperImpl(synchronizer, leaseStoreManager, this.lockTime, this.sleepTime); + DEFAULT_QUERY_PARTITIONS_MAX_BATCH_SIZE, + this.changeFeedProcessorOptions, + this.changeFeedMode); + + Bootstrapper bootstrapper = new BootstrapperImpl( + synchronizer, + leaseStoreManager, + this.lockTime, + this.sleepTime, + this.pkVersionLeaseStoreManager); PartitionSupervisorFactory partitionSupervisorFactory = new PartitionSupervisorFactoryImpl<>( factory, leaseStoreManager, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/FullFidelityChangeFeedProcessorImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/FullFidelityChangeFeedProcessorImpl.java index c0d89aa7cbd2..6f1bbab2f9dc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/FullFidelityChangeFeedProcessorImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/FullFidelityChangeFeedProcessorImpl.java @@ -37,4 +37,10 @@ CosmosChangeFeedRequestOptions createRequestOptionsForProcessingFromNow(FeedRang Class getPartitionProcessorItemType() { return ChangeFeedProcessorItem.class; } + + @Override + boolean canBootstrapFromPkVersionLeaseStore() { + // pkVersion does not support fullFidelity mode + return false; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/IncrementalChangeFeedProcessorImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/IncrementalChangeFeedProcessorImpl.java index 496a779132cc..0933b5c21870 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/IncrementalChangeFeedProcessorImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/IncrementalChangeFeedProcessorImpl.java @@ -35,4 +35,9 @@ CosmosChangeFeedRequestOptions createRequestOptionsForProcessingFromNow(FeedRang Class getPartitionProcessorItemType() { return JsonNode.class; } + + @Override + boolean canBootstrapFromPkVersionLeaseStore() { + return true; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreManagerImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreManagerImpl.java index 3ba65cd2fec0..7f85730f574f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreManagerImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreManagerImpl.java @@ -33,6 +33,7 @@ import java.time.Duration; import java.util.Collections; +import java.util.List; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -205,6 +206,11 @@ public Mono delete(Lease lease) { .then(); } + @Override + public Mono delete(List leases) { + throw new UnsupportedOperationException("Delete with leases are not supported for Change Feed epk version wire format"); + } + @Override public Mono acquire(Lease lease) { if (lease == null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizer.java index 3a780c80ded1..5aaaf23f4e67 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizer.java @@ -8,6 +8,8 @@ import com.azure.cosmos.implementation.changefeed.epkversion.feedRangeGoneHandler.FeedRangeGoneHandler; import reactor.core.publisher.Mono; +import java.util.List; + /** * READ DocDB partitions and create leases if they do not exist. */ @@ -19,6 +21,14 @@ public interface PartitionSynchronizer { */ Mono createMissingLeases(); + /*** + * Create epk leases based on pk version leases. + * + * @param pkVersionLeases partition key range based leases. + * @return a deferred computation of this operation. + */ + Mono createMissingLeases(List pkVersionLeases); + /*** * Get the feedRangeGone handler based on whether it is merge or split. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImpl.java index 441420da9d1d..834c1c6816d4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImpl.java @@ -2,24 +2,30 @@ // Licensed under the MIT License. package com.azure.cosmos.implementation.changefeed.epkversion; -import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient; import com.azure.cosmos.implementation.changefeed.Lease; import com.azure.cosmos.implementation.changefeed.LeaseContainer; import com.azure.cosmos.implementation.changefeed.LeaseManager; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedMode; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedState; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedStateV1; import com.azure.cosmos.implementation.changefeed.epkversion.feedRangeGoneHandler.FeedRangeGoneHandler; import com.azure.cosmos.implementation.changefeed.epkversion.feedRangeGoneHandler.FeedRangeGoneMergeHandler; import com.azure.cosmos.implementation.changefeed.epkversion.feedRangeGoneHandler.FeedRangeGoneSplitHandler; +import com.azure.cosmos.implementation.feedranges.FeedRangeContinuation; import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.ChangeFeedProcessorOptions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -29,24 +35,31 @@ class PartitionSynchronizerImpl implements PartitionSynchronizer { private final Logger logger = LoggerFactory.getLogger(PartitionSynchronizerImpl.class); private final ChangeFeedContextClient documentClient; - private final CosmosAsyncContainer collectionSelfLink; + private final String collectionSelfLink; private final LeaseContainer leaseContainer; private final LeaseManager leaseManager; + + private final ChangeFeedProcessorOptions changeFeedProcessorOptions; + private final ChangeFeedMode changeFeedMode; private final int degreeOfParallelism; private final int maxBatchSize; public PartitionSynchronizerImpl( - ChangeFeedContextClient documentClient, - CosmosAsyncContainer collectionSelfLink, - LeaseContainer leaseContainer, - LeaseManager leaseManager, - int degreeOfParallelism, - int maxBatchSize) { + ChangeFeedContextClient documentClient, + String collectionSelfLink, + LeaseContainer leaseContainer, + LeaseManager leaseManager, + int degreeOfParallelism, + int maxBatchSize, + ChangeFeedProcessorOptions changeFeedProcessorOptions, + ChangeFeedMode changeFeedMode) { this.documentClient = documentClient; this.collectionSelfLink = collectionSelfLink; this.leaseContainer = leaseContainer; this.leaseManager = leaseManager; + this.changeFeedProcessorOptions = changeFeedProcessorOptions; + this.changeFeedMode = changeFeedMode; this.degreeOfParallelism = degreeOfParallelism; this.maxBatchSize = maxBatchSize; } @@ -61,6 +74,15 @@ public Mono createMissingLeases() { }); } + @Override + public Mono createMissingLeases(List pkVersionLeases) { + return this.documentClient.getOverlappingRanges(PartitionKeyInternalHelper.FullRange) + .flatMap(pkRangeList -> this.createLeases(pkRangeList, pkVersionLeases).then()) + .doOnError(throwable -> { + logger.error("Create missing leases from pkVersionLeases failed", throwable); + }); + } + @Override public Mono getFeedRangeGoneHandler(Lease lease) { checkNotNull(lease, "Argument 'lease' can not be null"); @@ -137,4 +159,49 @@ private Flux createLeases(List partitionKeyRanges) { }, this.degreeOfParallelism); }); } + + private Flux createLeases(List partitionKeyRanges, List pkVersionLeases) { + return Mono.just(pkVersionLeases) + .flatMapMany(leaseList -> { + Map leaseByPartition = + leaseList.stream().collect(Collectors.toMap(lease -> lease.getLeaseToken(), lease -> lease)); + + return Flux.fromIterable(partitionKeyRanges) + .flatMap(pkRange -> { + if (leaseByPartition.containsKey(pkRange.getId())) { + // find a matching pk version lease, create epk version lease based on it + FeedRangeEpkImpl feedRangeEpk = new FeedRangeEpkImpl(pkRange.toRange()); + String leaseContinuationToken = this.getLeaseContinuationToken( + feedRangeEpk, + leaseByPartition.get(pkRange.getId()).getContinuationToken()); + return leaseManager.createLeaseIfNotExist(feedRangeEpk, leaseContinuationToken); + } + + return Mono.empty(); + + }, this.degreeOfParallelism); + }); + } + + private String getLeaseContinuationToken( + FeedRangeEpkImpl feedRangeEpk, + String etag) { + FeedRangeContinuation feedRangeContinuation = FeedRangeContinuation.create( + this.collectionSelfLink, + feedRangeEpk, + feedRangeEpk.getRange()); + feedRangeContinuation.replaceContinuation(etag); + + ChangeFeedState changeFeedState = new ChangeFeedStateV1( + this.collectionSelfLink, + feedRangeEpk, + this.changeFeedMode, + PartitionProcessorHelper.getStartFromSettings( + feedRangeEpk, + this.changeFeedProcessorOptions, + this.changeFeedMode), + feedRangeContinuation); + + return changeFeedState.toString(); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ServiceItemLeaseV1.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ServiceItemLeaseV1.java index a0b125755513..96525257ff09 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ServiceItemLeaseV1.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/ServiceItemLeaseV1.java @@ -150,17 +150,11 @@ public ChangeFeedState getContinuationState(String containerRid, ChangeFeedMode // Lease token are stored in Base64 encoded json - and contains the complete ChangeFeedState ChangeFeedState changeFeedState = ChangeFeedStateV1.fromString(this.continuationToken); - // Calculating this token from epk based lease format - // This token is then used to pass as lsn in form of etag. - String token = changeFeedState.getContinuation().getCurrentContinuationToken().getToken(); - // This token has extra quotes - token = token.replace("\"", ""); - return new ChangeFeedStateV1( containerRid, this.feedRangeInternal, changeFeedMode, - ChangeFeedStartFromInternal.createFromETagAndFeedRange(token, this.feedRangeInternal), + changeFeedState.getStartFromSettings(), changeFeedState.getContinuation()); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/IncrementalChangeFeedProcessorImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/IncrementalChangeFeedProcessorImpl.java index 70ac1d9178a5..219d66979742 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/IncrementalChangeFeedProcessorImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/IncrementalChangeFeedProcessorImpl.java @@ -87,8 +87,8 @@ public class IncrementalChangeFeedProcessorImpl implements ChangeFeedProcessor, private final ChangeFeedContextClient feedContextClient; private final ChangeFeedProcessorOptions changeFeedProcessorOptions; private final ChangeFeedObserverFactory observerFactory; - private volatile String databaseResourceId; - private volatile String collectionResourceId; + private volatile String databaseId; + private volatile String collectionId; private final ChangeFeedContextClient leaseContextClient; private PartitionLoadBalancingStrategy loadBalancingStrategy; private LeaseStoreManager leaseStoreManager; @@ -214,7 +214,7 @@ public Mono> getEstimatedLag() { final FeedRangeInternal feedRange = new FeedRangePartitionKeyRangeImpl(lease.getLeaseToken()); final CosmosChangeFeedRequestOptions options = ModelBridgeInternal.createChangeFeedRequestOptionsForChangeFeedState( - lease.getContinuationState(this.collectionResourceId, ChangeFeedMode.INCREMENTAL)); + lease.getContinuationState(this.collectionId, ChangeFeedMode.INCREMENTAL)); options.setMaxItemCount(1); return this.feedContextClient.createDocumentChangeFeedQuery( @@ -293,7 +293,7 @@ public Mono> getCurrentState() { final FeedRangeInternal feedRange = new FeedRangePartitionKeyRangeImpl(lease.getLeaseToken()); final CosmosChangeFeedRequestOptions options = ModelBridgeInternal.createChangeFeedRequestOptionsForChangeFeedState( - lease.getContinuationState(this.collectionResourceId, ChangeFeedMode.INCREMENTAL)); + lease.getContinuationState(this.collectionId, ChangeFeedMode.INCREMENTAL)); options.setMaxItemCount(1); return this.feedContextClient.createDocumentChangeFeedQuery( @@ -354,13 +354,13 @@ private Mono initializeCollectionPropertiesForBuild() { return this.feedContextClient .readDatabase(this.feedContextClient.getDatabaseClient(), null) .map( databaseResourceResponse -> { - this.databaseResourceId = databaseResourceResponse.getProperties().getId(); - return this.databaseResourceId; + this.databaseId = databaseResourceResponse.getProperties().getId(); + return this.databaseId; }) .flatMap( id -> this.feedContextClient .readContainer(this.feedContextClient.getContainerClient(), null) .map(documentCollectionResourceResponse -> { - this.collectionResourceId = documentCollectionResourceResponse.getProperties().getId(); + this.collectionId = documentCollectionResourceResponse.getProperties().getId(); return this; })); } @@ -411,8 +411,8 @@ private String getLeasePrefix() { "%s%s_%s_%s", optionsPrefix, uri.getHost(), - this.databaseResourceId, - this.collectionResourceId); + this.databaseId, + this.collectionId); } private Mono buildPartitionManager(LeaseStoreManager leaseStoreManager) { @@ -425,7 +425,7 @@ private Mono buildPartitionManager(LeaseStoreManager leaseStor leaseStoreManager, DEFAULT_DEGREE_OF_PARALLELISM, DEFAULT_QUERY_PARTITIONS_MAX_BATCH_SIZE, - this.collectionResourceId + this.collectionId ); Bootstrapper bootstrapper = new BootstrapperImpl(synchronizer, leaseStoreManager, this.lockTime, this.sleepTime); @@ -437,7 +437,7 @@ private Mono buildPartitionManager(LeaseStoreManager leaseStor this.changeFeedProcessorOptions, leaseStoreManager, this.feedContextClient.getContainerClient(), - this.collectionResourceId), + this.collectionId), this.changeFeedProcessorOptions, this.scheduler ); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreManagerImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreManagerImpl.java index 4bb4f4a383ae..35f850655084 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreManagerImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreManagerImpl.java @@ -20,6 +20,8 @@ import com.azure.cosmos.implementation.changefeed.ServiceItemLeaseUpdater; import com.azure.cosmos.implementation.changefeed.exceptions.LeaseLostException; import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl; +import com.azure.cosmos.models.CosmosBulkOperations; +import com.azure.cosmos.models.CosmosItemOperation; import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.PartitionKey; @@ -31,7 +33,9 @@ import reactor.core.publisher.Mono; import java.time.Duration; +import java.util.ArrayList; import java.util.Collections; +import java.util.List; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -40,7 +44,7 @@ * Provides flexible way to buildAsyncClient lease manager constructor parameters. * For the actual creation of lease manager instance, delegates to lease manager factory. */ -class LeaseStoreManagerImpl implements LeaseStoreManager, LeaseStoreManager.LeaseStoreManagerBuilderDefinition { +public class LeaseStoreManagerImpl implements LeaseStoreManager, LeaseStoreManager.LeaseStoreManagerBuilderDefinition { private final String LEASE_STORE_MANAGER_LEASE_SUFFIX = ".."; private final Logger logger = LoggerFactory.getLogger(LeaseStoreManagerImpl.class); @@ -203,6 +207,29 @@ public Mono delete(Lease lease) { .then(); } + @Override + public Mono delete(List leases) { + checkNotNull(leases, "Argument 'leases' can not be null"); + + List operations = new ArrayList<>(); + for (Lease lease : leases) { + operations.add(CosmosBulkOperations.getDeleteItemOperation(lease.getId(), new PartitionKey(lease.getId()))); + } + + return this.leaseDocumentClient.executeBulkOperations(Flux.defer(() -> Flux.fromIterable(operations))) + .flatMap(itemResponse -> { + if (itemResponse.getResponse() != null && itemResponse.getResponse().isSuccessStatusCode()) { + operations.remove(itemResponse.getOperation()); + } else { + logger.debug("Failed to delete pk version lease {}", itemResponse.getOperation().getId(), itemResponse.getException()); + } + + return Mono.empty(); + }) + .repeat(() -> operations.size() != 0) + .then(); + } + @Override public Mono acquire(Lease lease) { if (lease == null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/PartitionProcessorImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/PartitionProcessorImpl.java index a0c7ad2e5a8f..ea5ceb690631 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/PartitionProcessorImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/PartitionProcessorImpl.java @@ -4,6 +4,8 @@ import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.CosmosSchedulers; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.changefeed.CancellationToken; import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient; import com.azure.cosmos.implementation.changefeed.ChangeFeedObserver; @@ -42,8 +44,6 @@ */ class PartitionProcessorImpl implements PartitionProcessor { private static final Logger logger = LoggerFactory.getLogger(PartitionProcessorImpl.class); - - private static final int DefaultMaxItemCount = 100; private final ProcessorSettings settings; private final PartitionCheckpointer checkpointer; private final ChangeFeedObserver observer; @@ -69,6 +69,13 @@ public PartitionProcessorImpl(ChangeFeedObserver observer, ChangeFeedState state = settings.getStartState(); this.options = ModelBridgeInternal.createChangeFeedRequestOptionsForChangeFeedState(state); this.options.setMaxItemCount(settings.getMaxItemCount()); + + // For pk version, merge is not support, exclude it from the capabilities header + ImplementationBridgeHelpers.CosmosChangeFeedRequestOptionsHelper.getCosmosChangeFeedRequestOptionsAccessor() + .setHeader( + this.options, + HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES, + String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPORTED_CAPABILITIES_EXCLUDE_MERGE)); } @Override diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java index 033c8f52a5b7..7f42f9781fe3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java @@ -150,6 +150,8 @@ public GatewayAddressCache( // Set requested API version header for version enforcement. defaultRequestHeaders.put(HttpConstants.HttpHeaders.VERSION, HttpConstants.Versions.CURRENT_VERSION); + this.defaultRequestHeaders.put(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES, + String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPORTED_CAPABILITIES)); this.lastForcedRefreshMap = new ConcurrentHashMap<>(); this.globalEndpointManager = globalEndpointManager; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java index 991ed872855e..853d83c7b402 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java @@ -91,6 +91,9 @@ public HttpTransportClient(Configs configs, ConnectionPolicy connectionPolicy, U // Set requested API version header for version enforcement. this.defaultHeaders.put(HttpConstants.HttpHeaders.VERSION, HttpConstants.Versions.CURRENT_VERSION); this.defaultHeaders.put(HttpConstants.HttpHeaders.CACHE_CONTROL, HttpConstants.HeaderValues.NO_CACHE); + this.defaultHeaders.put( + HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES, + String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPORTED_CAPABILITIES)); if (userAgent == null) { userAgent = new UserAgentContainer(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java index 2cccad81c5a5..ff7cd8f33817 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java @@ -589,8 +589,8 @@ public enum RntbdRequestHeader implements RntbdHeader { PopulateIndexMetrics((short) 0x00A9, RntbdTokenType.Byte, false), IsClientEncrypted((short) 0x0087, RntbdTokenType.Byte, false), IntendedCollectionRid((short) 0x009D, RntbdTokenType.String, false), - CorrelatedActivityId((short) 0x00B0, RntbdTokenType.Guid, false); - + CorrelatedActivityId((short) 0x00B0, RntbdTokenType.Guid, false), + SDKSupportedCapabilities((short) 0x00A2, RntbdTokenType.ULong, false); public static final ImmutableMap map; public static final ImmutableSet set = Sets.immutableEnumSet(EnumSet.allOf(RntbdRequestHeader.class)); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java index 62ed25a15ca1..fd13e97a3906 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java @@ -119,6 +119,7 @@ final class RntbdRequestHeaders extends RntbdTokenStream { this.addIsClientEncrypted(headers); this.addIntendedCollectionRid(headers); this.addCorrelatedActivityId(headers); + this.addSDKSupportedCapabilities(headers); // Normal headers (Strings, Ints, Longs, etc.) @@ -161,6 +162,7 @@ final class RntbdRequestHeaders extends RntbdTokenStream { this.fillTokenFromHeader(headers, this::shouldBatchContinueOnError, HttpHeaders.SHOULD_BATCH_CONTINUE_ON_ERROR); this.fillTokenFromHeader(headers, this::isBatchOrdered, HttpHeaders.IS_BATCH_ORDERED); this.fillTokenFromHeader(headers, this::getCorrelatedActivityId, HttpHeaders.CORRELATED_ACTIVITY_ID); + this.fillTokenFromHeader(headers, this::getSDKSupportedCapabilities, HttpHeaders.SDK_SUPPORTED_CAPABILITIES); // Will be null in case of direct, which is fine - BE will use the value slice the connection context this. // When this is used in Gateway, the header value will be populated with the proxied HTTP request's header, @@ -603,6 +605,10 @@ private RntbdToken isBatchOrdered() { return this.get(RntbdRequestHeader.IsBatchOrdered); } + private RntbdToken getSDKSupportedCapabilities() { + return this.get(RntbdRequestHeader.SDKSupportedCapabilities); + } + private void addAimHeader(final Map headers) { final String value = headers.get(HttpHeaders.A_IM); @@ -1255,6 +1261,13 @@ private void addReturnPreference(final Map headers) { } } + private void addSDKSupportedCapabilities(final Map headers) { + final String value = headers.get(HttpHeaders.SDK_SUPPORTED_CAPABILITIES); + if (StringUtils.isNotEmpty(value)) { + this.getSDKSupportedCapabilities().setValue(Long.valueOf(value)); + } + } + private void fillTokenFromHeader(final Map headers, final Supplier supplier, final String name) { final String value = headers.get(name); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java index 9540b655c473..61c1fc5eff01 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java @@ -95,6 +95,8 @@ public enum SinceVersion { /** v4.35.0 */ V4_35_0, /** v4.37.0 */ - V4_37_0 + V4_37_0, + /** v4.40.0 */ + V4_40_0 } } diff --git a/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ChangeFeedProcessorLatestVersionChangesCodeSnippet.java b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ChangeFeedProcessorLatestVersionChangesCodeSnippet.java new file mode 100644 index 000000000000..c6cc944e1df8 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/samples/java/com/azure/cosmos/ChangeFeedProcessorLatestVersionChangesCodeSnippet.java @@ -0,0 +1,59 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.TestConfigurations; +import com.fasterxml.jackson.databind.JsonNode; + +public class ChangeFeedProcessorLatestVersionChangesCodeSnippet { + public void changeFeedProcessorBuilderCodeSnippet() { + String hostName = "test-host-name"; + CosmosAsyncClient cosmosAsyncClient = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .contentResponseOnWriteEnabled(true) + .consistencyLevel(ConsistencyLevel.SESSION) + .buildAsyncClient(); + CosmosAsyncDatabase cosmosAsyncDatabase = cosmosAsyncClient.getDatabase("testDb"); + CosmosAsyncContainer feedContainer = cosmosAsyncDatabase.getContainer("feedContainer"); + CosmosAsyncContainer leaseContainer = cosmosAsyncDatabase.getContainer("leaseContainer"); + // BEGIN: com.azure.cosmos.latestVersionChanges.builder + ChangeFeedProcessor changeFeedProcessor = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .handleLatestVersionChanges(docs -> { + for (JsonNode item : docs) { + // Implementation for handling and processing of each JsonNode item goes here + } + }) + .buildChangeFeedProcessor(); + // END: com.azure.cosmos.latestVersionChanges.builder + } + + public void handleChangesCodeSnippet() { + String hostName = "test-host-name"; + CosmosAsyncClient cosmosAsyncClient = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .contentResponseOnWriteEnabled(true) + .consistencyLevel(ConsistencyLevel.SESSION) + .buildAsyncClient(); + CosmosAsyncDatabase cosmosAsyncDatabase = cosmosAsyncClient.getDatabase("testDb"); + CosmosAsyncContainer feedContainer = cosmosAsyncDatabase.getContainer("feedContainer"); + CosmosAsyncContainer leaseContainer = cosmosAsyncDatabase.getContainer("leaseContainer"); + ChangeFeedProcessor changeFeedProcessor = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + // BEGIN: com.azure.cosmos.latestVersionChanges.handleChanges + .handleLatestVersionChanges(docs -> { + for (JsonNode item : docs) { + // Implementation for handling and processing of each JsonNode item goes here + } + }) + // END: com.azure.cosmos.latestVersionChanges.handleChanges + .buildChangeFeedProcessor(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/BootstrapperImplTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/BootstrapperImplTests.java new file mode 100644 index 000000000000..3509f5528423 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/BootstrapperImplTests.java @@ -0,0 +1,84 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.changefeed.epkversion; + +import com.azure.cosmos.implementation.changefeed.LeaseStore; +import com.azure.cosmos.implementation.changefeed.LeaseStoreManager; +import org.mockito.Mockito; +import org.testng.annotations.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.time.Duration; + +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.Mockito.times; + +public class BootstrapperImplTests { + + @Test(groups = "unit") + public void initializeStoreFromPkVersionLeaseStore() { + Duration lockTime = Duration.ofSeconds(5); + Duration expireTIme = Duration.ofSeconds(5); + + PartitionSynchronizer partitionSynchronizerMock = Mockito.mock(PartitionSynchronizer.class); + Mockito.when(partitionSynchronizerMock.createMissingLeases(Mockito.any())).thenReturn(Mono.empty()); + + LeaseStore leaseStoreMock = Mockito.mock(LeaseStore.class); + Mockito + .when(leaseStoreMock.isInitialized()) + .thenReturn(Mono.just(false)) + .thenReturn(Mono.just(true)); + Mockito.when(leaseStoreMock.acquireInitializationLock(lockTime)).thenReturn(Mono.just(true)); + Mockito.when(leaseStoreMock.markInitialized()).thenReturn(Mono.empty()); + + LeaseStoreManager pkVersionLeaseStoreManagerMock = Mockito.mock(LeaseStoreManager.class); + Mockito.when(pkVersionLeaseStoreManagerMock.isInitialized()).thenReturn(Mono.just(true)); + Mockito.when(pkVersionLeaseStoreManagerMock.getAllLeases()).thenReturn(Flux.empty()); + Mockito.when(pkVersionLeaseStoreManagerMock.delete(Mockito.anyList())).thenReturn(Mono.empty()); + BootstrapperImpl bootstrapper = new BootstrapperImpl( + partitionSynchronizerMock, + leaseStoreMock, + lockTime, + expireTIme, + pkVersionLeaseStoreManagerMock); + + bootstrapper.initialize().block(); + + Mockito.verify(pkVersionLeaseStoreManagerMock, times(1)).isInitialized(); + Mockito.verify(pkVersionLeaseStoreManagerMock, times(1)).getAllLeases(); + Mockito.verify(partitionSynchronizerMock, times(1)).createMissingLeases(anyList()); + Mockito.verify(pkVersionLeaseStoreManagerMock, Mockito.times(1)).delete(anyList()); + Mockito.verify(leaseStoreMock, times(2)).isInitialized(); + } + + @Test(groups = "unit") + public void initializeStoreFromScratch() { + Duration lockTime = Duration.ofSeconds(5); + Duration expireTIme = Duration.ofSeconds(5); + + PartitionSynchronizer partitionSynchronizerMock = Mockito.mock(PartitionSynchronizer.class); + Mockito.when(partitionSynchronizerMock.createMissingLeases()).thenReturn(Mono.empty()); + + LeaseStore leaseStoreMock = Mockito.mock(LeaseStore.class); + Mockito + .when(leaseStoreMock.isInitialized()) + .thenReturn(Mono.just(false)) + .thenReturn(Mono.just(true)); + Mockito.when(leaseStoreMock.acquireInitializationLock(lockTime)).thenReturn(Mono.just(true)); + Mockito.when(leaseStoreMock.markInitialized()).thenReturn(Mono.empty()); + + BootstrapperImpl bootstrapper = new BootstrapperImpl( + partitionSynchronizerMock, + leaseStoreMock, + lockTime, + expireTIme, + null); + + bootstrapper.initialize().block(); + Mockito.verify(partitionSynchronizerMock, times(0)).createMissingLeases(anyList()); + Mockito.verify(partitionSynchronizerMock, times(1)).createMissingLeases(); + Mockito.verify(leaseStoreMock, times(2)).isInitialized(); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImplTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImplTests.java index e04a5d92741f..ea1147797cd7 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImplTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/PartitionSynchronizerImplTests.java @@ -3,14 +3,18 @@ package com.azure.cosmos.implementation.changefeed.epkversion; -import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient; +import com.azure.cosmos.implementation.changefeed.Lease; import com.azure.cosmos.implementation.changefeed.LeaseContainer; import com.azure.cosmos.implementation.changefeed.LeaseManager; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedMode; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedState; +import com.azure.cosmos.implementation.changefeed.pkversion.ServiceItemLease; import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.ChangeFeedProcessorOptions; import org.assertj.core.api.Assertions; import org.mockito.ArgumentCaptor; import org.mockito.Mockito; @@ -31,20 +35,23 @@ import static org.mockito.Mockito.when; public class PartitionSynchronizerImplTests { - @Test + @Test(groups = "unit") public void createAllLeases() { ChangeFeedContextClient feedContextClientMock = Mockito.mock(ChangeFeedContextClient.class); - CosmosAsyncContainer feedContainerMock = Mockito.mock(CosmosAsyncContainer.class); + String feedContainerSelfLink = "/db/TestDB/coll/TestContainer"; LeaseContainer leaseContainerMock = Mockito.mock(LeaseContainer.class); LeaseManager leaseManagerMock = Mockito.mock(LeaseManager.class); + ChangeFeedProcessorOptions changeFeedProcessorOptions = Mockito.mock(ChangeFeedProcessorOptions.class); PartitionSynchronizerImpl partitionSynchronizer = new PartitionSynchronizerImpl( feedContextClientMock, - feedContainerMock, + feedContainerSelfLink, leaseContainerMock, leaseManagerMock, 1, - 1); + 1, + changeFeedProcessorOptions, + ChangeFeedMode.INCREMENTAL); List overlappingRanges = new ArrayList<>(); overlappingRanges.add(new PartitionKeyRange("1", "AA", "BB")); @@ -80,20 +87,23 @@ public void createAllLeases() { Assertions.assertThat(capturedEpkArguments.get(1).getRange()).isEqualTo(overlappingRanges.get(1).toRange()); } - @Test + @Test(groups = "unit") public void createMissingLeases() { ChangeFeedContextClient feedContextClientMock = Mockito.mock(ChangeFeedContextClient.class); - CosmosAsyncContainer feedContainerMock = Mockito.mock(CosmosAsyncContainer.class); + String feedContainerSelfLink = "/db/TestDB/coll/TestContainer"; LeaseContainer leaseContainerMock = Mockito.mock(LeaseContainer.class); LeaseManager leaseManagerMock = Mockito.mock(LeaseManager.class); + ChangeFeedProcessorOptions changeFeedProcessorOptions = Mockito.mock(ChangeFeedProcessorOptions.class); PartitionSynchronizerImpl partitionSynchronizer = new PartitionSynchronizerImpl( feedContextClientMock, - feedContainerMock, + feedContainerSelfLink, leaseContainerMock, leaseManagerMock, 1, - 1); + 1, + changeFeedProcessorOptions, + ChangeFeedMode.INCREMENTAL); List overlappingRanges = new ArrayList<>(); overlappingRanges.add(new PartitionKeyRange("1", "AA", "BB")); @@ -110,13 +120,13 @@ public void createMissingLeases() { new ServiceItemLeaseV1() .withLeaseToken("BB-DD") .withFeedRange(new FeedRangeEpkImpl(new Range<>("BB", "DD", true, false))); - childLease1.setId("TestLease-" + UUID.randomUUID()); + childLease2.setId("TestLease-" + UUID.randomUUID()); ServiceItemLeaseV1 childLease3 = new ServiceItemLeaseV1() .withLeaseToken("DD-EE") .withFeedRange(new FeedRangeEpkImpl(new Range<>("DD", "EE", true, false))); - childLease1.setId("TestLease-" + UUID.randomUUID()); + childLease3.setId("TestLease-" + UUID.randomUUID()); when(feedContextClientMock.getOverlappingRanges(PartitionKeyInternalHelper.FullRange)) .thenReturn(Mono.just(overlappingRanges)); @@ -136,4 +146,64 @@ public void createMissingLeases() { assertThat(feedRangeEpkArgumentCaptor.getAllValues().size()).isEqualTo(1); assertThat(feedRangeEpkArgumentCaptor.getAllValues().get(0).getRange()).isEqualTo(overlappingRanges.get(2).toRange()); } + + @Test(groups = "unit") + public void createMissingLeasesFromPkVersionLeases() { + ChangeFeedContextClient feedContextClientMock = Mockito.mock(ChangeFeedContextClient.class); + String feedContainerSelfLink = "/db/TestDB/coll/TestContainer"; + LeaseContainer leaseContainerMock = Mockito.mock(LeaseContainer.class); + LeaseManager leaseManagerMock = Mockito.mock(LeaseManager.class); + ChangeFeedProcessorOptions changeFeedProcessorOptions = Mockito.mock(ChangeFeedProcessorOptions.class); + + PartitionSynchronizerImpl partitionSynchronizer = new PartitionSynchronizerImpl( + feedContextClientMock, + feedContainerSelfLink, + leaseContainerMock, + leaseManagerMock, + 1, + 1, + changeFeedProcessorOptions, + ChangeFeedMode.INCREMENTAL); + + List overlappingRanges = new ArrayList<>(); + overlappingRanges.add(new PartitionKeyRange("1", "AA", "BB")); + overlappingRanges.add(new PartitionKeyRange("2", "BB", "DD")); + overlappingRanges.add(new PartitionKeyRange("3", "DD", "EE")); + + List pkVersionLeases = new ArrayList<>(); + + for (PartitionKeyRange partitionKeyRange : overlappingRanges) { + ServiceItemLease pkVersionLease = + new ServiceItemLease() + .withLeaseToken(partitionKeyRange.getId()) + .withETag(String.valueOf(partitionKeyRange.getId())); + pkVersionLease.setId("TestLease-" + UUID.randomUUID()); + pkVersionLeases.add(pkVersionLease); + } + + when(feedContextClientMock.getOverlappingRanges(PartitionKeyInternalHelper.FullRange)) + .thenReturn(Mono.just(overlappingRanges)); + + when(leaseManagerMock.createLeaseIfNotExist((FeedRangeEpkImpl) any(), any())) + .thenReturn(Mono.empty()); + + StepVerifier.create(partitionSynchronizer.createMissingLeases(pkVersionLeases)) + .verifyComplete(); + + // Verify the new lease will start from the lsn from pk version lease + ArgumentCaptor feedRangeEpkArgumentCaptor = ArgumentCaptor.forClass(FeedRangeEpkImpl.class); + ArgumentCaptor continuationTokenArgumentCaptor = ArgumentCaptor.forClass(String.class); + verify(leaseManagerMock, times(3)) + .createLeaseIfNotExist(feedRangeEpkArgumentCaptor.capture(), continuationTokenArgumentCaptor.capture()); + assertThat(feedRangeEpkArgumentCaptor.getAllValues().size()).isEqualTo(3); + assertThat(continuationTokenArgumentCaptor.getAllValues().size()).isEqualTo(3); + + for (int i = 0; i < pkVersionLeases.size(); i++) { + assertThat(feedRangeEpkArgumentCaptor.getAllValues().get(i).getRange()).isEqualTo(overlappingRanges.get(i).toRange()); + ChangeFeedState changeFeedState = ChangeFeedState.fromString(continuationTokenArgumentCaptor.getAllValues().get(i)); + assertThat(changeFeedState.getFeedRange()).isInstanceOf(FeedRangeEpkImpl.class); + assertThat(((FeedRangeEpkImpl)changeFeedState.getFeedRange()).getRange()).isEqualTo(overlappingRanges.get(i).toRange()); + assertThat(changeFeedState.getContinuation().getCurrentContinuationToken().getToken()).isEqualTo(pkVersionLeases.get(i).getContinuationToken()); + } + } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClientTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClientTest.java index 8cbb59d3bc7b..5497ae7b07f4 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClientTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClientTest.java @@ -151,7 +151,8 @@ public void validateDefaultHeaders() { assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.CACHE_CONTROL)).isEqualTo("no-cache"); assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.ACCEPT)).isEqualTo("application/json"); assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.VERSION)).isEqualTo(HttpConstants.Versions.CURRENT_VERSION); - + assertThat(httpRequest.headers().value(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES)) + .isEqualTo(String.valueOf(HttpConstants.SDKSupportedCapabilities.SUPPORTED_CAPABILITIES)); } @DataProvider(name = "fromMockedHttpResponseToExpectedDocumentClientException") diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index a2deb0652187..e5528246ce09 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -487,6 +487,7 @@ public Flux> bulkInsert(CosmosAsyncContainer cosmosCon return Flux.merge(Flux.fromIterable(result), concurrencyLevel); } + public List bulkInsertBlocking(CosmosAsyncContainer cosmosContainer, List documentDefinitionList) { return bulkInsert(cosmosContainer, documentDefinitionList, DEFAULT_BULK_INSERT_CONCURRENCY_LEVEL) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/changefeed/epkversion/ChangeFeedProcessorMigrateTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/changefeed/epkversion/ChangeFeedProcessorMigrateTests.java new file mode 100644 index 000000000000..960f425f3c66 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/changefeed/epkversion/ChangeFeedProcessorMigrateTests.java @@ -0,0 +1,250 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.rx.changefeed.epkversion; + +import com.azure.cosmos.ChangeFeedProcessor; +import com.azure.cosmos.ChangeFeedProcessorBuilder; +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosAsyncDatabase; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.InternalObjectNode; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.changefeed.common.LeaseVersion; +import com.azure.cosmos.models.ChangeFeedProcessorOptions; +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosContainerRequestOptions; +import com.azure.cosmos.rx.TestSuiteBase; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.commons.lang3.RandomStringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; +import reactor.core.scheduler.Schedulers; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Consumer; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * When customer has chosen to use handleLatestVersionChanges from handleChanges, + * internally SDK is going to change lease format from pkVersion to epkVersion. + * This test is to validate the migration works. + */ +public class ChangeFeedProcessorMigrateTests extends TestSuiteBase { + private final static Logger logger = LoggerFactory.getLogger(IncrementalChangeFeedProcessorTest.class); + private static final ObjectMapper OBJECT_MAPPER = Utils.getSimpleObjectMapper(); + private CosmosAsyncDatabase createdDatabase; + private final String hostName = RandomStringUtils.randomAlphabetic(6); + private final int FEED_COUNT = 10; + private final int CHANGE_FEED_PROCESSOR_TIMEOUT = 5000; + private final int FEED_COLLECTION_THROUGHPUT = 10100; + private final int LEASE_COLLECTION_THROUGHPUT = 400; + + private CosmosAsyncClient client; + + @Factory(dataProvider = "clientBuilders") + public ChangeFeedProcessorMigrateTests(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = { "emulator", "simple" }, timeOut = SETUP_TIMEOUT, alwaysRun = true) + public void before_ChangeFeedProcessorMigrateTests() { + client = getClientBuilder().buildAsyncClient(); + createdDatabase = getSharedCosmosDatabase(client); + } + + @AfterClass(groups = { "emulator", "simple" }, timeOut = 2 * SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + safeClose(client); + } + + @Test(groups = { "emulator" }, timeOut = 2 * TIMEOUT) + public void readFeedDocumentsBootstrapFromPkVersion() throws InterruptedException { + CosmosAsyncContainer createdFeedCollection = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + List createdDocuments = new ArrayList<>(); + Map receivedDocuments = new ConcurrentHashMap<>(); + setupReadFeedDocuments(createdDocuments, createdFeedCollection, FEED_COUNT); + + ChangeFeedProcessor changeFeedProcessor = + this.createDefaultChangeFeedProcessorBuilder(hostName, createdFeedCollection, createdLeaseCollection) + .handleChanges(changeFeedProcessorHandler(receivedDocuments)) + .buildChangeFeedProcessor(); + + try { + changeFeedProcessor.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .subscribe(); + } catch (Exception ex) { + logger.error("Change feed processor did not start in the expected time", ex); + throw ex; + } + + // Wait for the feed processor to receive and process the documents. + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + assertThat(changeFeedProcessor.isStarted()).as("Change Feed Processor instance is running").isTrue(); + + changeFeedProcessor + .stop() + .subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(CHANGE_FEED_PROCESSOR_TIMEOUT)) + .subscribe(); + + for (InternalObjectNode item : createdDocuments) { + assertThat(receivedDocuments.containsKey(item.getId())).as("Document with getId: " + item.getId()).isTrue(); + } + + String leaseItemCountQuery = "select VALUE COUNT(1) from c"; + int leaseContainerItemCount = createdLeaseCollection.queryItems(leaseItemCountQuery, Integer.class) + .collectList().block().get(0); + + // now switch to use handleLatestVersionChanges, validate it resumes from checkpoint from handChanges + logger.info("Switch to use handleLatestVersionChanges"); + + createdDocuments.clear(); + Map newReceivedDocuments = new ConcurrentHashMap<>(); + setupReadFeedDocuments(createdDocuments, createdFeedCollection, FEED_COUNT); + ChangeFeedProcessor changeFeedProcessor2 = + this.createDefaultChangeFeedProcessorBuilder(hostName, createdFeedCollection, createdLeaseCollection) + .handleLatestVersionChanges(changeFeedProcessorHandler(newReceivedDocuments)) + .buildChangeFeedProcessor(); + try { + changeFeedProcessor2.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(5 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .subscribe(); + } catch (Exception ex) { + logger.error("Change feed processor did not start in the expected time", ex); + throw ex; + } + + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + assertThat(changeFeedProcessor2.isStarted()).as("Change Feed Processor instance is running").isTrue(); + changeFeedProcessor2 + .stop() + .subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(CHANGE_FEED_PROCESSOR_TIMEOUT)) + .subscribe(); + + assertThat(newReceivedDocuments.size()).isEqualTo(createdDocuments.size()); + for (InternalObjectNode item : createdDocuments) { + assertThat(newReceivedDocuments.containsKey(item.getId())).as("Document with getId: " + item.getId()).isTrue(); + } + + // Validate lease container item + String query = "select * from c"; + List allItems = createdLeaseCollection.queryItems(query, JsonNode.class) + .collectList().block(); + assertThat(allItems.size()).isEqualTo(leaseContainerItemCount + 1); + int infoFileCount = 0; + for (JsonNode item : allItems) { + if (item.get("id").asText().endsWith(".info")) { + infoFileCount++; + } else { + assertThat(item.get("version").asInt()).isEqualTo(LeaseVersion.EPK_RANGE_BASED_LEASE.getVersionId()); + } + } + + assertThat(infoFileCount).isEqualTo(2); + + // Wait for the feed processor to shut down. + Thread.sleep(CHANGE_FEED_PROCESSOR_TIMEOUT); + } finally { + safeDeleteCollection(createdFeedCollection); + safeDeleteCollection(createdLeaseCollection); + } + } + + private ChangeFeedProcessorBuilder createDefaultChangeFeedProcessorBuilder( + String hostName, + CosmosAsyncContainer feedContainer, + CosmosAsyncContainer leaseContainer) { + return new ChangeFeedProcessorBuilder() + .hostName(hostName) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .options(new ChangeFeedProcessorOptions() + .setLeaseRenewInterval(Duration.ofSeconds(20)) + .setLeaseAcquireInterval(Duration.ofSeconds(10)) + .setLeaseExpirationInterval(Duration.ofSeconds(30)) + .setFeedPollDelay(Duration.ofSeconds(2)) + .setLeasePrefix("TEST") + .setMaxItemCount(10) + .setStartFromBeginning(true) + .setMaxScaleCount(0) // unlimited + ); + } + + private CosmosAsyncContainer createFeedCollection(int provisionedThroughput) { + CosmosContainerRequestOptions optionsFeedCollection = new CosmosContainerRequestOptions(); + return createCollection(createdDatabase, getCollectionDefinition(), optionsFeedCollection, provisionedThroughput); + } + + private CosmosAsyncContainer createLeaseCollection(int provisionedThroughput) { + CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); + CosmosContainerProperties collectionDefinition = new CosmosContainerProperties( + "leases_" + UUID.randomUUID(), + "/id"); + return createCollection(createdDatabase, collectionDefinition, options, provisionedThroughput); + } + + private void setupReadFeedDocuments( + List createdDocuments, + CosmosAsyncContainer feedCollection, + long count) { + List docDefList = new ArrayList<>(); + + for(int i = 0; i < count; i++) { + docDefList.add(getDocumentDefinition()); + } + + createdDocuments.addAll(bulkInsertBlocking(feedCollection, docDefList)); + waitIfNeededForReplicasToCatchUp(getClientBuilder()); + } + + private InternalObjectNode getDocumentDefinition() { + String uuid = UUID.randomUUID().toString(); + InternalObjectNode doc = new InternalObjectNode(String.format("{ " + + "\"id\": \"%s\", " + + "\"mypk\": \"%s\", " + + "\"sgmts\": [[6519456, 1471916863], [2498434, 1455671440]]" + + "}" + , uuid, uuid)); + return doc; + } + + + private Consumer> changeFeedProcessorHandler(Map receivedDocuments) { + return docs -> { + logger.info("START processing from thread in test {}", Thread.currentThread().getId()); + for (JsonNode item : docs) { + processItem(item, receivedDocuments); + } + logger.info("END processing from thread {}", Thread.currentThread().getId()); + }; + } + + private static synchronized void processItem(JsonNode item, Map receivedDocuments) { + try { + logger.info("RECEIVED {}", OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(item)); + } catch (JsonProcessingException e) { + logger.error("Failure in processing json [{}]", e.getMessage(), e); + } + receivedDocuments.put(item.get("id").asText(), item); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/changefeed/epkversion/IncrementalChangeFeedProcessorTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/changefeed/epkversion/IncrementalChangeFeedProcessorTest.java new file mode 100644 index 000000000000..7f96dfe2e9a3 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/changefeed/epkversion/IncrementalChangeFeedProcessorTest.java @@ -0,0 +1,1109 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.rx.changefeed.epkversion; + +import com.azure.cosmos.ChangeFeedProcessor; +import com.azure.cosmos.ChangeFeedProcessorBuilder; +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosAsyncDatabase; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.InternalObjectNode; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedState; +import com.azure.cosmos.implementation.changefeed.epkversion.ServiceItemLeaseV1; +import com.azure.cosmos.models.ChangeFeedProcessorOptions; +import com.azure.cosmos.models.ChangeFeedProcessorState; +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosContainerRequestOptions; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.SqlParameter; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.rx.TestSuiteBase; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.commons.lang3.RandomStringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testng.annotations.AfterClass; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; +import reactor.util.retry.Retry; + +import java.time.Duration; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Consumer; + +import static com.azure.cosmos.BridgeInternal.extractContainerSelfLink; +import static com.azure.cosmos.CosmosBridgeInternal.getContextClient; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; + +public class IncrementalChangeFeedProcessorTest extends TestSuiteBase { + private final static Logger logger = LoggerFactory.getLogger(IncrementalChangeFeedProcessorTest.class); + private static final ObjectMapper OBJECT_MAPPER = Utils.getSimpleObjectMapper(); + + private CosmosAsyncDatabase createdDatabase; + private final String hostName = RandomStringUtils.randomAlphabetic(6); + private final int FEED_COUNT = 10; + private final int CHANGE_FEED_PROCESSOR_TIMEOUT = 5000; + private final int FEED_COLLECTION_THROUGHPUT = 400; + private final int FEED_COLLECTION_THROUGHPUT_FOR_SPLIT = 10100; + private final int LEASE_COLLECTION_THROUGHPUT = 400; + + private CosmosAsyncClient client; + + private ChangeFeedProcessor changeFeedProcessor; + + @Factory(dataProvider = "clientBuilders") + public IncrementalChangeFeedProcessorTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @Test(groups = { "emulator" }, timeOut = 2 * TIMEOUT) + public void readFeedDocumentsStartFromBeginning() throws InterruptedException { + CosmosAsyncContainer createdFeedCollection = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + List createdDocuments = new ArrayList<>(); + Map receivedDocuments = new ConcurrentHashMap<>(); + setupReadFeedDocuments(createdDocuments, receivedDocuments, createdFeedCollection, FEED_COUNT); + + changeFeedProcessor = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .handleLatestVersionChanges(changeFeedProcessorHandler(receivedDocuments)) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .options(new ChangeFeedProcessorOptions() + .setLeaseRenewInterval(Duration.ofSeconds(20)) + .setLeaseAcquireInterval(Duration.ofSeconds(10)) + .setLeaseExpirationInterval(Duration.ofSeconds(30)) + .setFeedPollDelay(Duration.ofSeconds(2)) + .setLeasePrefix("TEST") + .setMaxItemCount(10) + .setStartFromBeginning(true) + .setMaxScaleCount(0) // unlimited + ) + .buildChangeFeedProcessor(); + + try { + changeFeedProcessor.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .subscribe(); + } catch (Exception ex) { + logger.error("Change feed processor did not start in the expected time", ex); + throw ex; + } + + // Wait for the feed processor to receive and process the documents. + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + assertThat(changeFeedProcessor.isStarted()).as("Change Feed Processor instance is running").isTrue(); + + changeFeedProcessor.stop().subscribeOn(Schedulers.boundedElastic()).timeout(Duration.ofMillis(CHANGE_FEED_PROCESSOR_TIMEOUT)).subscribe(); + + for (InternalObjectNode item : createdDocuments) { + assertThat(receivedDocuments.containsKey(item.getId())).as("Document with getId: " + item.getId()).isTrue(); + } + + // Wait for the feed processor to shutdown. + Thread.sleep(CHANGE_FEED_PROCESSOR_TIMEOUT); + } finally { + safeDeleteCollection(createdFeedCollection); + safeDeleteCollection(createdLeaseCollection); + + // Allow some time for the collections to be deleted before exiting. + Thread.sleep(500); + } + } + + @Test(groups = { "emulator" }, timeOut = 50 * CHANGE_FEED_PROCESSOR_TIMEOUT) + public void readFeedDocumentsStartFromCustomDate() throws InterruptedException { + CosmosAsyncContainer createdFeedCollection = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + List createdDocuments = new ArrayList<>(); + Map receivedDocuments = new ConcurrentHashMap<>(); + ChangeFeedProcessor changeFeedProcessor = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .handleLatestVersionChanges((List docs) -> { + logger.info("START processing from thread {}", Thread.currentThread().getId()); + for (JsonNode item : docs) { + processItem(item, receivedDocuments); + } + logger.info("END processing from thread {}", Thread.currentThread().getId()); + }) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .options(new ChangeFeedProcessorOptions() + .setLeaseRenewInterval(Duration.ofSeconds(20)) + .setLeaseAcquireInterval(Duration.ofSeconds(10)) + .setLeaseExpirationInterval(Duration.ofSeconds(30)) + .setFeedPollDelay(Duration.ofSeconds(1)) + .setLeasePrefix("TEST") + .setMaxItemCount(10) + .setStartTime(ZonedDateTime.now(ZoneOffset.UTC).minusDays(1).toInstant()) + .setMinScaleCount(1) + .setMaxScaleCount(3) + ) + .buildChangeFeedProcessor(); + + try { + changeFeedProcessor.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .subscribe(); + } catch (Exception ex) { + logger.error("Change feed processor did not start in the expected time", ex); + throw ex; + } + + setupReadFeedDocuments(createdDocuments, receivedDocuments, createdFeedCollection, FEED_COUNT); + + // Wait for the feed processor to receive and process the documents. + waitToReceiveDocuments(receivedDocuments, 40 * CHANGE_FEED_PROCESSOR_TIMEOUT, FEED_COUNT); + + assertThat(changeFeedProcessor.isStarted()).as("Change Feed Processor instance is running").isTrue(); + + changeFeedProcessor.stop().subscribeOn(Schedulers.boundedElastic()).timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)).subscribe(); + + for (InternalObjectNode item : createdDocuments) { + assertThat(receivedDocuments.containsKey(item.getId())).as("Document with getId: " + item.getId()).isTrue(); + } + + // Wait for the feed processor to shutdown. + Thread.sleep(CHANGE_FEED_PROCESSOR_TIMEOUT); + + } finally { + safeDeleteCollection(createdFeedCollection); + safeDeleteCollection(createdLeaseCollection); + + // Allow some time for the collections to be deleted before exiting. + Thread.sleep(500); + } + } + + @Test(groups = { "emulator" }, timeOut = 50 * CHANGE_FEED_PROCESSOR_TIMEOUT) + public void getCurrentState() throws InterruptedException { + CosmosAsyncContainer createdFeedCollection = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + List createdDocuments = new ArrayList<>(); + Map receivedDocuments = new ConcurrentHashMap<>(); + ChangeFeedProcessor changeFeedProcessorMain = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .handleLatestVersionChanges((List docs) -> { + logger.info("START processing from thread {}", Thread.currentThread().getId()); + for (JsonNode item : docs) { + processItem(item, receivedDocuments); + } + logger.info("END processing from thread {}", Thread.currentThread().getId()); + }) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .buildChangeFeedProcessor(); + + ChangeFeedProcessor changeFeedProcessorSideCart = new ChangeFeedProcessorBuilder() + .hostName("side-cart") + .handleLatestVersionChanges((List docs) -> { + fail("ERROR - we should not execute this handler"); + }) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .buildChangeFeedProcessor(); + + try { + changeFeedProcessorMain.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .then(Mono.just(changeFeedProcessorMain) + .delayElement(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .flatMap(value -> changeFeedProcessorMain.stop() + .subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + )) + .subscribe(); + } catch (Exception ex) { + logger.error("Change feed processor did not start and stopped in the expected time", ex); + throw ex; + } + + Thread.sleep(4 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + // Test for "zero" lag + List cfpCurrentState = changeFeedProcessorMain.getCurrentState() + .map(state -> { + try { + logger.info(OBJECT_MAPPER.writeValueAsString(state)); + } catch (JsonProcessingException ex) { + logger.error("Unexpected", ex); + } + return state; + }).block(); + + assertThat(cfpCurrentState.size()).isNotZero().as("Change Feed Processor number of leases should not be 0."); + + int totalLag = 0; + for (ChangeFeedProcessorState item : cfpCurrentState) { + totalLag += item.getEstimatedLag(); + } + + assertThat(totalLag).isEqualTo(0).as("Change Feed Processor Main estimated total lag at start"); + + // check the side cart CFP instance + List cfpCurrentStateSideCart = changeFeedProcessorSideCart.getCurrentState() + .map(state -> { + try { + logger.info(OBJECT_MAPPER.writeValueAsString(state)); + } catch (JsonProcessingException ex) { + logger.error("Unexpected", ex); + } + return state; + }).block(); + + assertThat(cfpCurrentStateSideCart.size()).isNotZero().as("Change Feed Processor side cart number of leases should not be 0."); + + totalLag = 0; + for (ChangeFeedProcessorState item : cfpCurrentStateSideCart) { + totalLag += item.getEstimatedLag(); + } + + assertThat(totalLag).isEqualTo(0).as("Change Feed Processor Side Cart estimated total lag at start"); + + + // Test for "FEED_COUNT total lag + setupReadFeedDocuments(createdDocuments, receivedDocuments, createdFeedCollection, FEED_COUNT); + + cfpCurrentState = changeFeedProcessorMain.getCurrentState() + .map(state -> { + try { + logger.info(OBJECT_MAPPER.writeValueAsString(state)); + } catch (JsonProcessingException ex) { + logger.error("Unexpected", ex); + } + return state; + }).block(); + + totalLag = 0; + for (ChangeFeedProcessorState item : cfpCurrentState) { + totalLag += item.getEstimatedLag(); + } + + assertThat(totalLag).isEqualTo(FEED_COUNT).as("Change Feed Processor Main estimated total lag"); + + // check the side cart CFP instance + cfpCurrentStateSideCart = changeFeedProcessorSideCart.getCurrentState() + .map(state -> { + try { + logger.info(OBJECT_MAPPER.writeValueAsString(state)); + } catch (JsonProcessingException ex) { + logger.error("Unexpected", ex); + } + return state; + }).block(); + + assertThat(cfpCurrentStateSideCart.size()).isNotZero().as("Change Feed Processor side cart number of leases should not be 0."); + + totalLag = 0; + for (ChangeFeedProcessorState item : cfpCurrentStateSideCart) { + totalLag += item.getEstimatedLag(); + } + + assertThat(totalLag).isEqualTo(FEED_COUNT).as("Change Feed Processor Side Cart estimated total lag"); + + + } finally { + safeDeleteCollection(createdFeedCollection); + safeDeleteCollection(createdLeaseCollection); + + // Allow some time for the collections to be deleted before exiting. + Thread.sleep(500); + } + } + + @Test(groups = { "emulator" }, timeOut = 50 * CHANGE_FEED_PROCESSOR_TIMEOUT) + public void staledLeaseAcquiring() throws InterruptedException { + final String ownerFirst = "Owner_First"; + final String ownerSecond = "Owner_Second"; + final String leasePrefix = "TEST"; + CosmosAsyncContainer createdFeedCollection = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + Map receivedDocuments = new ConcurrentHashMap<>(); + + ChangeFeedProcessor changeFeedProcessorFirst = new ChangeFeedProcessorBuilder() + .hostName(ownerFirst) + .handleLatestVersionChanges(docs -> { + logger.info("START processing from thread {} using host {}", Thread.currentThread().getId(), ownerFirst); + logger.info("END processing from thread {} using host {}", Thread.currentThread().getId(), ownerFirst); + }) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .options(new ChangeFeedProcessorOptions() + .setLeasePrefix(leasePrefix) + ) + .buildChangeFeedProcessor(); + + ChangeFeedProcessor changeFeedProcessorSecond = new ChangeFeedProcessorBuilder() + .hostName(ownerSecond) + .handleLatestVersionChanges((List docs) -> { + logger.info("START processing from thread {} using host {}", Thread.currentThread().getId(), ownerSecond); + for (JsonNode item : docs) { + processItem(item, receivedDocuments); + } + logger.info("END processing from thread {} using host {}", Thread.currentThread().getId(), ownerSecond); + }) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .options(new ChangeFeedProcessorOptions() + .setLeaseRenewInterval(Duration.ofSeconds(10)) + .setLeaseAcquireInterval(Duration.ofSeconds(5)) + .setLeaseExpirationInterval(Duration.ofSeconds(20)) + .setFeedPollDelay(Duration.ofSeconds(2)) + .setLeasePrefix(leasePrefix) + .setMaxItemCount(10) + .setStartFromBeginning(true) + .setMaxScaleCount(0) // unlimited + ) + .buildChangeFeedProcessor(); + + try { + changeFeedProcessorFirst.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .then(Mono.just(changeFeedProcessorFirst) + .delayElement(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .flatMap(value -> changeFeedProcessorFirst.stop() + .subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + )) + .doOnSuccess(aVoid -> { + try { + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + } catch (InterruptedException e) { + throw new RuntimeException("Interrupted exception", e); + } + logger.info("Update leases for Change feed processor in thread {} using host {}", Thread.currentThread().getId(), "Owner_first"); + + SqlParameter param = new SqlParameter(); + param.setName("@PartitionLeasePrefix"); + param.setValue(leasePrefix); + SqlQuerySpec querySpec = new SqlQuerySpec( + "SELECT * FROM c WHERE STARTSWITH(c.id, @PartitionLeasePrefix)", Collections.singletonList(param)); + + CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); + + createdLeaseCollection.queryItems(querySpec, cosmosQueryRequestOptions, InternalObjectNode.class).byPage() + .flatMap(documentFeedResponse -> reactor.core.publisher.Flux.fromIterable(documentFeedResponse.getResults())) + .flatMap(doc -> { + ServiceItemLeaseV1 leaseDocument = ServiceItemLeaseV1.fromDocument(doc); + leaseDocument.setOwner("TEMP_OWNER"); + CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + + return createdLeaseCollection.replaceItem(leaseDocument, leaseDocument.getId(), new PartitionKey(leaseDocument.getId()), options) + .map(CosmosItemResponse::getItem); + }) + .map(leaseDocument -> { + logger.info("QueryItems after Change feed processor processing; found host {}", leaseDocument.getOwner()); + return leaseDocument; + }) + .last() + .flatMap(leaseDocument -> { + logger.info("Start creating documents"); + List docDefList = new ArrayList<>(); + + for (int i = 0; i < FEED_COUNT; i++) { + docDefList.add(getDocumentDefinition()); + } + + return bulkInsert(createdFeedCollection, docDefList, FEED_COUNT) + .last() + .delayElement(Duration.ofMillis(1000)) + .flatMap(cosmosItemResponse -> { + logger.info("Start second Change feed processor"); + return changeFeedProcessorSecond.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)); + }); + }) + .subscribe(); + }) + .subscribe(); + } catch (Exception ex) { + logger.error("First change feed processor did not start in the expected time", ex); + throw ex; + } + + long remainingWork = 10 * CHANGE_FEED_PROCESSOR_TIMEOUT; + while (remainingWork > 0 && changeFeedProcessorFirst.isStarted() && !changeFeedProcessorSecond.isStarted()) { + remainingWork -= 100; + Thread.sleep(100); + } + + // Wait for the feed processor to receive and process the documents. + waitToReceiveDocuments(receivedDocuments, 10 * CHANGE_FEED_PROCESSOR_TIMEOUT, FEED_COUNT); + + assertThat(changeFeedProcessorSecond.isStarted()).as("Change Feed Processor instance is running").isTrue(); + + changeFeedProcessorSecond.stop().subscribeOn(Schedulers.boundedElastic()).timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)).subscribe(); + + // Wait for the feed processor to shutdown. + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + } finally { + safeDeleteCollection(createdFeedCollection); + safeDeleteCollection(createdLeaseCollection); + + // Allow some time for the collections to be deleted before exiting. + Thread.sleep(500); + } + } + + @Test(groups = { "emulator" }, timeOut = 50 * CHANGE_FEED_PROCESSOR_TIMEOUT) + public void ownerNullAcquiring() throws InterruptedException { + final String ownerFirst = "Owner_First"; + final String leasePrefix = "TEST"; + CosmosAsyncContainer createdFeedCollection = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + List createdDocuments = new ArrayList<>(); + Map receivedDocuments = new ConcurrentHashMap<>(); + + ChangeFeedProcessor changeFeedProcessorFirst = new ChangeFeedProcessorBuilder() + .hostName(ownerFirst) + .handleLatestVersionChanges(docs -> { + logger.info("START processing from thread {} using host {}", Thread.currentThread().getId(), ownerFirst); + for (JsonNode item : docs) { + try { + Thread.sleep(1000); + } catch (InterruptedException ignored) { + } + processItem(item, receivedDocuments); + } + logger.info("END processing from thread {} using host {}", Thread.currentThread().getId(), ownerFirst); + }) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .options(new ChangeFeedProcessorOptions() + .setStartFromBeginning(true) + .setLeasePrefix(leasePrefix) + .setLeaseRenewInterval(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .setLeaseAcquireInterval(Duration.ofMillis(5 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .setLeaseExpirationInterval(Duration.ofMillis(6 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .setFeedPollDelay(Duration.ofSeconds(5)) + ) + .buildChangeFeedProcessor(); + + try { + logger.info("Start more creating documents"); + List docDefList = new ArrayList<>(); + + for (int i = 0; i < FEED_COUNT; i++) { + docDefList.add(getDocumentDefinition()); + } + + bulkInsert(createdFeedCollection, docDefList, FEED_COUNT) + .last() + .flatMap(cosmosItemResponse -> { + logger.info("Start first Change feed processor"); + return changeFeedProcessorFirst.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)); + }) + .then( + Mono.just(changeFeedProcessorFirst) + .flatMap( value -> { + logger.info("Update leases for Change feed processor in thread {} using host {}", Thread.currentThread().getId(), "Owner_first"); + try { + Thread.sleep(CHANGE_FEED_PROCESSOR_TIMEOUT); + } catch (InterruptedException ignored) { + } + + logger.info("QueryItems before Change feed processor processing"); + + SqlParameter param1 = new SqlParameter(); + param1.setName("@PartitionLeasePrefix"); + param1.setValue(leasePrefix); + SqlParameter param2 = new SqlParameter(); + param2.setName("@Owner"); + param2.setValue(ownerFirst); + + SqlQuerySpec querySpec = new SqlQuerySpec( + "SELECT * FROM c WHERE STARTSWITH(c.id, @PartitionLeasePrefix) AND c.Owner=@Owner", Arrays.asList(param1, param2)); + + CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); + + return createdLeaseCollection.queryItems(querySpec, cosmosQueryRequestOptions, InternalObjectNode.class).byPage() + .flatMap(documentFeedResponse -> reactor.core.publisher.Flux.fromIterable(documentFeedResponse.getResults())) + .flatMap(doc -> { + ServiceItemLeaseV1 leaseDocument = ServiceItemLeaseV1.fromDocument(doc); + leaseDocument.setOwner(null); + CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + return createdLeaseCollection.replaceItem(leaseDocument, leaseDocument.getId(), new PartitionKey(leaseDocument.getId()), options) + .map(CosmosItemResponse::getItem); + }) + .map(leaseDocument -> { + logger.info("QueryItems after Change feed processor processing; current Owner is'{}'", leaseDocument.getOwner()); + return leaseDocument; + }) + .last() + .flatMap(leaseDocument -> { + logger.info("Start creating more documents"); + List docDefList1 = new ArrayList<>(); + + for (int i = 0; i < FEED_COUNT; i++) { + docDefList1.add(getDocumentDefinition()); + } + + return bulkInsert(createdFeedCollection, docDefList1, FEED_COUNT) + .last(); + }); + })) + .subscribe(); + } catch (Exception ex) { + logger.error("First change feed processor did not start in the expected time", ex); + throw ex; + } + + long remainingWork = 20 * CHANGE_FEED_PROCESSOR_TIMEOUT; + while (remainingWork > 0 && !changeFeedProcessorFirst.isStarted()) { + remainingWork -= 100; + Thread.sleep(100); + } + + // Wait for the feed processor to receive and process the documents. + waitToReceiveDocuments(receivedDocuments, 30 * CHANGE_FEED_PROCESSOR_TIMEOUT, 2 * FEED_COUNT); + + assertThat(changeFeedProcessorFirst.isStarted()).as("Change Feed Processor instance is running").isTrue(); + + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + changeFeedProcessorFirst.stop().subscribeOn(Schedulers.boundedElastic()).timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)).subscribe(); + + // Wait for the feed processor to shutdown. + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + } finally { + safeDeleteCollection(createdFeedCollection); + safeDeleteCollection(createdLeaseCollection); + + // Allow some time for the collections to be deleted before exiting. + Thread.sleep(500); + } + } + + @Test(groups = { "simple" }, timeOut = 160 * CHANGE_FEED_PROCESSOR_TIMEOUT) + public void readFeedDocumentsAfterSplit() throws InterruptedException { + CosmosAsyncContainer createdFeedCollectionForSplit = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(2 * LEASE_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseMonitorCollection = createLeaseMonitorCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + List createdDocuments = new ArrayList<>(); + Map receivedDocuments = new ConcurrentHashMap<>(); + LeaseStateMonitor leaseStateMonitor = new LeaseStateMonitor(); + + // create a monitoring CFP for ensuring the leases are updating as expected + ChangeFeedProcessor leaseMonitoringChangeFeedProcessor = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .handleLatestVersionChanges(leasesChangeFeedProcessorHandler(leaseStateMonitor)) + .feedContainer(createdLeaseCollection) + .leaseContainer(createdLeaseMonitorCollection) + .options(new ChangeFeedProcessorOptions() + .setLeasePrefix("MONITOR") + .setStartFromBeginning(true) + .setMaxItemCount(100) + .setLeaseExpirationInterval(Duration.ofMillis(10 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .setFeedPollDelay(Duration.ofMillis(200)) + ) + .buildChangeFeedProcessor(); + + // generate a first batch of documents + setupReadFeedDocuments(createdDocuments, receivedDocuments, createdFeedCollectionForSplit, FEED_COUNT); + + changeFeedProcessor = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .handleLatestVersionChanges(changeFeedProcessorHandler(receivedDocuments)) + .feedContainer(createdFeedCollectionForSplit) + .leaseContainer(createdLeaseCollection) + .options(new ChangeFeedProcessorOptions() + .setLeasePrefix("TEST") + .setStartFromBeginning(true) + .setMaxItemCount(10) + .setLeaseRenewInterval(Duration.ofSeconds(2)) + ) + .buildChangeFeedProcessor(); + + leaseMonitoringChangeFeedProcessor.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .onErrorResume(throwable -> { + logger.error("Change feed processor for lease monitoring did not start in the expected time", throwable); + return Mono.error(throwable); + }) + .then( + changeFeedProcessor.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .onErrorResume(throwable -> { + logger.error("Change feed processor did not start in the expected time", throwable); + return Mono.error(throwable); + }) + .doOnSuccess(aVoid -> { + // Wait for the feed processor to receive and process the first batch of documents. + try { + waitToReceiveDocuments(receivedDocuments, 2 * CHANGE_FEED_PROCESSOR_TIMEOUT, FEED_COUNT); + } catch (InterruptedException e) { + throw new RuntimeException("Interrupted exception", e); + } + }) + .then( + // increase throughput to force a single partition collection to go through a split + createdFeedCollectionForSplit + .readThroughput().subscribeOn(Schedulers.boundedElastic()) + .flatMap(currentThroughput -> + createdFeedCollectionForSplit + .replaceThroughput(ThroughputProperties.createManualThroughput(FEED_COLLECTION_THROUGHPUT_FOR_SPLIT)) + .subscribeOn(Schedulers.boundedElastic()) + ) + .onErrorResume(throwable -> { + System.out.println(throwable); + return Mono.empty(); + }) + .then() + ) + ) + .subscribe(); + + // Wait for the feed processor to receive and process the first batch of documents and apply throughput change. + Thread.sleep(4 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + // Retrieve the latest continuation token value. + long continuationToken = Long.MAX_VALUE; + for (JsonNode item : leaseStateMonitor.receivedLeases.values()) { + JsonNode tempToken = item.get("ContinuationToken"); + + long continuationTokenValue = 0; + if (tempToken != null && StringUtils.isNotEmpty(tempToken.asText())) { + ChangeFeedState changeFeedState = ChangeFeedState.fromString(tempToken.asText()); + continuationTokenValue = + Long.parseLong(changeFeedState.getContinuation().getCurrentContinuationToken().getToken().replace("\"", "")); + } + if (tempToken == null || continuationTokenValue == 0) { + logger.error("Found unexpected lease with continuation token value of null or 0"); + try { + logger.info("ERROR LEASE FOUND {}", OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(item)); + } catch (JsonProcessingException e) { + logger.error("Failure in processing json [{}]", e.getMessage(), e); + } + leaseStateMonitor.isContinuationTokenAdvancing = false; + } + else { + // keep the lowest continuation token value + if (continuationToken > continuationTokenValue) { + continuationToken = continuationTokenValue; + } + } + } + if (continuationToken == Long.MAX_VALUE) { + // something went wrong; we could not find any valid leases. + logger.error("Could not find any valid lease documents"); + leaseStateMonitor.isContinuationTokenAdvancing = false; + } + else { + leaseStateMonitor.parentContinuationToken = continuationToken; + } + leaseStateMonitor.isAfterLeaseInitialization = true; + + // Loop through reading the current partition count until we get a split + // This can take up to two minute or more. + String partitionKeyRangesPath = extractContainerSelfLink(createdFeedCollectionForSplit); + + AsyncDocumentClient contextClient = getContextClient(createdDatabase); + Flux.just(1).subscribeOn(Schedulers.boundedElastic()) + .flatMap(value -> { + logger.warn("Reading current throughput change."); + return contextClient.readPartitionKeyRanges(partitionKeyRangesPath, null); + }) + .map(partitionKeyRangeFeedResponse -> { + int count = partitionKeyRangeFeedResponse.getResults().size(); + + if (count < 2) { + logger.warn("Throughput change is pending."); + throw new RuntimeException("Throughput change is not done."); + } + return count; + }) + // this will timeout approximately after 30 minutes + .retryWhen(Retry.max(40).filter(throwable -> { + try { + logger.warn("Retrying..."); + // Splits are taking longer, so increasing sleep time between retries + Thread.sleep(10 * CHANGE_FEED_PROCESSOR_TIMEOUT); + } catch (InterruptedException e) { + throw new RuntimeException("Interrupted exception", e); + } + return true; + })) + .last() + .doOnSuccess(partitionCount -> { + leaseStateMonitor.isAfterSplits = true; + }) + .block(); + + assertThat(changeFeedProcessor.isStarted()).as("Change Feed Processor instance is running").isTrue(); + + // generate the second batch of documents + createReadFeedDocuments(createdDocuments, createdFeedCollectionForSplit, FEED_COUNT); + + // Wait for the feed processor to receive and process the second batch of documents. + waitToReceiveDocuments(receivedDocuments, 2 * CHANGE_FEED_PROCESSOR_TIMEOUT, FEED_COUNT * 2); + + changeFeedProcessor.stop().subscribeOn(Schedulers.boundedElastic()).timeout(Duration.ofMillis(CHANGE_FEED_PROCESSOR_TIMEOUT)).subscribe(); + leaseMonitoringChangeFeedProcessor.stop().subscribeOn(Schedulers.boundedElastic()).timeout(Duration.ofMillis(CHANGE_FEED_PROCESSOR_TIMEOUT)).subscribe(); + + int leaseCount = changeFeedProcessor.getCurrentState() .map(List::size).block(); + assertThat(leaseCount > 1).as("Found %d leases", leaseCount).isTrue(); + + assertThat(receivedDocuments.size()).isEqualTo(createdDocuments.size()); + for (InternalObjectNode item : createdDocuments) { + assertThat(receivedDocuments.containsKey(item.getId())).as("Document with getId: " + item.getId()).isTrue(); + } + + // check the continuation tokens have advanced after splits + assertThat(leaseStateMonitor.isContinuationTokenAdvancing && leaseStateMonitor.parentContinuationToken > 0) + .as("Continuation tokens for the leases after split should advance from parent value; parent: %d", leaseStateMonitor.parentContinuationToken).isTrue(); + + // Wait for the feed processor to shutdown. + Thread.sleep(2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + } finally { + System.out.println("Start to delete FeedCollectionForSplit"); + safeDeleteCollection(createdFeedCollectionForSplit); + safeDeleteCollection(createdLeaseCollection); + + // Allow some time for the collections to be deleted before exiting. + Thread.sleep(500); + } + } + + @Test(groups = { "emulator" }, timeOut = 20 * TIMEOUT) + public void inactiveOwnersRecovery() throws InterruptedException { + CosmosAsyncContainer createdFeedCollection = createFeedCollection(FEED_COLLECTION_THROUGHPUT); + CosmosAsyncContainer createdLeaseCollection = createLeaseCollection(LEASE_COLLECTION_THROUGHPUT); + + try { + List createdDocuments = new ArrayList<>(); + Map receivedDocuments = new ConcurrentHashMap<>(); + String leasePrefix = "TEST"; + + changeFeedProcessor = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .handleLatestVersionChanges(changeFeedProcessorHandler(receivedDocuments)) + .feedContainer(createdFeedCollection) + .leaseContainer(createdLeaseCollection) + .options(new ChangeFeedProcessorOptions() + .setLeaseRenewInterval(Duration.ofSeconds(1)) + .setLeaseAcquireInterval(Duration.ofSeconds(1)) + .setLeaseExpirationInterval(Duration.ofSeconds(5)) + .setFeedPollDelay(Duration.ofSeconds(1)) + .setLeasePrefix(leasePrefix) + .setMaxItemCount(100) + .setStartFromBeginning(true) + .setMaxScaleCount(0) // unlimited + //.setScheduler(Schedulers.boundedElastic()) + .setScheduler(Schedulers.newParallel("CFP parallel", + 10 * Schedulers.DEFAULT_POOL_SIZE, + true)) + ) + .buildChangeFeedProcessor(); + + setupReadFeedDocuments(createdDocuments, receivedDocuments, createdFeedCollection, FEED_COUNT); + + validateChangeFeedProcessing(changeFeedProcessor, createdDocuments, receivedDocuments,2 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + Thread.sleep(CHANGE_FEED_PROCESSOR_TIMEOUT); + + logger.info("Update leases with random owners"); + + SqlParameter param1 = new SqlParameter(); + param1.setName("@PartitionLeasePrefix"); + param1.setValue(leasePrefix); + + SqlQuerySpec querySpec = new SqlQuerySpec( + "SELECT * FROM c WHERE STARTSWITH(c.id, @PartitionLeasePrefix)", Arrays.asList(param1)); + + CosmosQueryRequestOptions cosmosQueryRequestOptions = new CosmosQueryRequestOptions(); + + createdLeaseCollection.queryItems(querySpec, cosmosQueryRequestOptions, InternalObjectNode.class).byPage() + .flatMap(documentFeedResponse -> reactor.core.publisher.Flux.fromIterable(documentFeedResponse.getResults())) + .flatMap(doc -> { + ServiceItemLeaseV1 leaseDocument = ServiceItemLeaseV1.fromDocument(doc); + leaseDocument.setOwner(RandomStringUtils.randomAlphabetic(10)); + CosmosItemRequestOptions options = new CosmosItemRequestOptions(); + return createdLeaseCollection.replaceItem(leaseDocument, leaseDocument.getId(), new PartitionKey(leaseDocument.getId()), options) + .map(CosmosItemResponse::getItem); + }) + .flatMap(leaseDocument -> createdLeaseCollection.readItem(leaseDocument.getId(), new PartitionKey(leaseDocument.getId()), InternalObjectNode.class)) + .map(doc -> { + ServiceItemLeaseV1 leaseDocument = ServiceItemLeaseV1.fromDocument(doc.getItem()); + logger.info("Change feed processor current Owner is'{}'", leaseDocument.getOwner()); + return leaseDocument; + }) + .blockLast(); + + createdDocuments.clear(); + receivedDocuments.clear(); + setupReadFeedDocuments(createdDocuments, receivedDocuments, createdFeedCollection, FEED_COUNT); + + validateChangeFeedProcessing(changeFeedProcessor, createdDocuments, receivedDocuments, 10 * CHANGE_FEED_PROCESSOR_TIMEOUT); + + // Wait for the feed processor to shutdown. + Thread.sleep(CHANGE_FEED_PROCESSOR_TIMEOUT); + } finally { + safeDeleteCollection(createdFeedCollection); + safeDeleteCollection(createdLeaseCollection); + + // Allow some time for the collections to be deleted before exiting. + Thread.sleep(500); + } + } + + void validateChangeFeedProcessing(ChangeFeedProcessor changeFeedProcessor, List createdDocuments, Map receivedDocuments, int sleepTime) throws InterruptedException { + try { + changeFeedProcessor.start().subscribeOn(Schedulers.boundedElastic()) + .timeout(Duration.ofMillis(2 * CHANGE_FEED_PROCESSOR_TIMEOUT)) + .subscribe(); + } catch (Exception ex) { + logger.error("Change feed processor did not start in the expected time", ex); + throw ex; + } + + // Wait for the feed processor to receive and process the documents. + Thread.sleep(sleepTime); + + assertThat(changeFeedProcessor.isStarted()).as("Change Feed Processor instance is running").isTrue(); + + List cfpCurrentState = changeFeedProcessor + .getCurrentState() + .map(state -> { + try { + logger.info(OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(state)); + } catch (JsonProcessingException ex) { + logger.error("Unexpected", ex); + } + return state; + }) + .block(); + + assertThat(cfpCurrentState).isNotNull().as("Change Feed Processor current state"); + + for (ChangeFeedProcessorState item : cfpCurrentState) { + assertThat(item.getHostName()).isEqualTo(hostName).as("Change Feed Processor ownership"); + } + + changeFeedProcessor.stop().subscribeOn(Schedulers.boundedElastic()).timeout(Duration.ofMillis(CHANGE_FEED_PROCESSOR_TIMEOUT)).subscribe(); + + for (InternalObjectNode item : createdDocuments) { + assertThat(receivedDocuments.containsKey(item.getId())).as("Document with getId: " + item.getId()).isTrue(); + } + } + + private Consumer> changeFeedProcessorHandler(Map receivedDocuments) { + return docs -> { + logger.info("START processing from thread in test {}", Thread.currentThread().getId()); + for (JsonNode item : docs) { + processItem(item, receivedDocuments); + } + logger.info("END processing from thread {}", Thread.currentThread().getId()); + }; + } + + private void waitToReceiveDocuments(Map receivedDocuments, long timeoutInMillisecond, long count) throws InterruptedException { + long remainingWork = timeoutInMillisecond; + while (remainingWork > 0 && receivedDocuments.size() < count) { + remainingWork -= 100; + Thread.sleep(100); + } + + assertThat(remainingWork > 0).as("Failed to receive all the feed documents").isTrue(); + } + + private Consumer> leasesChangeFeedProcessorHandler(LeaseStateMonitor leaseStateMonitor) { + return docs -> { + logger.info("LEASES processing from thread in test {}", Thread.currentThread().getId()); + for (JsonNode item : docs) { + try { + logger + .debug("LEASE RECEIVED {}", OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(item)); + } catch (JsonProcessingException e) { + logger.error("Failure in processing json [{}]", e.getMessage(), e); + } + + JsonNode leaseToken = item.get("LeaseToken"); + + if (leaseToken != null) { + JsonNode continuationTokenNode = item.get("ContinuationToken"); + if (continuationTokenNode == null) { + // Something catastrophic went wrong and the lease is malformed. + logger.error("Found invalid lease document"); + leaseStateMonitor.isContinuationTokenAdvancing = false; + } + else { + logger.info("LEASE {} with continuation {}", leaseToken.asText(), continuationTokenNode.asText()); + if (leaseStateMonitor.isAfterLeaseInitialization) { + ChangeFeedState changeFeedState = ChangeFeedState.fromString(continuationTokenNode.asText()); + long continuationToken = Long.parseLong(changeFeedState.getContinuation().getCurrentContinuationToken().getToken().replace("\"", "")); + if (leaseStateMonitor.parentContinuationToken > continuationToken) { + logger.error("Found unexpected continuation token that did not advance after the split; parent: {}, current: {}"); + leaseStateMonitor.isContinuationTokenAdvancing = false; + } + } + } + leaseStateMonitor.receivedLeases.put(item.get("id").asText(), item); + } + } + logger.info("LEASES processing from thread {}", Thread.currentThread().getId()); + }; + } + + @BeforeMethod(groups = { "emulator", "simple" }, timeOut = 2 * SETUP_TIMEOUT, alwaysRun = true) + public void beforeMethod() { + } + + @BeforeClass(groups = { "emulator", "simple" }, timeOut = SETUP_TIMEOUT, alwaysRun = true) + public void before_ChangeFeedProcessorTest() { + client = getClientBuilder().buildAsyncClient(); + createdDatabase = getSharedCosmosDatabase(client); + + // Following is code that when enabled locally it allows for a predicted database/collection name that can be + // checked in the Azure Portal +// try { +// client.getDatabase(databaseId).read() +// .map(cosmosDatabaseResponse -> cosmosDatabaseResponse.getDatabase()) +// .flatMap(database -> database.delete()) +// .onErrorResume(throwable -> { +// if (throwable instanceof com.azure.cosmos.CosmosClientException) { +// com.azure.cosmos.CosmosClientException clientException = (com.azure.cosmos.CosmosClientException) throwable; +// if (clientException.getStatusCode() == 404) { +// return Mono.empty(); +// } +// } +// return Mono.error(throwable); +// }).block(); +// Thread.sleep(500); +// } catch (Exception e){ +// log.warn("Database delete", e); +// } +// createdDatabase = createDatabase(client, databaseId); + } + + @AfterMethod(groups = { "emulator", "simple" }, timeOut = 3 * SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterMethod() { + } + + @AfterClass(groups = { "emulator", "simple" }, timeOut = 2 * SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { +// try { +// client.readAllDatabases() +// .flatMap(cosmosDatabaseProperties -> { +// CosmosAsyncDatabase cosmosDatabase = client.getDatabase(cosmosDatabaseProperties.getId()); +// return cosmosDatabase.delete(); +// }).blockLast(); +// Thread.sleep(500); +// } catch (Exception e){ } + + safeClose(client); + } + + private void setupReadFeedDocuments(List createdDocuments, Map receivedDocuments, CosmosAsyncContainer feedCollection, long count) { + List docDefList = new ArrayList<>(); + + for(int i = 0; i < count; i++) { + docDefList.add(getDocumentDefinition()); + } + + createdDocuments.addAll(bulkInsertBlocking(feedCollection, docDefList)); + waitIfNeededForReplicasToCatchUp(getClientBuilder()); + } + + private void createReadFeedDocuments(List createdDocuments, CosmosAsyncContainer feedCollection, long count) { + List docDefList = new ArrayList<>(); + + for(int i = 0; i < count; i++) { + docDefList.add(getDocumentDefinition()); + } + + createdDocuments.addAll(bulkInsertBlocking(feedCollection, docDefList)); + waitIfNeededForReplicasToCatchUp(getClientBuilder()); + } + + private InternalObjectNode getDocumentDefinition() { + String uuid = UUID.randomUUID().toString(); + InternalObjectNode doc = new InternalObjectNode(String.format("{ " + + "\"id\": \"%s\", " + + "\"mypk\": \"%s\", " + + "\"sgmts\": [[6519456, 1471916863], [2498434, 1455671440]]" + + "}" + , uuid, uuid)); + return doc; + } + + private CosmosAsyncContainer createFeedCollection(int provisionedThroughput) { + CosmosContainerRequestOptions optionsFeedCollection = new CosmosContainerRequestOptions(); + return createCollection(createdDatabase, getCollectionDefinition(), optionsFeedCollection, provisionedThroughput); + } + + private CosmosAsyncContainer createLeaseCollection(int provisionedThroughput) { + CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); + CosmosContainerProperties collectionDefinition = new CosmosContainerProperties( + "leases_" + UUID.randomUUID(), + "/id"); + return createCollection(createdDatabase, collectionDefinition, options, provisionedThroughput); + } + + private CosmosAsyncContainer createLeaseMonitorCollection(int provisionedThroughput) { + CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); + CosmosContainerProperties collectionDefinition = new CosmosContainerProperties( + "monitor_" + UUID.randomUUID(), + "/id"); + return createCollection(createdDatabase, collectionDefinition, options, provisionedThroughput); + } + + private static synchronized void processItem(JsonNode item, Map receivedDocuments) { + try { + logger.info("RECEIVED {}", OBJECT_MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(item)); + } catch (JsonProcessingException e) { + logger.error("Failure in processing json [{}]", e.getMessage(), e); + } + receivedDocuments.put(item.get("id").asText(), item); + } + + class LeaseStateMonitor { + public Map receivedLeases = new ConcurrentHashMap<>(); + public volatile boolean isAfterLeaseInitialization = false; + public volatile boolean isAfterSplits = false; + public volatile long parentContinuationToken = 0; + public volatile boolean isContinuationTokenAdvancing = true; + } +}