Create FutureManagerRunner and convert FDBDatabaseRunnerImpl to use it. (#3347)

This refactoring moves the future management (completing of futures once
the runner closes) out of the FDBDatabaseRunnerImpl into a subclass of
TransactionalRunner.
This allows the new subclass of TransactionalRunner to be used (e.g in
the case where we want another throttling implementation).

Resolves #3346

Author: ohadzeliger

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/FDBDatabaseRunnerImpl.java

@@ -29,6 +29,7 @@
 import com.apple.foundationdb.record.logging.KeyValueLogMessage;
 import com.apple.foundationdb.record.logging.LogMessageKeys;
 import com.apple.foundationdb.record.provider.foundationdb.runners.ExponentialDelay;
+import com.apple.foundationdb.record.provider.foundationdb.runners.FutureAutoClose;
 import com.apple.foundationdb.record.provider.foundationdb.runners.TransactionalRunner;
 import com.apple.foundationdb.record.provider.foundationdb.synchronizedsession.SynchronizedSessionRunner;
 import com.apple.foundationdb.record.util.Result;
@@ -38,7 +39,6 @@
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -57,6 +57,7 @@ public class FDBDatabaseRunnerImpl implements FDBDatabaseRunner {
     @Nonnull
     private final FDBDatabase database;
     private final TransactionalRunner transactionalRunner;
+    private final FutureAutoClose futureManager;
     @Nonnull
     private FDBRecordContextConfig.Builder contextConfigBuilder;
     @Nonnull
@@ -67,22 +68,19 @@ public class FDBDatabaseRunnerImpl implements FDBDatabaseRunner {
     private long initialDelayMillis;
 
     private boolean closed;
-    @Nonnull
-    private final List<CompletableFuture<?>> futuresToCompleteExceptionally;
 
     @API(API.Status.INTERNAL)
     FDBDatabaseRunnerImpl(@Nonnull FDBDatabase database, FDBRecordContextConfig.Builder contextConfigBuilder) {
         this.database = database;
         this.contextConfigBuilder = contextConfigBuilder;
         this.executor = database.newContextExecutor(contextConfigBuilder.getMdcContext());
         this.transactionalRunner = new TransactionalRunner(database, contextConfigBuilder);
+        this.futureManager = new FutureAutoClose();
 
         final FDBDatabaseFactory factory = database.getFactory();
         this.maxAttempts = factory.getMaxAttempts();
         this.maxDelayMillis = factory.getMaxDelayMillis();
         this.initialDelayMillis = factory.getInitialDelayMillis();
-
-        futuresToCompleteExceptionally = new ArrayList<>();
     }
 
     @Override
@@ -225,7 +223,7 @@ LogMessageKeys.MAX_ATTEMPTS, getMaxAttempts(),
                     if (getTimer() != null) {
                         future = getTimer().instrument(FDBStoreTimer.Events.RETRY_DELAY, future, executor);
                     }
-                    addFutureToCompleteExceptionally(future);
+                    futureManager.registerFuture(future);
                     return future.thenApply(vignore -> {
                         currAttempt++;
                         return true;
@@ -240,8 +238,7 @@ LogMessageKeys.MAX_ATTEMPTS, getMaxAttempts(),
         @SuppressWarnings("squid:S1181")
         public CompletableFuture<T> runAsync(@Nonnull final Function<? super FDBRecordContext, CompletableFuture<? extends T>> retriable,
                                              @Nonnull final BiFunction<? super T, Throwable, Result<? extends T, ? extends Throwable>> handlePostTransaction) {
-            CompletableFuture<T> future = new CompletableFuture<>();
-            addFutureToCompleteExceptionally(future);
+            CompletableFuture<T> future = futureManager.newFuture();
             AsyncUtil.whileTrue(() -> {
                 try {
                     return transactionalRunner.runAsync(currAttempt != 0, retriable)
@@ -314,13 +311,25 @@ public synchronized void close() {
             return;
         }
         closed = true;
-        if (!futuresToCompleteExceptionally.stream().allMatch(CompletableFuture::isDone)) {
-            final Exception exception = new RunnerClosed();
-            for (CompletableFuture<?> future : futuresToCompleteExceptionally) {
-                future.completeExceptionally(exception);
+        // Ensure we call both close() methods, capturing all exceptions
+        RuntimeException caught = null;
+        try {
+            futureManager.close();
+        } catch (RuntimeException e) {
+            caught = e;
+        }
+        try {
+            transactionalRunner.close();
+        } catch (RuntimeException e) {
+            if (caught != null) {
+                caught.addSuppressed(e);
+            } else {
+                caught = e;
             }
         }
-        transactionalRunner.close();
+        if (caught != null) {
+            throw caught;
+        }
     }
 
     @Override
@@ -337,15 +346,4 @@ public SynchronizedSessionRunner startSynchronizedSession(@Nonnull Subspace lock
     public SynchronizedSessionRunner joinSynchronizedSession(@Nonnull Subspace lockSubspace, @Nonnull UUID sessionId, long leaseLengthMillis) {
         return SynchronizedSessionRunner.joinSession(lockSubspace, sessionId, leaseLengthMillis, this);
     }
-
-    private synchronized void addFutureToCompleteExceptionally(@Nonnull CompletableFuture<?> future) {
-        if (closed) {
-            final RunnerClosed exception = new RunnerClosed();
-            future.completeExceptionally(exception);
-            throw exception;
-        }
-        futuresToCompleteExceptionally.removeIf(CompletableFuture::isDone);
-        futuresToCompleteExceptionally.add(future);
-    }
-
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/runners/FutureAutoClose.java

@@ -0,0 +1,90 @@
+/*
+ * FutureAutoClose.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.apple.foundationdb.record.provider.foundationdb.runners;
+
+import com.apple.foundationdb.annotation.API;
+import com.apple.foundationdb.record.provider.foundationdb.FDBDatabaseRunner;
+
+import javax.annotation.Nonnull;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * A helper that completes futures once closed.
+ * <p>
+ * This is meant to be used together with other {@link FDBDatabaseRunner}s and help manage any {@link CompletableFuture} instances
+ * once the runner closes. This helper should be lifecycle coupled to the runners it accompanies.
+ * <p>
+ * Futures created externally can be registered with this class. This class can also create new futures (and register them).
+ * When {@link #close()} is called, any registered {@link CompletableFuture} which is still incomplete will be completed exceptionally
+ * with a {@link FDBDatabaseRunner.RunnerClosed} exception.
+ * <p>
+ * It is the responsibility of the users of the runner to ensure futures that it creates are registered.
+ * <p>
+ * Normally, only the top level (root) of the future chain needs registration. Completing this future
+ * will cause the completion to trickle down to the dependent futures.
+ */
+@API(API.Status.INTERNAL)
+public class FutureAutoClose implements AutoCloseable {
+    @Nonnull
+    private final List<CompletableFuture<?>> futuresToClose;
+    private boolean closed;
+
+    public FutureAutoClose() {
+        futuresToClose = new ArrayList<>();
+        closed = false;
+    }
+
+    public <T> CompletableFuture<T> newFuture() {
+        return registerFuture(new CompletableFuture<>());
+    }
+
+    public synchronized <T> CompletableFuture<T> registerFuture(CompletableFuture<T> future) {
+        if (isClosed()) {
+            final FDBDatabaseRunner.RunnerClosed exception = new FDBDatabaseRunner.RunnerClosed();
+            future.completeExceptionally(exception);
+            throw exception;
+        }
+        futuresToClose.removeIf(CompletableFuture::isDone);
+        futuresToClose.add(future);
+        return future;
+    }
+
+    @Override
+    public synchronized void close() {
+        if (!isClosed()) {
+            closed = true;
+
+            if (!futuresToClose.stream().allMatch(CompletableFuture::isDone)) {
+                final Exception exception = new FDBDatabaseRunner.RunnerClosed();
+                for (CompletableFuture<?> future : futuresToClose) {
+                    future.completeExceptionally(exception);
+                }
+            }
+            futuresToClose.clear();
+        }
+    }
+
+    public boolean isClosed() {
+        return closed;
+    }
+}

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/provider/foundationdb/runners/TransactionalRunner.java

@@ -175,5 +175,4 @@ public synchronized void close() {
         contextsToClose.forEach(FDBRecordContext::close);
         this.closed = true;
     }
-
 }