diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseWritesWithoutContentResponseTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseWritesWithoutContentResponseTests.java new file mode 100644 index 0000000000000..caa4cddffb6a0 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseWritesWithoutContentResponseTests.java @@ -0,0 +1,509 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.changefeed.epkversion; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.ChangeFeedProcessor; +import com.azure.cosmos.ChangeFeedProcessorBuilder; +import com.azure.cosmos.ConsistencyLevel; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosBridgeInternal; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.InternalObjectNode; +import com.azure.cosmos.implementation.ResourceResponse; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.implementation.Document; +import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient; +import com.azure.cosmos.implementation.changefeed.Lease; +import com.azure.cosmos.implementation.changefeed.LeaseStoreManager; +import com.azure.cosmos.implementation.changefeed.common.PartitionedByIdCollectionRequestOptionsFactory; +import com.azure.cosmos.implementation.changefeed.exceptions.LeaseLostException; +import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl; +import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.models.ChangeFeedProcessorItem; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.PartitionKey; +import com.fasterxml.jackson.databind.JsonNode; +import org.mockito.Mockito; +import org.testng.annotations.Test; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collections; +import java.util.List; +import java.util.function.Consumer; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Validates that change feed processor lease writes (epk version) are issued with content response on write disabled + * and that the lease concurrency token (ETag) is correctly taken from the response headers. + */ +public class LeaseWritesWithoutContentResponseTests { + private static final String OWNER = "host-1"; + private static final String LEASE_PREFIX = "prefix"; + + @Test(groups = "unit") + public void updateLeaseWithoutContentResponseUsesETagFromResponseHeader() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces(clientMock, noContentResponse("etag-2")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, + cachedLease.getId(), + new PartitionKey(cachedLease.getId()), + new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken(token("continuation-2")); + return lease; + }).block(); + + assertThat(updatedLease).isSameAs(cachedLease); + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-2"); + assertThat(updatedLease.getContinuationToken()).isEqualTo(token("continuation-2")); + assertThat(updatedLease.getOwner()).isEqualTo(OWNER); + + assertThat(replaces).hasSize(1); + assertThat(replaces.get(0).ifMatchETag).isEqualTo("etag-1"); + assertThat(replaces.get(0).contentResponseOnWriteEnabled).isFalse(); + verify(clientMock, never()).readItem(anyString(), any(), any(), any()); + } + + @Test(groups = "unit") + public void sequentialUpdatesUseLatestETagAsIfMatch() { + // Regression test for the stale If-Match ETag issue - each update must use the ETag returned by the previous + // write, otherwise every subsequent checkpoint would fail with 412 and require a recovery read. + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces( + clientMock, + noContentResponse("etag-2"), + noContentResponse("etag-3"), + noContentResponse("etag-4")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + for (int i = 2; i <= 4; i++) { + String continuation = token("continuation-" + i); + Lease updatedLease = updater.updateLease( + cachedLease, + cachedLease.getId(), + new PartitionKey(cachedLease.getId()), + new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken(continuation); + return lease; + }).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-" + i); + assertThat(updatedLease.getContinuationToken()).isEqualTo(continuation); + } + + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-2", "etag-3"); + assertThat(replaces).allMatch(write -> !write.contentResponseOnWriteEnabled); + verify(clientMock, never()).readItem(anyString(), any(), any(), any()); + } + + @Test(groups = "unit") + public void updateLeaseDisablesContentResponseEvenWithNullRequestOptions() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces(clientMock, noContentResponse("etag-2")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), null, lease -> lease).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-2"); + assertThat(replaces).hasSize(1); + assertThat(replaces.get(0).ifMatchETag).isEqualTo("etag-1"); + assertThat(replaces.get(0).contentResponseOnWriteEnabled).isFalse(); + } + + @Test(groups = "unit") + public void updateLeaseWithContentResponseStillUsesDocument() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + CosmosItemResponse response = contentResponse( + leaseDocument("lease-id", "etag-from-body", OWNER, token("continuation-from-body")), + "etag-from-header"); + recordReplaces(clientMock, response); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-from-body"); + // Legacy behavior: setServiceItemLease keeps the cached (submitted) continuation token. + assertThat(updatedLease.getContinuationToken()).isEqualTo(token("continuation-1")); + } + + @Test(groups = "unit") + public void updateLeaseWithoutETagKeepsCachedTokenAndRecoversOnNextUpdate() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces( + clientMock, + noContentResponse(null), + preconditionFailed(), + noContentResponse("etag-4")); + List reads = recordReads(clientMock, leaseDocument("lease-id", "etag-3", OWNER, token("continuation-2"))); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + Lease firstUpdate = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease).block(); + assertThat(firstUpdate.getConcurrencyToken()).isEqualTo("etag-1"); + + Lease secondUpdate = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken(token("continuation-3")); + return lease; + }).block(); + + assertThat(secondUpdate.getConcurrencyToken()).isEqualTo("etag-4"); + assertThat(secondUpdate.getContinuationToken()).isEqualTo(token("continuation-3")); + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-1", "etag-3"); + assertThat(reads).hasSize(1); + } + + @Test(groups = "unit") + public void updateLeasePreconditionFailedRefreshesETagAndRetries() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces(clientMock, preconditionFailed(), noContentResponse("etag-3")); + List reads = recordReads(clientMock, leaseDocument("lease-id", "etag-2", OWNER, token("continuation-1"))); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken(token("continuation-2")); + return lease; + }).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-3"); + assertThat(updatedLease.getContinuationToken()).isEqualTo(token("continuation-2")); + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-2"); + assertThat(replaces).allMatch(write -> !write.contentResponseOnWriteEnabled); + assertThat(reads).containsExactly("lease-id"); + } + + @Test(groups = "unit") + public void updateLeasePreconditionFailedWithNewOwnerThrowsLeaseLost() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordReplaces(clientMock, preconditionFailed()); + recordReads(clientMock, leaseDocument("lease-id", "etag-2", "host-2", token("continuation-1"))); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + StepVerifier.create(updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease)) + .expectError(LeaseLostException.class) + .verify(); + } + + @Test(groups = "unit") + public void updateLeaseConflictOrNotFoundThrowsLeaseLost() { + for (int statusCode : new int[] { 409, 404 }) { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordReplaces(clientMock, BridgeInternal.createCosmosException(statusCode)); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLeaseV1 cachedLease = createLease("etag-1"); + + StepVerifier.create(updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease)) + .expectError(LeaseLostException.class) + .verify(); + } + } + + @Test(groups = "unit") + public void leaseStoreManagerAcquireThenUpdatePropertiesChainsETags() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces( + clientMock, + noContentResponse("etag-2"), + noContentResponse("etag-3")); + + LeaseStoreManager leaseStoreManager = createLeaseStoreManager(clientMock); + ServiceItemLeaseV1 lease = createLease("etag-1"); + lease.setOwner("host-0"); + + Lease acquiredLease = leaseStoreManager.acquire(lease).block(); + assertThat(acquiredLease.getOwner()).isEqualTo(OWNER); + assertThat(acquiredLease.getConcurrencyToken()).isEqualTo("etag-2"); + + acquiredLease.setProperties(Collections.singletonMap("key", "value")); + Lease updatedLease = leaseStoreManager.updateProperties(acquiredLease).block(); + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-3"); + assertThat(updatedLease.getProperties()).containsEntry("key", "value"); + + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-2"); + assertThat(replaces).allMatch(write -> !write.contentResponseOnWriteEnabled); + verify(clientMock, never()).readItem(anyString(), any(), any(), any()); + } + + @Test(groups = "unit") + public void createLeaseIfNotExistWithoutContentResponseUsesETagFromResponseHeader() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List creates = recordCreates(clientMock, noContentResponse("etag-created")); + + LeaseStoreManager leaseStoreManager = createLeaseStoreManager(clientMock); + FeedRangeEpkImpl feedRange = new FeedRangeEpkImpl(new Range<>("AA", "BB", true, false)); + + Lease lease = leaseStoreManager.createLeaseIfNotExist(feedRange, token("continuation-0")).block(); + + assertThat(lease).isNotNull(); + assertThat(lease.getId()).isEqualTo(LEASE_PREFIX + "..AA-BB"); + assertThat(lease.getLeaseToken()).isEqualTo("AA-BB"); + assertThat(lease.getContinuationToken()).isEqualTo(token("continuation-0")); + assertThat(lease.getConcurrencyToken()).isEqualTo("etag-created"); + assertThat(lease.toString()).isNotNull(); + assertThat(creates).hasSize(1); + assertThat(creates.get(0).contentResponseOnWriteEnabled).isFalse(); + } + + @Test(groups = "unit") + public void createLeaseIfNotExistConflictReturnsEmpty() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordCreates(clientMock, BridgeInternal.createCosmosException(409)); + + LeaseStoreManager leaseStoreManager = createLeaseStoreManager(clientMock); + FeedRangeEpkImpl feedRange = new FeedRangeEpkImpl(new Range<>("AA", "BB", true, false)); + + StepVerifier.create(leaseStoreManager.createLeaseIfNotExist(feedRange, token("continuation-0"))) + .verifyComplete(); + } + + @Test(groups = "unit") + public void leaseStoreInitializationWritesWithoutContentResponse() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + CosmosAsyncContainer leaseContainer = mock(CosmosAsyncContainer.class); + List creates = recordCreates( + clientMock, + noContentResponse("lock-etag"), + noContentResponse("marker-etag")); + List deletes = new ArrayList<>(); + doAnswer(invocation -> { + CosmosItemRequestOptions options = invocation.getArgument(2); + deletes.add(options); + return Mono.just(noContentResponse(null)); + }).when(clientMock).deleteItem(anyString(), any(), any()); + + LeaseStoreImpl leaseStore = new LeaseStoreImpl( + clientMock, LEASE_PREFIX, leaseContainer, new PartitionedByIdCollectionRequestOptionsFactory()); + + StepVerifier.create(leaseStore.acquireInitializationLock(Duration.ofSeconds(30))) + .expectNext(true) + .verifyComplete(); + StepVerifier.create(leaseStore.markInitialized()) + .expectNext(true) + .verifyComplete(); + leaseStore.releaseInitializationLock().block(); + + assertThat(creates).hasSize(2); + assertThat(creates).allMatch(write -> !write.contentResponseOnWriteEnabled); + assertThat(creates.get(0).documentId).isEqualTo(LEASE_PREFIX + ".lock"); + assertThat(creates.get(1).documentId).isEqualTo(LEASE_PREFIX + ".info"); + assertThat(deletes).hasSize(1); + assertThat(deletes.get(0).getIfMatchETag()).isEqualTo("lock-etag"); + } + + @Test(groups = "unit") + public void leaseStoreAcquireInitializationLockConflictReturnsFalse() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordCreates(clientMock, BridgeInternal.createCosmosException(409)); + + LeaseStoreImpl leaseStore = new LeaseStoreImpl( + clientMock, LEASE_PREFIX, mock(CosmosAsyncContainer.class), new PartitionedByIdCollectionRequestOptionsFactory()); + + StepVerifier.create(leaseStore.acquireInitializationLock(Duration.ofSeconds(30))) + .expectNext(false) + .verifyComplete(); + } + + @Test(groups = "unit") + public void changeFeedProcessorAllowsLeaseClientWithContentResponseOnWriteDisabled() { + CosmosAsyncContainer feedContainer = mockContainer(true); + CosmosAsyncContainer leaseContainer = mockContainer(false); + + Consumer> handler = docs -> { }; + ChangeFeedProcessor incrementalProcessor = new ChangeFeedProcessorBuilder() + .hostName(OWNER) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .handleChanges(handler) + .buildChangeFeedProcessor(); + assertThat(incrementalProcessor).isNotNull(); + + Consumer> latestVersionHandler = docs -> { }; + ChangeFeedProcessor latestVersionProcessor = new ChangeFeedProcessorBuilder() + .hostName(OWNER) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .handleLatestVersionChanges(latestVersionHandler) + .buildChangeFeedProcessor(); + assertThat(latestVersionProcessor).isNotNull(); + + ChangeFeedProcessor allVersionsProcessor = new ChangeFeedProcessorBuilder() + .hostName(OWNER) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .handleAllVersionsAndDeletesChanges(latestVersionHandler) + .buildChangeFeedProcessor(); + assertThat(allVersionsProcessor).isNotNull(); + } + + static CosmosAsyncContainer mockContainer(boolean contentResponseOnWriteEnabled) { + CosmosAsyncContainer container = mock(CosmosAsyncContainer.class, Mockito.RETURNS_DEEP_STUBS); + AsyncDocumentClient documentClient = mock(AsyncDocumentClient.class); + when(documentClient.isContentResponseOnWriteEnabled()).thenReturn(contentResponseOnWriteEnabled); + when(documentClient.getConsistencyLevel()).thenReturn(ConsistencyLevel.SESSION); + when(CosmosBridgeInternal.getContextClient(container)).thenReturn(documentClient); + return container; + } + + private static LeaseStoreManager createLeaseStoreManager(ChangeFeedContextClient clientMock) { + return LeaseStoreManagerImpl.builder() + .leasePrefix(LEASE_PREFIX) + .leaseCollectionLink(mock(CosmosAsyncContainer.class)) + .leaseContextClient(clientMock) + .requestOptionsFactory(new PartitionedByIdCollectionRequestOptionsFactory()) + .hostName(OWNER) + .build(); + } + + private static ServiceItemLeaseV1 createLease(String etag) { + ServiceItemLeaseV1 lease = new ServiceItemLeaseV1() + .withId("lease-id") + .withLeaseToken("AA-BB") + .withFeedRange(new FeedRangeEpkImpl(new Range<>("AA", "BB", true, false))) + .withContinuationToken(token("continuation-1")) + .withOwner(OWNER) + .withETag(etag); + return lease; + } + + private static InternalObjectNode leaseDocument(String id, String etag, String owner, String continuationToken) { + InternalObjectNode document = new InternalObjectNode(); + document.setId(id); + document.set("_etag", etag); + document.set("LeaseToken", "AA-BB"); + document.set("Owner", owner); + document.set("ContinuationToken", continuationToken); + document.set("version", 1); + document.set("_ts", "0"); + return document; + } + + @SuppressWarnings("unchecked") + private static CosmosItemResponse noContentResponse(String etag) { + CosmosItemResponse response = mock(CosmosItemResponse.class); + when(response.getETag()).thenReturn(etag); + return response; + } + + @SuppressWarnings("unchecked") + private static CosmosItemResponse contentResponse(InternalObjectNode document, String etag) { + ResourceResponse resourceResponse = mock(ResourceResponse.class); + when(resourceResponse.hasPayload()).thenReturn(true); + when(resourceResponse.getBody()).thenReturn(document.getPropertyBag()); + when(resourceResponse.getETag()).thenReturn(etag); + return ImplementationBridgeHelpers.CosmosItemResponseHelper.getCosmosItemResponseBuilderAccessor() + .createCosmosItemResponse(resourceResponse, Object.class, CosmosItemSerializer.DEFAULT_SERIALIZER); + } + + // Epk version leases store the continuation token as base64 encoded change feed state. + private static String token(String value) { + return Base64.getEncoder().encodeToString(value.getBytes(StandardCharsets.UTF_8)); + } + + private static CosmosException preconditionFailed() { + return BridgeInternal.createCosmosException(412); + } + + private static Mono toMono(Object result) { + if (result instanceof Throwable) { + return Mono.error((Throwable) result); + } + return Mono.just(result); + } + + private static List recordReplaces(ChangeFeedContextClient clientMock, Object... results) { + List writes = new ArrayList<>(); + doAnswer(invocation -> { + CosmosItemRequestOptions options = invocation.getArgument(3); + writes.add(new RecordedWrite(invocation.getArgument(0), options)); + return toMono(results[Math.min(writes.size(), results.length) - 1]); + }).when(clientMock).replaceItem(anyString(), any(), any(), any()); + return writes; + } + + private static List recordCreates(ChangeFeedContextClient clientMock, Object... results) { + List writes = new ArrayList<>(); + doAnswer(invocation -> { + Object document = invocation.getArgument(1); + String id = document instanceof Lease + ? ((Lease) document).getId() + : ((InternalObjectNode) document).getId(); + CosmosItemRequestOptions options = invocation.getArgument(2); + writes.add(new RecordedWrite(id, options)); + return toMono(results[Math.min(writes.size(), results.length) - 1]); + }).when(clientMock).createItem(any(), any(), any(), anyBoolean()); + return writes; + } + + @SuppressWarnings("unchecked") + private static List recordReads(ChangeFeedContextClient clientMock, InternalObjectNode document) { + List reads = new ArrayList<>(); + doAnswer(invocation -> { + reads.add(invocation.getArgument(0)); + return Mono.just(contentResponse(document, document.getETag())); + }).when(clientMock).readItem(anyString(), any(), any(), eq(InternalObjectNode.class)); + return reads; + } + + private static final class RecordedWrite { + private final String documentId; + private final String ifMatchETag; + private final boolean contentResponseOnWriteEnabled; + + private RecordedWrite(String documentId, CosmosItemRequestOptions options) { + this.documentId = documentId; + this.ifMatchETag = options.getIfMatchETag(); + Boolean enabled = options.isContentResponseOnWriteEnabled(); + this.contentResponseOnWriteEnabled = enabled == null || enabled; + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseWritesWithoutContentResponseTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseWritesWithoutContentResponseTests.java new file mode 100644 index 0000000000000..395fb2c762716 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseWritesWithoutContentResponseTests.java @@ -0,0 +1,476 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.changefeed.pkversion; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.ChangeFeedProcessor; +import com.azure.cosmos.ChangeFeedProcessorBuilder; +import com.azure.cosmos.ConsistencyLevel; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosBridgeInternal; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.InternalObjectNode; +import com.azure.cosmos.implementation.ResourceResponse; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.implementation.Document; +import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient; +import com.azure.cosmos.implementation.changefeed.Lease; +import com.azure.cosmos.implementation.changefeed.LeaseStoreManager; +import com.azure.cosmos.implementation.changefeed.common.PartitionedByIdCollectionRequestOptionsFactory; +import com.azure.cosmos.implementation.changefeed.exceptions.LeaseLostException; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.PartitionKey; +import com.fasterxml.jackson.databind.JsonNode; +import org.mockito.Mockito; +import org.testng.annotations.Test; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.function.Consumer; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Validates that change feed processor lease writes (partition key range id version) are issued with content response on write disabled + * and that the lease concurrency token (ETag) is correctly taken from the response headers. + */ +public class LeaseWritesWithoutContentResponseTests { + private static final String OWNER = "host-1"; + private static final String LEASE_PREFIX = "prefix"; + + @Test(groups = "unit") + public void updateLeaseWithoutContentResponseUsesETagFromResponseHeader() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces(clientMock, noContentResponse("etag-2")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, + cachedLease.getId(), + new PartitionKey(cachedLease.getId()), + new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken("continuation-2"); + return lease; + }).block(); + + assertThat(updatedLease).isSameAs(cachedLease); + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-2"); + assertThat(updatedLease.getContinuationToken()).isEqualTo("continuation-2"); + assertThat(updatedLease.getOwner()).isEqualTo(OWNER); + + assertThat(replaces).hasSize(1); + assertThat(replaces.get(0).ifMatchETag).isEqualTo("etag-1"); + assertThat(replaces.get(0).contentResponseOnWriteEnabled).isFalse(); + verify(clientMock, never()).readItem(anyString(), any(), any(), any()); + } + + @Test(groups = "unit") + public void sequentialUpdatesUseLatestETagAsIfMatch() { + // Regression test for the stale If-Match ETag issue - each update must use the ETag returned by the previous + // write, otherwise every subsequent checkpoint would fail with 412 and require a recovery read. + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces( + clientMock, + noContentResponse("etag-2"), + noContentResponse("etag-3"), + noContentResponse("etag-4")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + for (int i = 2; i <= 4; i++) { + String continuation = "continuation-" + i; + Lease updatedLease = updater.updateLease( + cachedLease, + cachedLease.getId(), + new PartitionKey(cachedLease.getId()), + new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken(continuation); + return lease; + }).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-" + i); + assertThat(updatedLease.getContinuationToken()).isEqualTo(continuation); + } + + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-2", "etag-3"); + assertThat(replaces).allMatch(write -> !write.contentResponseOnWriteEnabled); + verify(clientMock, never()).readItem(anyString(), any(), any(), any()); + } + + @Test(groups = "unit") + public void updateLeaseDisablesContentResponseEvenWithNullRequestOptions() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces(clientMock, noContentResponse("etag-2")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), null, lease -> lease).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-2"); + assertThat(replaces).hasSize(1); + assertThat(replaces.get(0).ifMatchETag).isEqualTo("etag-1"); + assertThat(replaces.get(0).contentResponseOnWriteEnabled).isFalse(); + } + + @Test(groups = "unit") + public void updateLeaseWithContentResponseStillUsesDocument() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + CosmosItemResponse response = contentResponse( + leaseDocument("lease-id", "etag-from-body", OWNER, "continuation-from-body"), + "etag-from-header"); + recordReplaces(clientMock, response); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-from-body"); + // Legacy behavior: setServiceItemLease keeps the cached (submitted) continuation token. + assertThat(updatedLease.getContinuationToken()).isEqualTo("continuation-1"); + } + + @Test(groups = "unit") + public void updateLeaseWithoutETagKeepsCachedTokenAndRecoversOnNextUpdate() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces( + clientMock, + noContentResponse(null), + preconditionFailed(), + noContentResponse("etag-4")); + List reads = recordReads(clientMock, leaseDocument("lease-id", "etag-3", OWNER, "continuation-2")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + Lease firstUpdate = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease).block(); + assertThat(firstUpdate.getConcurrencyToken()).isEqualTo("etag-1"); + + Lease secondUpdate = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken("continuation-3"); + return lease; + }).block(); + + assertThat(secondUpdate.getConcurrencyToken()).isEqualTo("etag-4"); + assertThat(secondUpdate.getContinuationToken()).isEqualTo("continuation-3"); + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-1", "etag-3"); + assertThat(reads).hasSize(1); + } + + @Test(groups = "unit") + public void updateLeasePreconditionFailedRefreshesETagAndRetries() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces(clientMock, preconditionFailed(), noContentResponse("etag-3")); + List reads = recordReads(clientMock, leaseDocument("lease-id", "etag-2", OWNER, "continuation-1")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + Lease updatedLease = updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> { + lease.setContinuationToken("continuation-2"); + return lease; + }).block(); + + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-3"); + assertThat(updatedLease.getContinuationToken()).isEqualTo("continuation-2"); + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-2"); + assertThat(replaces).allMatch(write -> !write.contentResponseOnWriteEnabled); + assertThat(reads).containsExactly("lease-id"); + } + + @Test(groups = "unit") + public void updateLeasePreconditionFailedWithNewOwnerThrowsLeaseLost() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordReplaces(clientMock, preconditionFailed()); + recordReads(clientMock, leaseDocument("lease-id", "etag-2", "host-2", "continuation-1")); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + StepVerifier.create(updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease)) + .expectError(LeaseLostException.class) + .verify(); + } + + @Test(groups = "unit") + public void updateLeaseConflictOrNotFoundThrowsLeaseLost() { + for (int statusCode : new int[] { 409, 404 }) { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordReplaces(clientMock, BridgeInternal.createCosmosException(statusCode)); + + DocumentServiceLeaseUpdaterImpl updater = new DocumentServiceLeaseUpdaterImpl(clientMock); + ServiceItemLease cachedLease = createLease("etag-1"); + + StepVerifier.create(updater.updateLease( + cachedLease, cachedLease.getId(), new PartitionKey(cachedLease.getId()), new CosmosItemRequestOptions(), + lease -> lease)) + .expectError(LeaseLostException.class) + .verify(); + } + } + + @Test(groups = "unit") + public void leaseStoreManagerAcquireThenUpdatePropertiesChainsETags() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List replaces = recordReplaces( + clientMock, + noContentResponse("etag-2"), + noContentResponse("etag-3")); + + LeaseStoreManager leaseStoreManager = createLeaseStoreManager(clientMock); + ServiceItemLease lease = createLease("etag-1"); + lease.setOwner("host-0"); + + Lease acquiredLease = leaseStoreManager.acquire(lease).block(); + assertThat(acquiredLease.getOwner()).isEqualTo(OWNER); + assertThat(acquiredLease.getConcurrencyToken()).isEqualTo("etag-2"); + + acquiredLease.setProperties(Collections.singletonMap("key", "value")); + Lease updatedLease = leaseStoreManager.updateProperties(acquiredLease).block(); + assertThat(updatedLease.getConcurrencyToken()).isEqualTo("etag-3"); + assertThat(updatedLease.getProperties()).containsEntry("key", "value"); + + assertThat(replaces).extracting(write -> write.ifMatchETag).containsExactly("etag-1", "etag-2"); + assertThat(replaces).allMatch(write -> !write.contentResponseOnWriteEnabled); + verify(clientMock, never()).readItem(anyString(), any(), any(), any()); + } + + @Test(groups = "unit") + public void createLeaseIfNotExistWithoutContentResponseUsesETagFromResponseHeader() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + List creates = recordCreates(clientMock, noContentResponse("etag-created")); + + LeaseStoreManager leaseStoreManager = createLeaseStoreManager(clientMock); + Lease lease = leaseStoreManager.createLeaseIfNotExist("0", "continuation-0").block(); + + assertThat(lease).isNotNull(); + assertThat(lease.getId()).isEqualTo(LEASE_PREFIX + "..0"); + assertThat(lease.getLeaseToken()).isEqualTo("0"); + assertThat(lease.getContinuationToken()).isEqualTo("continuation-0"); + assertThat(lease.getConcurrencyToken()).isEqualTo("etag-created"); + assertThat(lease.toString()).isNotNull(); + assertThat(creates).hasSize(1); + assertThat(creates.get(0).contentResponseOnWriteEnabled).isFalse(); + } + + @Test(groups = "unit") + public void createLeaseIfNotExistConflictReturnsEmpty() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordCreates(clientMock, BridgeInternal.createCosmosException(409)); + + LeaseStoreManager leaseStoreManager = createLeaseStoreManager(clientMock); + StepVerifier.create(leaseStoreManager.createLeaseIfNotExist("0", "continuation-0")) + .verifyComplete(); + } + + @Test(groups = "unit") + public void leaseStoreInitializationWritesWithoutContentResponse() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + CosmosAsyncContainer leaseContainer = mock(CosmosAsyncContainer.class); + List creates = recordCreates( + clientMock, + noContentResponse("lock-etag"), + noContentResponse("marker-etag")); + List deletes = new ArrayList<>(); + doAnswer(invocation -> { + CosmosItemRequestOptions options = invocation.getArgument(2); + deletes.add(options); + return Mono.just(noContentResponse(null)); + }).when(clientMock).deleteItem(anyString(), any(), any()); + + LeaseStoreImpl leaseStore = new LeaseStoreImpl( + clientMock, LEASE_PREFIX, leaseContainer, new PartitionedByIdCollectionRequestOptionsFactory()); + + StepVerifier.create(leaseStore.acquireInitializationLock(Duration.ofSeconds(30))) + .expectNext(true) + .verifyComplete(); + StepVerifier.create(leaseStore.markInitialized()) + .expectNext(true) + .verifyComplete(); + leaseStore.releaseInitializationLock().block(); + + assertThat(creates).hasSize(2); + assertThat(creates).allMatch(write -> !write.contentResponseOnWriteEnabled); + assertThat(creates.get(0).documentId).isEqualTo(LEASE_PREFIX + ".lock"); + assertThat(creates.get(1).documentId).isEqualTo(LEASE_PREFIX + ".info"); + assertThat(deletes).hasSize(1); + assertThat(deletes.get(0).getIfMatchETag()).isEqualTo("lock-etag"); + } + + @Test(groups = "unit") + public void leaseStoreAcquireInitializationLockConflictReturnsFalse() { + ChangeFeedContextClient clientMock = mock(ChangeFeedContextClient.class); + recordCreates(clientMock, BridgeInternal.createCosmosException(409)); + + LeaseStoreImpl leaseStore = new LeaseStoreImpl( + clientMock, LEASE_PREFIX, mock(CosmosAsyncContainer.class), new PartitionedByIdCollectionRequestOptionsFactory()); + + StepVerifier.create(leaseStore.acquireInitializationLock(Duration.ofSeconds(30))) + .expectNext(false) + .verifyComplete(); + } + + @Test(groups = "unit") + public void changeFeedProcessorAllowsLeaseClientWithContentResponseOnWriteDisabled() { + CosmosAsyncContainer feedContainer = mockContainer(true); + CosmosAsyncContainer leaseContainer = mockContainer(false); + + Consumer> handler = docs -> { }; + ChangeFeedProcessor processor = new ChangeFeedProcessorBuilder() + .hostName(OWNER) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .handleChanges(handler) + .buildChangeFeedProcessor(); + assertThat(processor).isNotNull(); + } + + static CosmosAsyncContainer mockContainer(boolean contentResponseOnWriteEnabled) { + CosmosAsyncContainer container = mock(CosmosAsyncContainer.class, Mockito.RETURNS_DEEP_STUBS); + AsyncDocumentClient documentClient = mock(AsyncDocumentClient.class); + when(documentClient.isContentResponseOnWriteEnabled()).thenReturn(contentResponseOnWriteEnabled); + when(documentClient.getConsistencyLevel()).thenReturn(ConsistencyLevel.SESSION); + when(CosmosBridgeInternal.getContextClient(container)).thenReturn(documentClient); + return container; + } + + private static LeaseStoreManager createLeaseStoreManager(ChangeFeedContextClient clientMock) { + return LeaseStoreManagerImpl.builder() + .leasePrefix(LEASE_PREFIX) + .leaseCollectionLink(mock(CosmosAsyncContainer.class)) + .leaseContextClient(clientMock) + .requestOptionsFactory(new PartitionedByIdCollectionRequestOptionsFactory()) + .hostName(OWNER) + .build(); + } + + private static ServiceItemLease createLease(String etag) { + ServiceItemLease lease = new ServiceItemLease() + .withId("lease-id") + .withLeaseToken("0") + .withContinuationToken("continuation-1") + .withOwner(OWNER) + .withETag(etag); + return lease; + } + + private static InternalObjectNode leaseDocument(String id, String etag, String owner, String continuationToken) { + InternalObjectNode document = new InternalObjectNode(); + document.setId(id); + document.set("_etag", etag); + document.set("LeaseToken", "0"); + document.set("Owner", owner); + document.set("ContinuationToken", continuationToken); + document.set("_ts", "0"); + return document; + } + + @SuppressWarnings("unchecked") + private static CosmosItemResponse noContentResponse(String etag) { + CosmosItemResponse response = mock(CosmosItemResponse.class); + when(response.getETag()).thenReturn(etag); + return response; + } + + @SuppressWarnings("unchecked") + private static CosmosItemResponse contentResponse(InternalObjectNode document, String etag) { + ResourceResponse resourceResponse = mock(ResourceResponse.class); + when(resourceResponse.hasPayload()).thenReturn(true); + when(resourceResponse.getBody()).thenReturn(document.getPropertyBag()); + when(resourceResponse.getETag()).thenReturn(etag); + return ImplementationBridgeHelpers.CosmosItemResponseHelper.getCosmosItemResponseBuilderAccessor() + .createCosmosItemResponse(resourceResponse, Object.class, CosmosItemSerializer.DEFAULT_SERIALIZER); + } + + private static CosmosException preconditionFailed() { + return BridgeInternal.createCosmosException(412); + } + + private static Mono toMono(Object result) { + if (result instanceof Throwable) { + return Mono.error((Throwable) result); + } + return Mono.just(result); + } + + private static List recordReplaces(ChangeFeedContextClient clientMock, Object... results) { + List writes = new ArrayList<>(); + doAnswer(invocation -> { + CosmosItemRequestOptions options = invocation.getArgument(3); + writes.add(new RecordedWrite(invocation.getArgument(0), options)); + return toMono(results[Math.min(writes.size(), results.length) - 1]); + }).when(clientMock).replaceItem(anyString(), any(), any(), any()); + return writes; + } + + private static List recordCreates(ChangeFeedContextClient clientMock, Object... results) { + List writes = new ArrayList<>(); + doAnswer(invocation -> { + Object document = invocation.getArgument(1); + String id = document instanceof Lease + ? ((Lease) document).getId() + : ((InternalObjectNode) document).getId(); + CosmosItemRequestOptions options = invocation.getArgument(2); + writes.add(new RecordedWrite(id, options)); + return toMono(results[Math.min(writes.size(), results.length) - 1]); + }).when(clientMock).createItem(any(), any(), any(), anyBoolean()); + return writes; + } + + @SuppressWarnings("unchecked") + private static List recordReads(ChangeFeedContextClient clientMock, InternalObjectNode document) { + List reads = new ArrayList<>(); + doAnswer(invocation -> { + reads.add(invocation.getArgument(0)); + return Mono.just(contentResponse(document, document.getETag())); + }).when(clientMock).readItem(anyString(), any(), any(), eq(InternalObjectNode.class)); + return reads; + } + + private static final class RecordedWrite { + private final String documentId; + private final String ifMatchETag; + private final boolean contentResponseOnWriteEnabled; + + private RecordedWrite(String documentId, CosmosItemRequestOptions options) { + this.documentId = documentId; + this.ifMatchETag = options.getIfMatchETag(); + Boolean enabled = options.isContentResponseOnWriteEnabled(); + this.contentResponseOnWriteEnabled = enabled == null || enabled; + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/changefeed/ChangeFeedProcessorLeaseWithoutContentResponseTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/changefeed/ChangeFeedProcessorLeaseWithoutContentResponseTest.java new file mode 100644 index 0000000000000..51df495115cf5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/changefeed/ChangeFeedProcessorLeaseWithoutContentResponseTest.java @@ -0,0 +1,338 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.rx.changefeed; + +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.CosmosDiagnosticsContext; +import com.azure.cosmos.implementation.InternalObjectNode; +import com.azure.cosmos.models.ChangeFeedProcessorItem; +import com.azure.cosmos.models.ChangeFeedProcessorOptions; +import com.azure.cosmos.models.CosmosClientTelemetryConfig; +import com.azure.cosmos.models.CosmosContainerProperties; +import com.azure.cosmos.models.CosmosContainerRequestOptions; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.rx.TestSuiteBase; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Consumer; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Validates that the change feed processor works end to end when lease writes do not return response content - both + * when the lease container client has content response on write disabled (which used to be rejected) and with the + * change feed processor now always disabling content response for its own lease writes. + *

