From dd1323b31e5e4148c2c38c5b467c5daa853a9316 Mon Sep 17 00:00:00 2001
From: milismsft-mac <adrianmi@microsoft.com>
Date: Tue, 10 Dec 2019 12:48:26 -0800
Subject: [PATCH] Cosmos ChangeFeedProcessor fixes - fix to terminate the
 current processing threads when the lease renewal worker thread has failed. -
 fix to out of bounds exception when the leases load balancer thread is trying
 to aquire a lease.

---
 .../implementation/EqualPartitionsBalancingStrategy.java        | 2 +-
 .../internal/changefeed/implementation/LeaseRenewerImpl.java    | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/EqualPartitionsBalancingStrategy.java b/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/EqualPartitionsBalancingStrategy.java
index 302e7f5acc1f..cb19f309af9c 100644
--- a/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/EqualPartitionsBalancingStrategy.java
+++ b/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/EqualPartitionsBalancingStrategy.java
@@ -73,7 +73,7 @@ public List<Lease> selectLeasesToTake(List<Lease> allLeases) {
             return new ArrayList<Lease>();
 
         if (expiredLeases.size() > 0) {
-            return expiredLeases.subList(0, partitionsNeededForMe);
+            return expiredLeases.subList(0, Math.min(expiredLeases.size(), partitionsNeededForMe));
         }
 
         Lease stolenLease = getLeaseToSteal(workerToPartitionCount, target, partitionsNeededForMe, allPartitions);
diff --git a/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/LeaseRenewerImpl.java b/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/LeaseRenewerImpl.java
index 87328513f8fb..917f071561e3 100644
--- a/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/LeaseRenewerImpl.java
+++ b/sdk/cosmos/microsoft-azure-cosmos/src/main/java/com/azure/data/cosmos/internal/changefeed/implementation/LeaseRenewerImpl.java
@@ -65,8 +65,10 @@ public Mono<Void> run(CancellationToken cancellationToken) {
             .doOnError(throwable -> {
                 if (throwable instanceof LeaseLostException) {
                     logger.info("Partition {}: renew lease loop failed.", this.lease.getLeaseToken(), throwable);
+                    this.resultException = (LeaseLostException) throwable;
                 } else {
                     logger.error("Partition {}: renew lease loop failed.", this.lease.getLeaseToken(), throwable);
+                    this.resultException = new RuntimeException(throwable);
                 }
             });
     }