Skip to content

Enable configurable multi-lease acquisition per cycle in change feed processor #47606

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stas-openai wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Enable configurable multi-lease acquisition per cycle in change feed processor #47606

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
95a3b8e
introduce MaxLeasesToAcquirePerCycle
stas-openai Dec 28, 2025
0ca5605
address copilot's comments (restructure code + only shuffle when
stas-openai Dec 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 121 additions & 0 deletions .../azure/cosmos/implementation/changefeed/common/EqualPartitionsBalancingStrategyTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.cosmos.implementation.changefeed.common;

import com.azure.cosmos.implementation.changefeed.Lease;
import com.azure.cosmos.implementation.changefeed.epkversion.ServiceItemLeaseV1;
import org.testng.annotations.Test;

import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;

public class EqualPartitionsBalancingStrategyTests {

@Test(groups = "unit")
public void expiredLeases_legacyClampsToOneWhenMultipleWorkers() {
String hostName = "me";
List<Lease> allLeases = new ArrayList<>();

// Multiple workers exist because there are existing owners in the lease set.
for (int i = 0; i < 10; i++) {
allLeases.add(newLease("unowned-" + i, null));
allLeases.add(newLease("old1-" + i, "old1", Instant.now()));
allLeases.add(newLease("old2-" + i, "old2", Instant.now()));
}

EqualPartitionsBalancingStrategy strategy =
new EqualPartitionsBalancingStrategy(hostName, 0, 0, Duration.ofSeconds(60));

List<Lease> leasesToTake = strategy.selectLeasesToTake(allLeases);
assertEquals(leasesToTake.size(), 1);
}

@Test(groups = "unit")
public void expiredLeases_allowsMultipleWhenConfigured() {
String hostName = "me";
List<Lease> allLeases = new ArrayList<>();

for (int i = 0; i < 10; i++) {
allLeases.add(newLease("unowned-" + i, null));
allLeases.add(newLease("old1-" + i, "old1", Instant.now()));
allLeases.add(newLease("old2-" + i, "old2", Instant.now()));
}

EqualPartitionsBalancingStrategy strategy =
new EqualPartitionsBalancingStrategy(hostName, 0, 0, Duration.ofSeconds(60), 5);

List<Lease> leasesToTake = strategy.selectLeasesToTake(allLeases);
assertEquals(leasesToTake.size(), 5);
assertUniqueLeaseTokens(leasesToTake);
}

@Test(groups = "unit")
public void stealLeases_legacyClampsToOne() {
String hostName = "me";
List<Lease> allLeases = new ArrayList<>();

// No expired leases: everything is owned by a single other worker.
for (int i = 0; i < 30; i++) {
allLeases.add(newLease("old-" + i, "old", Instant.now()));
}

EqualPartitionsBalancingStrategy strategy =
new EqualPartitionsBalancingStrategy(hostName, 0, 0, Duration.ofSeconds(60));

List<Lease> leasesToTake = strategy.selectLeasesToTake(allLeases);
assertEquals(leasesToTake.size(), 1);
assertEquals(leasesToTake.get(0).getOwner(), "old");
}

@Test(groups = "unit")
public void stealLeases_stillClampsToOneWhenConfigured() {
String hostName = "me";
List<Lease> allLeases = new ArrayList<>();

for (int i = 0; i < 30; i++) {
allLeases.add(newLease("old-" + i, "old", Instant.now()));
}

EqualPartitionsBalancingStrategy strategy =
new EqualPartitionsBalancingStrategy(hostName, 0, 0, Duration.ofSeconds(60), 5);

List<Lease> leasesToTake = strategy.selectLeasesToTake(allLeases);
// Multi-acquire is only for unused/expired leases; stealing intentionally keeps the legacy 1-lease-per-cycle behavior.
assertEquals(leasesToTake.size(), 1);
assertUniqueLeaseTokens(leasesToTake);

for (Lease lease : leasesToTake) {
assertEquals(lease.getOwner(), "old");
}
}

private static ServiceItemLeaseV1 newLease(String token, String owner) {
return newLease(token, owner, null);
}

private static ServiceItemLeaseV1 newLease(String token, String owner, Instant timestamp) {
ServiceItemLeaseV1 lease = new ServiceItemLeaseV1()
.withLeaseToken(token)
.withOwner(owner);
if (timestamp != null) {
lease.withTimestamp(timestamp);
}
lease.setId("lease-" + token);
return lease;
}

private static void assertUniqueLeaseTokens(List<Lease> leases) {
Set<String> tokens = new HashSet<>();
for (Lease lease : leases) {
assertTrue(tokens.add(lease.getLeaseToken()), "Duplicate lease token: " + lease.getLeaseToken());
}
}
}
54 changes: 54 additions & 0 deletions ...com/azure/cosmos/implementation/changefeed/epkversion/PartitionLoadBalancerImplTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import java.time.Duration;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.UUID;

