[improve][offload] Coalesce automatic offload triggers to reduce retry loops and ledger scans

managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/AutomaticOffloadTriggerController.java

@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.bookkeeper.mledger.impl;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Coalesces repeated automatic offload triggers into at most one active run and one follow-up run.
+ */
+final class AutomaticOffloadTriggerController {
+    private static final int IDLE = 0;
+    private static final int RUNNING = 1;
+    private static final int RUNNING_WITH_PENDING_TRIGGER = 2;
+
+    private final AtomicInteger state = new AtomicInteger(IDLE);
+
+    /**
+     * Records an automatic offload trigger.
+     *
+     * @return true when the caller must start a new automatic offload run
+     */
+    boolean requestRun() {
+        while (true) {
+            int current = state.get();
+            switch (current) {
+                case IDLE:
+                    if (state.compareAndSet(IDLE, RUNNING)) {
+                        return true;
+                    }
+                    break;
+                case RUNNING:
+                    if (state.compareAndSet(RUNNING, RUNNING_WITH_PENDING_TRIGGER)) {
+                        return false;
+                    }
+                    break;
+                case RUNNING_WITH_PENDING_TRIGGER:
+                    return false;
+                default:
+                    throw new IllegalStateException("Unknown automatic offload trigger state: " + current);
+            }
+        }
+    }
+
+    /**
+     * Records completion of the current automatic offload run.
+     *
+     * @return true when the caller must immediately start one coalesced follow-up run
+     */
+    boolean completeRun() {
+        while (true) {
+            int current = state.get();
+            switch (current) {
+                case IDLE:
+                    return false;
+                case RUNNING:
+                    if (state.compareAndSet(RUNNING, IDLE)) {
+                        return false;
+                    }
+                    break;
+                case RUNNING_WITH_PENDING_TRIGGER:
+                    if (state.compareAndSet(RUNNING_WITH_PENDING_TRIGGER, RUNNING)) {
+                        return true;
+                    }
+                    break;
+                default:
+                    throw new IllegalStateException("Unknown automatic offload trigger state: " + current);
+            }
+        }
+    }
+}

managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/ManagedLedgerFactoryImpl.java

@@ -19,7 +19,7 @@
 package org.apache.bookkeeper.mledger.impl;
 
 import static org.apache.bookkeeper.mledger.ManagedLedgerException.getManagedLedgerException;
-import static org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl.NULL_OFFLOAD_PROMISE;
+import static org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl.AUTOMATIC_OFFLOAD_TRIGGER;
 import static org.apache.pulsar.common.util.Runnables.catchingAndLoggingThrowables;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Predicates;
@@ -496,7 +496,7 @@ public void initializeComplete() {
                                     future.complete(newledger);
                                     // May need to trigger offloading
                                     if (config.isTriggerOffloadOnTopicLoad()) {
-                                        newledger.maybeOffloadInBackground(NULL_OFFLOAD_PROMISE);
+                                        newledger.maybeOffloadInBackground(AUTOMATIC_OFFLOAD_TRIGGER);
                                     }
                                 });
                             }

managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/ManagedLedgerImpl.java

