Revert "Remove CompletableFuture from RecoverySourceHandler (#93665)"

This reverts commit 67509c0.

Author: DaveCTurner

server/src/main/java/org/elasticsearch/indices/recovery/RecoverySourceHandler.java

@@ -79,6 +79,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -199,7 +200,12 @@ public void recoverToTarget(ActionListener<RecoveryResponse> listener) {
                 retentionLeaseRef.set(
                     shard.getRetentionLeases().get(ReplicationTracker.getPeerRecoveryRetentionLeaseId(targetShardRouting))
                 );
-            }, shardId + " validating recovery target [" + request.targetAllocationId() + "] registered ", shard, cancellableThreads);
+            },
+                shardId + " validating recovery target [" + request.targetAllocationId() + "] registered ",
+                shard,
+                cancellableThreads,
+                logger
+            );
             final Closeable retentionLock = shard.acquireHistoryRetentionLock();
             resources.add(retentionLock);
             final long startingSeqNo;
@@ -286,7 +292,7 @@ && isTargetSameHistory()
                             logger.debug("no peer-recovery retention lease for " + request.targetAllocationId());
                             deleteRetentionLeaseStep.onResponse(null);
                         }
-                    }, shardId + " removing retention lease for [" + request.targetAllocationId() + "]", shard, cancellableThreads);
+                    }, shardId + " removing retention lease for [" + request.targetAllocationId() + "]", shard, cancellableThreads, logger);
 
                     deleteRetentionLeaseStep.whenComplete(ignored -> {
                         assert Transports.assertNotTransportThread(RecoverySourceHandler.this + "[phase1]");
@@ -317,7 +323,8 @@ && isTargetSameHistory()
                     () -> shard.initiateTracking(request.targetAllocationId()),
                     shardId + " initiating tracking of " + request.targetAllocationId(),
                     shard,
-                    cancellableThreads
+                    cancellableThreads,
+                    logger
                 );
 
                 final long endingSeqNo = shard.seqNoStats().getMaxSeqNo();
@@ -398,21 +405,44 @@ static void runUnderPrimaryPermit(
         CancellableThreads.Interruptible runnable,
         String reason,
         IndexShard primary,
-        CancellableThreads cancellableThreads
+        CancellableThreads cancellableThreads,
+        Logger logger
     ) {
         cancellableThreads.execute(() -> {
-            final var permit = new ListenableActionFuture<Releasable>();
-            primary.acquirePrimaryOperationPermit(permit, ThreadPool.Names.SAME, reason);
-            try (var ignored = FutureUtils.get(permit)) {
+            CompletableFuture<Releasable> permit = new CompletableFuture<>();
+
+            // this wrapping looks unnecessary necessary, see #93290; TODO remove it
+            final ActionListener<Releasable> onAcquired = new ActionListener<Releasable>() {
+                @Override
+                public void onResponse(Releasable releasable) {
+                    if (permit.complete(releasable) == false) {
+                        releasable.close();
+                    }
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    permit.completeExceptionally(e);
+                }
+            };
+            primary.acquirePrimaryOperationPermit(onAcquired, ThreadPool.Names.SAME, reason);
+            try (Releasable ignored = FutureUtils.get(permit)) {
                 // check that the IndexShard still has the primary authority. This needs to be checked under operation permit to prevent
                 // races, as IndexShard will switch its authority only when it holds all operation permits, see IndexShard.relocated()
                 if (primary.isRelocatedPrimary()) {
                     throw new IndexShardRelocatedException(primary.shardId());
                 }
                 runnable.run();
             } finally {
-                // add a listener to release the permit because we might have been interrupted while waiting (double-releasing is ok)
-                permit.addListener(ActionListener.wrap(Releasable::close, e -> {}));
+                // just in case we got an exception (likely interrupted) while waiting for the get
+                permit.whenComplete((r, e) -> {
+                    if (r != null) {
+                        r.close();
+                    }
+                    if (e != null) {
+                        logger.trace("suppressing exception on completion (it was already bubbled up or the operation was aborted)", e);
+                    }
+                });
             }
         });
     }
@@ -980,7 +1010,7 @@ void createRetentionLease(final long startingSeqNo, ActionListener<RetentionLeas
                 addRetentionLeaseStep.addListener(listener.map(rr -> newLease));
                 logger.trace("created retention lease with estimated checkpoint of [{}]", estimatedGlobalCheckpoint);
             }
-        }, shardId + " establishing retention lease for [" + request.targetAllocationId() + "]", shard, cancellableThreads);
+        }, shardId + " establishing retention lease for [" + request.targetAllocationId() + "]", shard, cancellableThreads, logger);
     }
 
     boolean hasSameLegacySyncId(Store.MetadataSnapshot source, Store.MetadataSnapshot target) {
@@ -1221,7 +1251,8 @@ void finalizeRecovery(long targetLocalCheckpoint, long trimAboveSeqNo, ActionLis
             () -> shard.markAllocationIdAsInSync(request.targetAllocationId(), targetLocalCheckpoint),
             shardId + " marking " + request.targetAllocationId() + " as in sync",
             shard,
-            cancellableThreads
+            cancellableThreads,
+            logger
         );
         final long globalCheckpoint = shard.getLastKnownGlobalCheckpoint(); // this global checkpoint is persisted in finalizeRecovery
         final StepListener<Void> finalizeListener = new StepListener<>();
@@ -1232,7 +1263,8 @@ void finalizeRecovery(long targetLocalCheckpoint, long trimAboveSeqNo, ActionLis
                 () -> shard.updateGlobalCheckpointForShard(request.targetAllocationId(), globalCheckpoint),
                 shardId + " updating " + request.targetAllocationId() + "'s global checkpoint",
                 shard,
-                cancellableThreads
+                cancellableThreads,
+                logger
             );
 
             if (request.isPrimaryRelocation()) {

server/src/test/java/org/elasticsearch/indices/recovery/RecoverySourceHandlerTests.java

@@ -798,7 +798,7 @@ public void testCancellationsDoesNotLeakPrimaryPermits() throws Exception {
         Thread cancelingThread = new Thread(() -> cancellableThreads.cancel("test"));
         cancelingThread.start();
         try {
-            RecoverySourceHandler.runUnderPrimaryPermit(() -> {}, "test", shard, cancellableThreads);
+            RecoverySourceHandler.runUnderPrimaryPermit(() -> {}, "test", shard, cancellableThreads, logger);
         } catch (CancellableThreads.ExecutionCancelledException e) {
             // expected.
         }