Code review comment: I combined both cases. Earlier we only do the delay on success case, now we will also need to do on error case. Since they are using the same logic, so I combined into one

Code before review:
private Mono<Void> run(CancellationToken cancellationToken) {
return Flux.just(this)
.flatMap(value -> this.leaseContainer.getAllLeases())
.collectList()
.flatMap(allLeases -> {
if (cancellationToken.isCancellationRequested()) return Mono.empty();
List<Lease> leasesToTake = this.partitionLoadBalancingStrategy.selectLeasesToTake(allLeases);
if (leasesToTake.size() > 0) {
this.logger.info("Found {} total leases, taking ownership of {}", allLeases.size(), leasesToTake.size());
}
if (cancellationToken.isCancellationRequested()) return Mono.empty();
return Flux.fromIterable(leasesToTake)
.limitRate(1)
.flatMap(lease -> {
if (cancellationToken.isCancellationRequested()) return Mono.empty();
return this.partitionController.addOrUpdateLease(lease);
})
.then();
})
.onErrorResume(throwable -> {
logger.warn("Unexpected exception thrown while trying to acquire available leases", throwable);
return Mono.empty();
})
.then(
Mono.just(this)
.flatMap(value -> {
if (cancellationToken.isCancellationRequested()) {
return Mono.empty();
}
Instant stopTimer = Instant.now().plus(this.leaseAcquireInterval);
return Mono.just(value)
.delayElement(Duration.ofMillis(100), CosmosSchedulers.COSMOS_PARALLEL)
.repeat(() -> {
Instant currentTime = Instant.now();
return !cancellationToken.isCancellationRequested() && currentTime.isBefore(stopTimer);
})
.then();
})
)
.repeat(() -> !cancellationToken.isCancellationRequested())
.then()
.onErrorResume(throwable -> {
logger.info("Partition load balancer task stopped.");
return this.stop();
});
}

Line under review: .then();

Code after addressing the review:
private Mono<Void> run(CancellationToken cancellationToken) {
return Flux.just(this)
.flatMap(value -> this.leaseContainer.getAllLeases())
.collectList()
.flatMap(allLeases -> {
if (cancellationToken.isCancellationRequested()) return Mono.empty();
List<Lease> leasesToTake = this.partitionLoadBalancingStrategy.selectLeasesToTake(allLeases);
if (leasesToTake.size() > 0) {
this.logger.info("Found {} total leases, taking ownership of {}", allLeases.size(), leasesToTake.size());
}
if (cancellationToken.isCancellationRequested()) return Mono.empty();
return Flux.fromIterable(leasesToTake)
.limitRate(1)
.flatMap(lease -> {
if (cancellationToken.isCancellationRequested()) return Mono.empty();
return this.partitionController.addOrUpdateLease(lease);
})
.then();
})
.onErrorResume(throwable -> {
logger.warn("Unexpected exception thrown while trying to acquire available leases", throwable);
return Mono.empty();
})
.then(
Mono.just(this)
.flatMap(value -> {
if (cancellationToken.isCancellationRequested()) {
return Mono.empty();
}
Instant stopTimer = Instant.now().plus(this.leaseAcquireInterval);
return Mono.just(value)
.delayElement(Duration.ofMillis(100), CosmosSchedulers.COSMOS_PARALLEL)
.repeat(() -> {
Instant currentTime = Instant.now();
return !cancellationToken.isCancellationRequested() && currentTime.isBefore(stopTimer);
})
.then();
})
)
.repeat(() -> !cancellationToken.isCancellationRequested())
.then()
.onErrorResume(throwable -> {
logger.info("Partition load balancer task stopped.");
return this.stop();
});
}