Expand Down Expand Up @@ -79,4 +80,57 @@ public void run(boolean loadBalancingSucceeded) throws InterruptedException {
.subscribeOn(Schedulers.boundedElastic())
.subscribe();
}

@Test(groups = "unit")
public void run_multipleLeasesReturnedByStrategy_areAllAttempted() {
PartitionController partitionControllerMock = Mockito.mock(PartitionController.class);
LeaseContainer leaseContainerMock = Mockito.mock(LeaseContainer.class);
PartitionLoadBalancingStrategy partitionLoadBalancingStrategyMock = Mockito.mock(PartitionLoadBalancingStrategy.class);

ServiceItemLeaseV1 lease1 = new ServiceItemLeaseV1().withLeaseToken("1");
lease1.setId("TestLease-" + UUID.randomUUID());
ServiceItemLeaseV1 lease2 = new ServiceItemLeaseV1().withLeaseToken("2");
lease2.setId("TestLease-" + UUID.randomUUID());
ServiceItemLeaseV1 lease3 = new ServiceItemLeaseV1().withLeaseToken("3");
lease3.setId("TestLease-" + UUID.randomUUID());

List<Lease> allLeases = Arrays.asList(lease1, lease2, lease3);
Mockito.when(leaseContainerMock.getAllLeases()).thenReturn(Flux.fromIterable(allLeases));

// PartitionLoadBalancer collects leases into a new list instance each cycle, so don't match on identity.
Mockito.when(partitionLoadBalancingStrategyMock.selectLeasesToTake(Mockito.anyList()))
.thenReturn(allLeases)
.thenReturn(Collections.emptyList());

// Ensure the "happy path" doesn't get swallowed by the load balancer's error handler.
Mockito.when(partitionControllerMock.addOrUpdateLease(Mockito.any()))
.thenAnswer(invocation -> Mono.just((Lease) invocation.getArgument(0)));
Mockito.when(partitionControllerMock.shutdown()).thenReturn(Mono.empty());

PartitionLoadBalancerImpl partitionLoadBalancerImpl =
new PartitionLoadBalancerImpl(
partitionControllerMock,
leaseContainerMock,
partitionLoadBalancingStrategyMock,
Duration.ofSeconds(1),
Schedulers.boundedElastic(),
null
);

partitionLoadBalancerImpl
.start()
.timeout(Duration.ofMillis(PARTITION_LOAD_BALANCER_TIMEOUT))
.subscribeOn(Schedulers.boundedElastic())
.subscribe();

// Wait until the first balancing cycle attempts all leases returned by the strategy.
Mockito.verify(partitionControllerMock, Mockito.timeout(PARTITION_LOAD_BALANCER_TIMEOUT).times(allLeases.size()))
.addOrUpdateLease(Mockito.any());

partitionLoadBalancerImpl
.stop()
.timeout(Duration.ofMillis(PARTITION_LOAD_BALANCER_TIMEOUT))
.subscribeOn(Schedulers.boundedElastic())
.subscribe();
}
}
51 changes: 51 additions & 0 deletions .../com/azure/cosmos/implementation/changefeed/pkversion/PartitionLoadBalancerImplTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import java.time.Duration;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.UUID;

Expand Down Expand Up @@ -78,4 +79,54 @@ public void run(boolean loadBalancingSucceeded) throws InterruptedException {
.subscribeOn(Schedulers.boundedElastic())
.subscribe();
}