+ * Each lease is checkpointed many times sequentially, so any stale If-Match ETag kept on the cached lease after a + * write (the server's new ETag is only available in the response headers) would surface as 412 responses on the + * lease container. + */ +public class ChangeFeedProcessorLeaseWithoutContentResponseTest extends TestSuiteBase { + private static final int DOCUMENTS_PER_WAVE = 30; + private static final int MAX_ITEM_COUNT = 3; + private static final int FEED_THROUGHPUT = 400; + private static final int LEASE_THROUGHPUT = 400; + private static final Duration WAIT_TIMEOUT = Duration.ofMinutes(2); + private static final Duration LIFECYCLE_TIMEOUT = Duration.ofSeconds(60); + private static final long CHECKPOINT_SETTLE_TIME_MS = 3000; + private static final int TEST_TIMEOUT = 10 * 60 * 1000; + + private enum Mode { + PK_RANGE_ID_VERSION, + LATEST_VERSION, + ALL_VERSIONS_AND_DELETES + } + + private final Map> leaseOperationsByContainer = new ConcurrentHashMap<>(); + + private CosmosAsyncClient client; + private CosmosAsyncClient leaseClientWithoutContentResponse; + private CosmosAsyncDatabase createdDatabase; + + @Factory(dataProvider = "clientBuildersWithSessionConsistency") + public ChangeFeedProcessorLeaseWithoutContentResponseTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @DataProvider(name = "incrementalModes") + public static Object[][] incrementalModes() { + return new Object[][] { + { Mode.PK_RANGE_ID_VERSION }, + { Mode.LATEST_VERSION } + }; + } + + @BeforeClass(groups = { "query", "long-emulator" }, timeOut = SETUP_TIMEOUT) + public void before_ChangeFeedProcessorLeaseWithoutContentResponseTest() { + client = getClientBuilder().buildAsyncClient(); + createdDatabase = getSharedCosmosDatabase(client); + + CosmosClientTelemetryConfig telemetryConfig = new CosmosClientTelemetryConfig() + .diagnosticsHandler((diagnosticsContext, traceContext) -> recordLeaseOperation(diagnosticsContext)); + + leaseClientWithoutContentResponse = copyCosmosClientBuilder(getClientBuilder()) + .contentResponseOnWriteEnabled(false) + .clientTelemetryConfig(telemetryConfig) + .buildAsyncClient(); + } + + @AfterClass(groups = { "query", "long-emulator" }, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + safeClose(leaseClientWithoutContentResponse); + safeClose(client); + } + + @Test(groups = { "query" }, dataProvider = "incrementalModes", timeOut = TEST_TIMEOUT) + public void incrementalLeaseWritesWithoutContentResponse(Mode mode) throws InterruptedException { + runScenario(mode); + } + + @Test(groups = { "long-emulator" }, timeOut = TEST_TIMEOUT) + public void allVersionsAndDeletesLeaseWritesWithoutContentResponse() throws InterruptedException { + runScenario(Mode.ALL_VERSIONS_AND_DELETES); + } + + private void runScenario(Mode mode) throws InterruptedException { + CosmosContainerProperties feedDefinition = mode == Mode.ALL_VERSIONS_AND_DELETES + ? getCollectionDefinitionWithFullFidelity() + : getCollectionDefinition(); + CosmosAsyncContainer feedContainer = createCollection( + createdDatabase, feedDefinition, new CosmosContainerRequestOptions(), FEED_THROUGHPUT); + CosmosAsyncContainer leaseContainerWithContent = createCollection( + createdDatabase, + new CosmosContainerProperties("leases_" + UUID.randomUUID(), "/id"), + new CosmosContainerRequestOptions(), + LEASE_THROUGHPUT); + CosmosAsyncContainer leaseContainer = leaseClientWithoutContentResponse + .getDatabase(createdDatabase.getId()) + .getContainer(leaseContainerWithContent.getId()); + ConcurrentLinkedQueue leaseOperations = new ConcurrentLinkedQueue<>(); + leaseOperationsByContainer.put(leaseContainer.getId(), leaseOperations); + + String leasePrefix = "noContent" + mode.ordinal(); + Map receivedDocuments = new ConcurrentHashMap<>(); + List createdDocumentIds = new ArrayList<>(); + ChangeFeedProcessor firstProcessor = null; + ChangeFeedProcessor secondProcessor = null; + + try { + firstProcessor = buildProcessor(mode, "host-1", leasePrefix, feedContainer, leaseContainer, receivedDocuments); + start(firstProcessor); + waitForOwnedLeases(leaseContainerWithContent, "host-1"); + if (mode == Mode.ALL_VERSIONS_AND_DELETES) { + // AllVersionsAndDeletes always starts from "now" - give the processor time to issue its first read. + Thread.sleep(5000); + } + + // Two waves, each producing many sequential checkpoints on the same lease(s). + createdDocumentIds.addAll(createDocuments(feedContainer, DOCUMENTS_PER_WAVE)); + waitForDocuments(receivedDocuments, createdDocumentIds); + createdDocumentIds.addAll(createDocuments(feedContainer, DOCUMENTS_PER_WAVE)); + waitForDocuments(receivedDocuments, createdDocumentIds); + // Change feed processing is at-least-once - let the checkpoint for the last batch complete before stopping. + Thread.sleep(CHECKPOINT_SETTLE_TIME_MS); + + // Snapshot the lease operations while the processor still owns the leases. + List operationsWhileRunning = new ArrayList<>(leaseOperations); + stop(firstProcessor); + + assertThat(operationsWhileRunning.stream().filter(op -> op.statusCode == 412).collect(Collectors.toList())) + .as("No lease write should fail with 412 - the cached lease must carry the server's latest ETag") + .isEmpty(); + long successfulReplaces = operationsWhileRunning.stream() + .filter(op -> "Replace".equalsIgnoreCase(op.operationType) && op.statusCode == 200) + .count(); + int minimumCheckpoints = (2 * DOCUMENTS_PER_WAVE) / MAX_ITEM_COUNT; + assertThat(successfulReplaces) + .as("Each lease should have been checkpointed many times sequentially") + .isGreaterThanOrEqualTo(minimumCheckpoints); + assertThat(operationsWhileRunning.stream().anyMatch(op -> "Create".equalsIgnoreCase(op.operationType) && op.statusCode == 201)) + .as("Lease store bootstrap (lock, leases, marker) should have created documents") + .isTrue(); + + assertExactlyOnce(receivedDocuments, createdDocumentIds); + List leases = readLeases(leaseContainerWithContent); + assertThat(leases).isNotEmpty(); + for (JsonNode lease : leases) { + assertThat(lease.path("ContinuationToken").asText(null)) + .as("Lease %s should have a persisted continuation token", lease.path("id").asText()) + .isNotEmpty(); + } + + // A new processor instance must resume from the persisted continuation tokens: nothing is re-delivered + // and only new documents are processed. + secondProcessor = buildProcessor(mode, "host-2", leasePrefix, feedContainer, leaseContainer, receivedDocuments); + start(secondProcessor); + waitForOwnedLeases(leaseContainerWithContent, "host-2"); + createdDocumentIds.addAll(createDocuments(feedContainer, 5)); + waitForDocuments(receivedDocuments, createdDocumentIds); + Thread.sleep(CHECKPOINT_SETTLE_TIME_MS); + stop(secondProcessor); + + assertExactlyOnce(receivedDocuments, createdDocumentIds); + } finally { + stop(firstProcessor); + stop(secondProcessor); + leaseOperationsByContainer.remove(leaseContainer.getId()); + safeDeleteCollection(feedContainer); + safeDeleteCollection(leaseContainerWithContent); + } + } + + private ChangeFeedProcessor buildProcessor( + Mode mode, + String hostName, + String leasePrefix, + CosmosAsyncContainer feedContainer, + CosmosAsyncContainer leaseContainer, + Map receivedDocuments) { + + // Long renew/expiration intervals keep the renewer from racing with checkpoints on the same lease, so the + // only lease writes during the test are deterministic: bootstrap, acquire, sequential checkpoints, release. + ChangeFeedProcessorOptions options = new ChangeFeedProcessorOptions() + .setLeasePrefix(leasePrefix) + .setMaxItemCount(MAX_ITEM_COUNT) + .setFeedPollDelay(Duration.ofMillis(500)) + .setLeaseAcquireInterval(Duration.ofSeconds(2)) + .setLeaseRenewInterval(Duration.ofMinutes(5)) + .setLeaseExpirationInterval(Duration.ofMinutes(10)); + if (mode != Mode.ALL_VERSIONS_AND_DELETES) { + options.setStartFromBeginning(true); + } + + ChangeFeedProcessorBuilder builder = new ChangeFeedProcessorBuilder() + .hostName(hostName) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .options(options); + + Consumer onDocument = id -> receivedDocuments.computeIfAbsent(id, key -> new AtomicInteger()).incrementAndGet(); + switch (mode) { + case PK_RANGE_ID_VERSION: + builder.handleChanges((List docs) -> docs.forEach(doc -> onDocument.accept(doc.get("id").asText()))); + break; + case LATEST_VERSION: + builder.handleLatestVersionChanges((List items) -> + items.forEach(item -> onDocument.accept(item.getCurrent().get("id").asText()))); + break; + case ALL_VERSIONS_AND_DELETES: + builder.handleAllVersionsAndDeletesChanges((List items) -> + items.forEach(item -> onDocument.accept(item.getCurrent().get("id").asText()))); + break; + default: + throw new IllegalArgumentException("Unexpected mode " + mode); + } + + return builder.buildChangeFeedProcessor(); + } + + private void recordLeaseOperation(CosmosDiagnosticsContext diagnosticsContext) { + ConcurrentLinkedQueue operations = leaseOperationsByContainer.get(diagnosticsContext.getContainerName()); + if (operations != null && "Document".equalsIgnoreCase(diagnosticsContext.getResourceType())) { + operations.add(new LeaseOperation(diagnosticsContext.getOperationType(), diagnosticsContext.getStatusCode())); + } + } + + private List createDocuments(CosmosAsyncContainer feedContainer, int count) { + List documents = new ArrayList<>(); + for (int i = 0; i < count; i++) { + String id = UUID.randomUUID().toString(); + documents.add(new InternalObjectNode(String.format("{ \"id\": \"%s\", \"mypk\": \"%s\" }", id, id))); + } + + List created = insertAllItemsBlocking(feedContainer, documents, false); + waitIfNeededForReplicasToCatchUp(getClientBuilder()); + return created.stream().map(InternalObjectNode::getId).collect(Collectors.toList()); + } + + private static void waitForDocuments(Map receivedDocuments, List expectedIds) + throws InterruptedException { + + Instant deadline = Instant.now().plus(WAIT_TIMEOUT); + while (!receivedDocuments.keySet().containsAll(expectedIds) && Instant.now().isBefore(deadline)) { + Thread.sleep(250); + } + + assertThat(receivedDocuments.keySet()).containsAll(expectedIds); + } + + private static void waitForOwnedLeases(CosmosAsyncContainer leaseContainer, String hostName) throws InterruptedException { + Instant deadline = Instant.now().plus(WAIT_TIMEOUT); + while (Instant.now().isBefore(deadline)) { + List leases = readLeases(leaseContainer); + if (!leases.isEmpty() && leases.stream().allMatch(lease -> hostName.equals(lease.path("Owner").asText(null)))) { + return; + } + Thread.sleep(500); + } + + throw new AssertionError("Leases were not acquired by " + hostName + " within " + WAIT_TIMEOUT); + } + + private static List readLeases(CosmosAsyncContainer leaseContainer) { + return leaseContainer + .queryItems("SELECT * FROM c WHERE IS_DEFINED(c.LeaseToken)", new CosmosQueryRequestOptions(), ObjectNode.class) + .collectList() + .block() + .stream() + .map(JsonNode.class::cast) + .collect(Collectors.toList()); + } + + private static void assertExactlyOnce(Map receivedDocuments, List expectedIds) { + assertThat(receivedDocuments.keySet()).containsExactlyInAnyOrderElementsOf(expectedIds); + for (Map.Entry entry : receivedDocuments.entrySet()) { + assertThat(entry.getValue().get()) + .as("Document %s should be delivered exactly once (no lease loss / continuation regression)", entry.getKey()) + .isEqualTo(1); + } + } + + private static void start(ChangeFeedProcessor processor) { + processor.start().subscribeOn(Schedulers.boundedElastic()).timeout(LIFECYCLE_TIMEOUT).block(); + } + + private static void stop(ChangeFeedProcessor processor) { + if (processor != null && processor.isStarted()) { + processor.stop().timeout(LIFECYCLE_TIMEOUT).onErrorResume(throwable -> Mono.empty()).block(); + } + } + + private static final class LeaseOperation { + private final String operationType; + private final int statusCode; + + private LeaseOperation(String operationType, int statusCode) { + this.operationType = operationType; + this.statusCode = statusCode; + } + + @Override + public String toString() { + return operationType + ":" + statusCode; + } + } +} diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index c50751f8e4960..d302dead78c98 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -15,6 +15,7 @@ * Updated Netty to `4.2.19.Final`, Reactor Netty to `1.3.7`, and Reactor Core to `3.8.7`. Migrated direct TCP event loops to the Netty 4.2 I/O handler API while preserving pooled buffers and explicit TLS validation settings. * Added a compact `ppaf` bookmark to each data-plane attempt in `CosmosDiagnostics`, containing the current per-partition write region, failed regions, and the time it was designated, or an empty object when no override is active. +* Change Feed Processor lease writes no longer request the document in the response, reducing response bandwidth and client response-processing overhead without changing RU charges. The new lease ETag is taken from the response headers, and the lease container's client no longer needs `contentResponseOnWriteEnabled(true)`. - See [PR 50697](https://github.com/Azure/azure-sdk-for-java/pull/50697) ### 4.83.0 (2026-09-22) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedHelper.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedHelper.java index 542440c7d99f5..5edf7b0f0cd4c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedHelper.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/common/ChangeFeedHelper.java @@ -2,6 +2,8 @@ // Licensed under the MIT License. package com.azure.cosmos.implementation.changefeed.common; +import com.azure.cosmos.models.CosmosItemRequestOptions; + import java.net.MalformedURLException; import java.net.URI; import java.net.URISyntaxException; @@ -24,6 +26,20 @@ public class ChangeFeedHelper { public static final int HTTP_STATUS_CODE_TOO_MANY_REQUESTS = 429; public static final int HTTP_STATUS_CODE_INTERNAL_SERVER_ERROR = 500; + /** + * Disables content response on write for a lease container write request. Lease state is preserved locally + * and only the new ETag (concurrency token) is taken from the response headers, so the response body is + * never needed. This overrides the content-response-on-write setting of the lease client. + * + * @param options the request options to update; a new instance is created when null. + * @return the request options with content response on write disabled. + */ + public static CosmosItemRequestOptions withContentResponseOnWriteDisabled(CosmosItemRequestOptions options) { + CosmosItemRequestOptions effectiveOptions = options != null ? options : new CosmosItemRequestOptions(); + effectiveOptions.setContentResponseOnWriteEnabled(false); + return effectiveOptions; + } + public static String getDatabaseLink(String databaseName) { return String.format("/dbs/%s", databaseName); } 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 c726d3d2cde67..98ee2871376bf 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 @@ -145,10 +145,6 @@ private void validateChangeFeedProcessorOptions(ChangeFeedProcessorOptions chang private void validateLeaseContainer(CosmosAsyncContainer leaseContainer) { checkNotNull(leaseContainer, "Argument 'leaseContainer' can not be null"); - if (!getContextClient(leaseContainer).isContentResponseOnWriteEnabled()) { - throw new IllegalArgumentException("leaseClient: content response on write setting must be enabled"); - } - ConsistencyLevel consistencyLevel = getContextClient(leaseContainer).getConsistencyLevel(); if (consistencyLevel == ConsistencyLevel.CONSISTENT_PREFIX || consistencyLevel == ConsistencyLevel.EVENTUAL) { logger.warn("leaseClient consistency level setting are less then expected which is SESSION"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/DocumentServiceLeaseUpdaterImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/DocumentServiceLeaseUpdaterImpl.java index a05af0e0c4207..843f0b56c767a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/DocumentServiceLeaseUpdaterImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/DocumentServiceLeaseUpdaterImpl.java @@ -11,7 +11,10 @@ import com.azure.cosmos.implementation.changefeed.ServiceItemLeaseUpdater; import com.azure.cosmos.implementation.changefeed.exceptions.LeaseConflictException; import com.azure.cosmos.implementation.changefeed.exceptions.LeaseLostException; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedHelper; import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.PartitionKey; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -61,10 +64,7 @@ public Mono updateLease( return Mono.just(this) .flatMap( value -> this.tryReplaceLease(cachedLease, itemId, partitionKey, requestOptions)) - .map(leaseDocument -> { - cachedLease.setServiceItemLease(ServiceItemLeaseV1.fromDocument(leaseDocument)); - return cachedLease; - }) + .map(replaceResponse -> this.applyReplaceResponse(cachedLease, replaceResponse)) .hasElement() .flatMap(hasItems -> { if (hasItems) { @@ -134,7 +134,29 @@ public Mono updateLease( }); } - private Mono tryReplaceLease( + private Lease applyReplaceResponse(Lease cachedLease, CosmosItemResponse replaceResponse) { + InternalObjectNode leaseDocument = BridgeInternal.getProperties(replaceResponse); + if (leaseDocument != null) { + cachedLease.setServiceItemLease(ServiceItemLeaseV1.fromDocument(leaseDocument)); + return cachedLease; + } + + // Lease writes are issued with content response on write disabled - the submitted lease state is what got + // persisted, so only the new concurrency token needs to be taken from the response headers. Not doing so + // would leave a stale If-Match ETag on the cached lease and cause the next update to fail with 412. + String newETag = replaceResponse.getETag(); + if (StringUtils.isEmpty(newETag)) { + logger.warn( + "Lease with token {}: Replace succeeded but the response has no ETag; the next update will refresh the lease.", + cachedLease.getLeaseToken()); + } else { + cachedLease.setConcurrencyToken(newETag); + } + + return cachedLease; + } + + private Mono> tryReplaceLease( Lease lease, String itemId, PartitionKey partitionKey, @@ -144,7 +166,6 @@ private Mono tryReplaceLease( partitionKey, lease, this.getCreateIfMatchOptions(cosmosItemRequestOptions, lease)) - .map(cosmosItemResponse -> BridgeInternal.getProperties(cosmosItemResponse)) .onErrorResume(re -> { if (re instanceof CosmosException) { CosmosException ex = (CosmosException) re; @@ -168,8 +189,10 @@ private Mono tryReplaceLease( } private CosmosItemRequestOptions getCreateIfMatchOptions(CosmosItemRequestOptions createIfMatchOptions, Lease lease) { - createIfMatchOptions.setIfMatchETag(lease.getConcurrencyToken()); + CosmosItemRequestOptions effectiveOptions = + ChangeFeedHelper.withContentResponseOnWriteDisabled(createIfMatchOptions); + effectiveOptions.setIfMatchETag(lease.getConcurrencyToken()); - return createIfMatchOptions; + return effectiveOptions; } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreImpl.java index 05041099f6c9c..c4cb1210c3729 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/epkversion/LeaseStoreImpl.java @@ -77,7 +77,8 @@ public Mono markInitialized() { return this.client.createItem( this.leaseCollectionLink, containerDocument, - this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument)), + ChangeFeedHelper.withContentResponseOnWriteDisabled( + this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument))), false) .map( item -> { return true; @@ -108,15 +109,14 @@ public Mono acquireInitializationLock(Duration lockExpirationTime) { return this.client.createItem( this.leaseCollectionLink, containerDocument, - this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument)), + ChangeFeedHelper.withContentResponseOnWriteDisabled( + this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument))), false) .map(documentResourceResponse -> { - if (BridgeInternal.getProperties(documentResourceResponse) != null) { - this.lockETag = BridgeInternal.getProperties(documentResourceResponse).getETag(); - return true; - } else { - return false; - } + // The lock is created with content response on write disabled, so the ETag comes from the + // response headers - a successful create always means the lock was acquired. + this.lockETag = documentResourceResponse.getETag(); + return true; }) .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { 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 b636b27b20c71..28b542985577e 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 @@ -180,7 +180,8 @@ public Mono createLeaseIfNotExist(FeedRangeEpkImpl feedRange, String cont return this.leaseDocumentClient.createItem( this.settings.getLeaseCollectionLink(), documentServiceLease, - this.requestOptionsFactory.createItemRequestOptions(documentServiceLease), + ChangeFeedHelper.withContentResponseOnWriteDisabled( + this.requestOptionsFactory.createItemRequestOptions(documentServiceLease)), false) .onErrorResume( ex -> { if (ex instanceof CosmosException) { @@ -204,11 +205,16 @@ public Mono createLeaseIfNotExist(FeedRangeEpkImpl feedRange, String cont } InternalObjectNode document = BridgeInternal.getProperties(documentResourceResponse); + if (document != null) { + return documentServiceLease + .withId(document.getId()) + .withETag(document.getETag()) + .withTs(document.getString(Constants.Properties.LAST_MODIFIED)); + } - return documentServiceLease - .withId(document.getId()) - .withETag(document.getETag()) - .withTs(document.getString(Constants.Properties.LAST_MODIFIED)); + // Lease writes are issued with content response on write disabled - the submitted lease is what got + // persisted, so only the concurrency token needs to be taken from the response headers. + return documentServiceLease.withETag(documentResourceResponse.getETag()); }); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/DocumentServiceLeaseUpdaterImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/DocumentServiceLeaseUpdaterImpl.java index 8b10662392be6..86c7e135ac681 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/DocumentServiceLeaseUpdaterImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/DocumentServiceLeaseUpdaterImpl.java @@ -5,7 +5,10 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.InternalObjectNode; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.changefeed.common.ChangeFeedHelper; import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.implementation.changefeed.ChangeFeedContextClient; import com.azure.cosmos.implementation.changefeed.Lease; @@ -56,10 +59,7 @@ public Mono updateLease(final Lease cachedLease, String itemId, Partition return Mono.just(this) .flatMap( value -> this.tryReplaceLease(cachedLease, itemId, partitionKey, requestOptions)) - .map(leaseDocument -> { - cachedLease.setServiceItemLease(ServiceItemLease.fromDocument(leaseDocument)); - return cachedLease; - }) + .map(replaceResponse -> this.applyReplaceResponse(cachedLease, replaceResponse)) .hasElement() .flatMap(hasItems -> { if (hasItems) { @@ -127,13 +127,34 @@ public Mono updateLease(final Lease cachedLease, String itemId, Partition }); } - private Mono tryReplaceLease( + private Lease applyReplaceResponse(Lease cachedLease, CosmosItemResponse replaceResponse) { + InternalObjectNode leaseDocument = BridgeInternal.getProperties(replaceResponse); + if (leaseDocument != null) { + cachedLease.setServiceItemLease(ServiceItemLease.fromDocument(leaseDocument)); + return cachedLease; + } + + // Lease writes are issued with content response on write disabled - the submitted lease state is what got + // persisted, so only the new concurrency token needs to be taken from the response headers. Not doing so + // would leave a stale If-Match ETag on the cached lease and cause the next update to fail with 412. + String newETag = replaceResponse.getETag(); + if (StringUtils.isEmpty(newETag)) { + logger.warn( + "Lease with token {}: Replace succeeded but the response has no ETag; the next update will refresh the lease.", + cachedLease.getLeaseToken()); + } else { + cachedLease.setConcurrencyToken(newETag); + } + + return cachedLease; + } + + private Mono> tryReplaceLease( Lease lease, String itemId, PartitionKey partitionKey, CosmosItemRequestOptions cosmosItemRequestOptions) throws LeaseLostException { return this.client.replaceItem(itemId, partitionKey, lease, this.getCreateIfMatchOptions(cosmosItemRequestOptions, lease)) - .map(cosmosItemResponse -> BridgeInternal.getProperties(cosmosItemResponse)) .onErrorResume(re -> { if (re instanceof CosmosException) { CosmosException ex = (CosmosException) re; @@ -157,7 +178,10 @@ private Mono tryReplaceLease( } private CosmosItemRequestOptions getCreateIfMatchOptions(CosmosItemRequestOptions createIfMatchOptions, Lease lease) { - createIfMatchOptions.setIfMatchETag(lease.getConcurrencyToken()); - return createIfMatchOptions; + CosmosItemRequestOptions effectiveOptions = + ChangeFeedHelper.withContentResponseOnWriteDisabled(createIfMatchOptions); + effectiveOptions.setIfMatchETag(lease.getConcurrencyToken()); + + return effectiveOptions; } } 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 930a8b44bfb1b..4096fda6dbc08 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 @@ -135,10 +135,6 @@ private void validateChangeFeedProcessorOptions(ChangeFeedProcessorOptions chang private void validateLeaseContainer(CosmosAsyncContainer leaseContainer) { checkNotNull(leaseContainer, "Argument 'leaseContainer' can not be null"); - if (!getContextClient(leaseContainer).isContentResponseOnWriteEnabled()) { - throw new IllegalArgumentException("leaseClient: content response on write setting must be enabled"); - } - ConsistencyLevel consistencyLevel = getContextClient(leaseContainer).getConsistencyLevel(); if (consistencyLevel == ConsistencyLevel.CONSISTENT_PREFIX || consistencyLevel == ConsistencyLevel.EVENTUAL) { logger.warn("leaseClient consistency level setting are less then expected which is SESSION"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreImpl.java index 084302074e620..2b2f55befa788 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/changefeed/pkversion/LeaseStoreImpl.java @@ -77,7 +77,8 @@ public Mono markInitialized() { return this.client.createItem( this.leaseCollectionLink, containerDocument, - this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument)), + ChangeFeedHelper.withContentResponseOnWriteDisabled( + this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument))), false) .map( item -> true) .onErrorResume(throwable -> { @@ -106,15 +107,14 @@ public Mono acquireInitializationLock(Duration lockExpirationTime) { return this.client.createItem( this.leaseCollectionLink, containerDocument, - this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument)), + ChangeFeedHelper.withContentResponseOnWriteDisabled( + this.requestOptionsFactory.createItemRequestOptions(ServiceItemLeaseV1.fromDocument(containerDocument))), false) .map(documentResourceResponse -> { - if (BridgeInternal.getProperties(documentResourceResponse) != null) { - this.lockETag = BridgeInternal.getProperties(documentResourceResponse).getETag(); - return true; - } else { - return false; - } + // The lock is created with content response on write disabled, so the ETag comes from the + // response headers - a successful create always means the lock was acquired. + this.lockETag = documentResourceResponse.getETag(); + return true; }) .onErrorResume(throwable -> { if (throwable instanceof CosmosException) { 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 642d7e578e44a..86f357c8bc1c0 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 @@ -176,7 +176,8 @@ public Mono createLeaseIfNotExist(String leaseToken, String continuationT return this.leaseDocumentClient.createItem( this.settings.getLeaseCollectionLink(), documentServiceLease, - this.requestOptionsFactory.createItemRequestOptions(documentServiceLease), + ChangeFeedHelper.withContentResponseOnWriteDisabled( + this.requestOptionsFactory.createItemRequestOptions(documentServiceLease)), false) .onErrorResume( ex -> { if (ex instanceof CosmosException) { @@ -200,11 +201,16 @@ public Mono createLeaseIfNotExist(String leaseToken, String continuationT } InternalObjectNode document = BridgeInternal.getProperties(documentResourceResponse); + if (document != null) { + return documentServiceLease + .withId(document.getId()) + .withETag(document.getETag()) + .withTs(document.getString(Constants.Properties.LAST_MODIFIED)); + } - return documentServiceLease - .withId(document.getId()) - .withETag(document.getETag()) - .withTs(document.getString(Constants.Properties.LAST_MODIFIED)); + // Lease writes are issued with content response on write disabled - the submitted lease is what got + // persisted, so only the concurrency token needs to be taken from the response headers. + return documentServiceLease.withETag(documentResourceResponse.getETag()); }); }