@@ -231,8 +231,20 @@ public Logger getLogger() {
     protected final CallbackMutex trimmerMutex = new CallbackMutex();
 
     protected final CallbackMutex offloadMutex = new CallbackMutex();
-    public static final CompletableFuture<Position> NULL_OFFLOAD_PROMISE = CompletableFuture
+    private final AutomaticOffloadTriggerController automaticOffloadTriggerController =
+            new AutomaticOffloadTriggerController();
+    // Identity sentinel for automatic offload requests. The completed Position value is not used.
+    public static final CompletableFuture<Position> AUTOMATIC_OFFLOAD_TRIGGER = CompletableFuture
             .completedFuture(PositionFactory.LATEST);
+
+    private enum OffloadRequestSource {
+        AUTOMATIC,
+        EXPLICIT
+    }
+
+    private record OffloadThresholds(long thresholdInBytes, long thresholdInSeconds) {
+    }
+
     @VisibleForTesting
     @Getter
     protected volatile LedgerHandle currentLedger;
@@ -1968,7 +1980,7 @@ synchronized void ledgerClosed(final LedgerHandle lh, Long lastAddConfirmed) {
 
         trimConsumedLedgersInBackground();
 
-        maybeOffloadInBackground(NULL_OFFLOAD_PROMISE);
+        maybeOffloadInBackground(AUTOMATIC_OFFLOAD_TRIGGER);
 
         createLedgerAfterClosed();
     }
@@ -2804,22 +2816,73 @@ private void scheduleDeferredTrimming(boolean isTruncate, CompletableFuture<?> p
     }
 
     public void maybeOffloadInBackground(CompletableFuture<Position> promise) {
-        if (getOffloadPoliciesIfAppendable().isEmpty()) {
+        if (promise == AUTOMATIC_OFFLOAD_TRIGGER) {
+            if (automaticOffloadTriggerController.requestRun()) {
+                startAutomaticOffload();
+            }
+            return;
+        }
+
+        maybeOffloadInBackground(promise, OffloadRequestSource.EXPLICIT);
+    }
+
+    private void startAutomaticOffload() {
+        CompletableFuture<Position> automaticOffloadCompletion = new CompletableFuture<>();
+        automaticOffloadCompletion.whenComplete((res, ex) -> finishAutomaticOffload(ex));
+        try {
+            maybeOffloadInBackground(automaticOffloadCompletion, OffloadRequestSource.AUTOMATIC);
+        } catch (RuntimeException e) {
+            automaticOffloadCompletion.completeExceptionally(e);
+        }
+    }
+
+    private void maybeOffloadInBackground(CompletableFuture<Position> promise, OffloadRequestSource source) {
+        Optional<OffloadThresholds> offloadThresholds = getOffloadThresholds();
+        if (offloadThresholds.isEmpty()) {
+            if (source == OffloadRequestSource.AUTOMATIC) {
+                promise.complete(PositionFactory.LATEST);
+            }
             return;
         }
 
-        final OffloadPolicies policies = config.getLedgerOffloader().getOffloadPolicies();
+        OffloadThresholds thresholds = offloadThresholds.get();
+        try {
+            executor.execute(() -> maybeOffload(thresholds.thresholdInBytes(), thresholds.thresholdInSeconds(),
+                    promise, source));
+        } catch (RuntimeException e) {
+            promise.completeExceptionally(e);
+        }
+    }
+
+    private Optional<OffloadThresholds> getOffloadThresholds() {
+        Optional<OffloadPolicies> optionalOffloadPolicies = getOffloadPoliciesIfAppendable();
+        if (optionalOffloadPolicies.isEmpty()) {
+            return Optional.empty();
+        }
+
+        final OffloadPolicies policies = optionalOffloadPolicies.get();
         final long offloadThresholdInBytes =
                 Optional.ofNullable(policies.getManagedLedgerOffloadThresholdInBytes()).orElse(-1L);
         final long offloadThresholdInSeconds =
                 Optional.ofNullable(policies.getManagedLedgerOffloadThresholdInSeconds()).orElse(-1L);
         if (offloadThresholdInBytes >= 0 || offloadThresholdInSeconds >= 0) {
-            executor.execute(() -> maybeOffload(offloadThresholdInBytes, offloadThresholdInSeconds, promise));
+            return Optional.of(new OffloadThresholds(offloadThresholdInBytes, offloadThresholdInSeconds));
+        }
+
+        return Optional.empty();
+    }
+
+    private void finishAutomaticOffload(Throwable exception) {
+        if (exception != null) {
+            log.debug().exception(exception).log("Failed to automatically offload ledgers");
+        }
+        if (automaticOffloadTriggerController.completeRun()) {
+            startAutomaticOffload();
         }
     }
 
     private void maybeOffload(long offloadThresholdInBytes, long offloadThresholdInSeconds,
-                              CompletableFuture<Position> finalPromise) {
+                              CompletableFuture<Position> finalPromise, OffloadRequestSource source) {
         if (getOffloadPoliciesIfAppendable().isEmpty()) {
             String msg = String.format("[%s] Nothing to offload due to offloader or offloadPolicies is NULL", name);
             finalPromise.completeExceptionally(new IllegalArgumentException(msg));
@@ -2834,8 +2897,12 @@ private void maybeOffload(long offloadThresholdInBytes, long offloadThresholdInS
         }
 
         if (!offloadMutex.tryLock()) {
-            scheduledExecutor.schedule(() -> maybeOffloadInBackground(finalPromise),
-                    100, TimeUnit.MILLISECONDS);
+            try {
+                scheduledExecutor.schedule(() -> maybeOffloadInBackground(finalPromise, source),
+                        100, TimeUnit.MILLISECONDS);
+            } catch (RuntimeException e) {
+                finalPromise.completeExceptionally(e);
+            }
             return;
         }
 
@@ -2926,12 +2993,11 @@ void internalTrimConsumedLedgers(CompletableFuture<?> promise) {
 
     private Optional<OffloadPolicies> getOffloadPoliciesIfAppendable() {
         LedgerOffloader ledgerOffloader = config.getLedgerOffloader();
-        if (ledgerOffloader == null
-                || !ledgerOffloader.isAppendable()
-                || ledgerOffloader.getOffloadPolicies() == null) {
+        if (ledgerOffloader == null || !ledgerOffloader.isAppendable()) {
             return Optional.empty();
         }
-        return Optional.ofNullable(ledgerOffloader.getOffloadPolicies());
+        OffloadPolicies offloadPolicies = ledgerOffloader.getOffloadPolicies();
+        return Optional.ofNullable(offloadPolicies);
     }
 
     @VisibleForTesting

managed-ledger/src/test/java/org/apache/bookkeeper/mledger/impl/AutomaticOffloadTriggerControllerTest.java

@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.bookkeeper.mledger.impl;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import org.testng.annotations.Test;
+
+public class AutomaticOffloadTriggerControllerTest {
+
+    @Test
+    public void triggersCoalesceWhileRunIsActive() {
+        AutomaticOffloadTriggerController controller = new AutomaticOffloadTriggerController();
+
+        assertThat(controller.requestRun()).isTrue();
+        assertThat(controller.requestRun()).isFalse();
+        assertThat(controller.requestRun()).isFalse();
+    }
+
+    @Test
+    public void pendingTriggerSchedulesOneFollowUpRun() {
+        AutomaticOffloadTriggerController controller = new AutomaticOffloadTriggerController();
+
+        assertThat(controller.requestRun()).isTrue();
+        assertThat(controller.requestRun()).isFalse();
+
+        assertThat(controller.completeRun()).isTrue();
+        assertThat(controller.completeRun()).isFalse();
+        assertThat(controller.requestRun()).isTrue();
+    }
+
+    @Test(timeOut = 30000)
+    public void concurrentTriggerAndCompletionAlwaysReserveOneFollowUpRun() throws Exception {
+        ExecutorService executor = Executors.newFixedThreadPool(2);
+        try {
+            for (int i = 0; i < 1000; i++) {
+                AutomaticOffloadTriggerController controller = new AutomaticOffloadTriggerController();
+                assertThat(controller.requestRun()).isTrue();
+
+                // Completion and a new trigger can race; exactly one side must reserve the follow-up run.
+                CyclicBarrier barrier = new CyclicBarrier(3);
+                Future<Boolean> completeResult = executor.submit(() -> {
+                    barrier.await(5, TimeUnit.SECONDS);
+                    return controller.completeRun();
+                });
+                Future<Boolean> triggerResult = executor.submit(() -> {
+                    barrier.await(5, TimeUnit.SECONDS);
+                    return controller.requestRun();
+                });
+
+                barrier.await(5, TimeUnit.SECONDS);
+                boolean followUpReservedByComplete = completeResult.get(5, TimeUnit.SECONDS);
+                boolean followUpReservedByTrigger = triggerResult.get(5, TimeUnit.SECONDS);
+
+                assertThat(followUpReservedByComplete)
+                        .as("iteration %s must reserve exactly one follow-up run", i)
+                        .isNotEqualTo(followUpReservedByTrigger);
+                assertThat(controller.completeRun()).isFalse();
+            }
+        } finally {
+            executor.shutdownNow();
+            executor.awaitTermination(5, TimeUnit.SECONDS);
+        }
+    }
+}