@Test(groups = "unit")
public void run_multipleLeasesReturnedByStrategy_areAllAttempted() {
PartitionController partitionControllerMock = Mockito.mock(PartitionController.class);
LeaseContainer leaseContainerMock = Mockito.mock(LeaseContainer.class);
PartitionLoadBalancingStrategy partitionLoadBalancingStrategyMock = Mockito.mock(PartitionLoadBalancingStrategy.class);

ServiceItemLease lease1 = new ServiceItemLease().withLeaseToken("1");
lease1.setId("TestLease-" + UUID.randomUUID());
ServiceItemLease lease2 = new ServiceItemLease().withLeaseToken("2");
lease2.setId("TestLease-" + UUID.randomUUID());
ServiceItemLease lease3 = new ServiceItemLease().withLeaseToken("3");
lease3.setId("TestLease-" + UUID.randomUUID());

List<Lease> allLeases = Arrays.asList(lease1, lease2, lease3);
Mockito.when(leaseContainerMock.getAllLeases()).thenReturn(Flux.fromIterable(allLeases));

// PartitionLoadBalancer collects leases into a new list instance each cycle, so don't match on identity.
Mockito.when(partitionLoadBalancingStrategyMock.selectLeasesToTake(Mockito.anyList()))
.thenReturn(allLeases)
.thenReturn(Collections.emptyList());

Mockito.when(partitionControllerMock.addOrUpdateLease(Mockito.any()))
.thenAnswer(invocation -> Mono.just((Lease) invocation.getArgument(0)));
Mockito.when(partitionControllerMock.shutdown()).thenReturn(Mono.empty());

PartitionLoadBalancerImpl partitionLoadBalancerImpl =
new PartitionLoadBalancerImpl(
partitionControllerMock,
leaseContainerMock,
partitionLoadBalancingStrategyMock,
Duration.ofSeconds(1),
Schedulers.boundedElastic()
);

partitionLoadBalancerImpl
.start()
.timeout(Duration.ofMillis(PARTITION_LOAD_BALANCER_TIMEOUT))
.subscribeOn(Schedulers.boundedElastic())
.subscribe();

Mockito.verify(partitionControllerMock, Mockito.timeout(PARTITION_LOAD_BALANCER_TIMEOUT).times(allLeases.size()))
.addOrUpdateLease(Mockito.any());

partitionLoadBalancerImpl
.stop()
.timeout(Duration.ofMillis(PARTITION_LOAD_BALANCER_TIMEOUT))
.subscribeOn(Schedulers.boundedElastic())
.subscribe();
}
}
1 change: 1 addition & 0 deletions sdk/cosmos/azure-cosmos/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
### 4.77.0-beta.1 (Unreleased)

#### Features Added
* Added `ChangeFeedProcessorOptions#setMaxLeasesToAcquirePerCycle(int)` to allow faster acquisition of unused/expired leases during scale-out and rolling deployments (default `0` preserves legacy behavior).
Copy link
Member

@FabianMeiswinkel FabianMeiswinkel Dec 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NIT - template should be followed (including link to the PR)


#### Breaking Changes

Expand Down
59 changes: 40 additions & 19 deletions ...a/com/azure/cosmos/implementation/changefeed/common/EqualPartitionsBalancingStrategy.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,18 @@ public class EqualPartitionsBalancingStrategy implements PartitionLoadBalancingS
private final int minPartitionCount;
private final int maxPartitionCount;
private final Duration leaseExpirationInterval;
private final int maxLeasesToAcquirePerCycle;

public EqualPartitionsBalancingStrategy(String hostName, int minPartitionCount, int maxPartitionCount, Duration leaseExpirationInterval) {
this(hostName, minPartitionCount, maxPartitionCount, leaseExpirationInterval, 0);
}

public EqualPartitionsBalancingStrategy(
String hostName,
int minPartitionCount,
int maxPartitionCount,
Duration leaseExpirationInterval,
int maxLeasesToAcquirePerCycle) {
if (hostName == null) {
throw new IllegalArgumentException("hostName");
}
Expand All @@ -35,6 +45,10 @@ public EqualPartitionsBalancingStrategy(String hostName, int minPartitionCount,
this.minPartitionCount = minPartitionCount;
this.maxPartitionCount = maxPartitionCount;
this.leaseExpirationInterval = leaseExpirationInterval;
if (maxLeasesToAcquirePerCycle < 0) {
throw new IllegalArgumentException("maxLeasesToAcquirePerCycle cannot be negative");
}
this.maxLeasesToAcquirePerCycle = maxLeasesToAcquirePerCycle;
}

@Override
Expand All @@ -58,42 +72,49 @@ public List<Lease> selectLeasesToTake(List<Lease> allLeases) {

if (expiredLeases.size() > 0) {
// We should try to pick at least one expired lease even if already overbooked when maximum partition count is not set.
// If other CFP instances are running, limit the number of expired leases to acquire to maximum 1 (non-greedy acquiring).
if ((this.maxPartitionCount == 0 && partitionsNeededForMe <= 0) || (partitionsNeededForMe > 1 && workerToPartitionCount.size() > 1)) {
// If other CFP instances are running, legacy behavior limits the number of expired leases to acquire to maximum 1
// (non-greedy acquiring) to reduce collisions. A configured maxLeasesToAcquirePerCycle overrides this clamp.
if (this.maxPartitionCount == 0 && partitionsNeededForMe <= 0) {
partitionsNeededForMe = 1;
} else if (partitionsNeededForMe > 1 && workerToPartitionCount.size() > 1 && this.maxLeasesToAcquirePerCycle == 0) {
// Legacy behavior: clamp to 1 when multiple workers exist.
partitionsNeededForMe = 1;
}

if (this.maxLeasesToAcquirePerCycle > 0) {
partitionsNeededForMe = Math.min(partitionsNeededForMe, this.maxLeasesToAcquirePerCycle);
}

if (partitionsNeededForMe <= 0) {
return new ArrayList<>();
}

// Try to minimize potential collisions between different CFP instances trying to pick the same lease.
// For multiple acquisitions, take a random subset.
Random random = new Random();
Collections.shuffle(expiredLeases, random);

if (partitionsNeededForMe == 1) {
// Try to minimize potential collisions between different CFP instances trying to pick the same lease.
Random random = new Random();
Lease expiredLease = expiredLeases.get(random.nextInt(expiredLeases.size()));
Lease expiredLease = expiredLeases.get(0);
this.logger.info("Found unused or expired lease {} (owner was {}); previous lease count for instance owner {} is {}, count of leases to target is {} and maxScaleCount {} ",
expiredLease.getLeaseToken(), expiredLease.getOwner(), this.hostName, myCount, partitionsNeededForMe, this.maxPartitionCount);

return Collections.singletonList(expiredLease);
} else {
for (Lease lease : expiredLeases) {
this.logger.info("Found unused or expired lease {} (owner was {}); previous lease count for instance owner {} is {} and maxScaleCount {} ",
lease.getLeaseToken(), lease.getOwner(), this.hostName, myCount, this.maxPartitionCount);
}
}

// If we reach here with partitionsNeededForMe < 0, then it means the change feed processor instances has owned leases >= the maxScaleCount.
// Then in this case, the change feed processor instance will not pick up any new leases.
if (partitionsNeededForMe <= 0)
{
return new ArrayList<>();
}
this.logger.info("Found {} unused or expired leases; previous lease count for instance owner {} is {}, count of leases to target is {} and maxScaleCount {} ",
Copy link
Member

@jeet1995 jeet1995 Dec 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These log lines can include maxLeasesToAcquirePerCycle as well.

expiredLeases.size(), this.hostName, myCount, partitionsNeededForMe, this.maxPartitionCount);

return expiredLeases.subList(0, Math.min(partitionsNeededForMe, expiredLeases.size()));
}

if (partitionsNeededForMe <= 0)
return new ArrayList<Lease>();
if (partitionsNeededForMe <= 0) {
return new ArrayList<>();
}

// Intentionally keep the legacy behavior for stealing: attempt to steal at most 1 lease per cycle.
Lease stolenLease = getLeaseToSteal(workerToPartitionCount, target, partitionsNeededForMe, allPartitions);
List<Lease> stolenLeases = new ArrayList<>();

if (stolenLease != null) {
stolenLeases.add(stolenLease);
}
Expand Down
3 changes: 2 additions & 1 deletion ...va/com/azure/cosmos/implementation/changefeed/epkversion/ChangeFeedProcessorImplBase.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -443,7 +443,8 @@ private Mono<PartitionManager> buildPartitionManager(LeaseStoreManager leaseStor
this.hostName,
this.changeFeedProcessorOptions.getMinScaleCount(),
this.changeFeedProcessorOptions.getMaxScaleCount(),
this.changeFeedProcessorOptions.getLeaseExpirationInterval());
this.changeFeedProcessorOptions.getLeaseExpirationInterval(),
this.changeFeedProcessorOptions.getMaxLeasesToAcquirePerCycle());
}

PartitionController partitionController = new PartitionControllerImpl(leaseStoreManager, leaseStoreManager, partitionSupervisorFactory, synchronizer, scheduler);
Expand Down
9 changes: 5 additions & 4 deletions .../azure/cosmos/implementation/changefeed/pkversion/IncrementalChangeFeedProcessorImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -469,10 +469,11 @@ private Mono<PartitionManager> buildPartitionManager(LeaseStoreManager leaseStor

if (this.loadBalancingStrategy == null) {
this.loadBalancingStrategy = new EqualPartitionsBalancingStrategy(
this.hostName,
this.changeFeedProcessorOptions.getMinScaleCount(),
this.changeFeedProcessorOptions.getMaxScaleCount(),
this.changeFeedProcessorOptions.getLeaseExpirationInterval());
this.hostName,
this.changeFeedProcessorOptions.getMinScaleCount(),
this.changeFeedProcessorOptions.getMaxScaleCount(),
this.changeFeedProcessorOptions.getLeaseExpirationInterval(),
this.changeFeedProcessorOptions.getMaxLeasesToAcquirePerCycle());
}

PartitionController partitionController =
Expand Down
Loading
Loading