[improve][broker] Expose interface for the replicator in ManagedLedger instead of cast the class (#26298)

Author: zymap

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

@@ -18,12 +18,14 @@
  */
 package org.apache.bookkeeper.mledger;
 
+import static org.apache.pulsar.common.util.Runnables.catchingAndLoggingThrowables;
 import com.google.common.collect.Range;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
 import org.apache.bookkeeper.common.annotation.InterfaceAudience;
 import org.apache.bookkeeper.common.annotation.InterfaceStability;
@@ -879,6 +881,20 @@ default void skipNonRecoverableLedger(long ledgerId){}
      */
     ManagedLedger getManagedLedger();
 
+    /**
+     * Schedule a continuation of a read callback.
+     *
+     * <p>Implementations that deliver read callbacks on a dedicated execution context should override this method
+     * to run the continuation on that same execution context.
+     *
+     * @param callback the callback continuation
+     * @param delay the delay before executing the continuation
+     * @param unit the time unit of the delay
+     */
+    default void scheduleReadCallback(Runnable callback, long delay, TimeUnit unit) {
+        CompletableFuture.delayedExecutor(delay, unit).execute(catchingAndLoggingThrowables(callback));
+    }
+
     /**
      * Get last individual deleted range.
      * @return range

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

@@ -26,6 +26,7 @@
 import static org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl.DEFAULT_LEDGER_DELETE_RETRIES;
 import static org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl.createManagedLedgerException;
 import static org.apache.bookkeeper.mledger.util.Errors.isNoSuchLedgerExistsException;
+import static org.apache.pulsar.common.util.Runnables.catchingAndLoggingThrowables;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.Collections2;
@@ -4021,6 +4022,12 @@ public ManagedLedger getManagedLedger() {
         return this.ledger;
     }
 
+    @Override
+    public void scheduleReadCallback(Runnable callback, long delay, TimeUnit unit) {
+        ledger.getScheduledExecutor().schedule(
+                catchingAndLoggingThrowables(() -> ledger.getExecutor().execute(callback)), delay, unit);
+    }
+
     @Override
     public Range<Position> getLastIndividualDeletedRange() {
         lock.readLock().lock();

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

@@ -100,6 +100,7 @@
 import org.apache.bookkeeper.client.api.LedgerEntries;
 import org.apache.bookkeeper.client.api.ReadHandle;
 import org.apache.bookkeeper.common.util.OrderedExecutor;
+import org.apache.bookkeeper.common.util.OrderedScheduler;
 import org.apache.bookkeeper.mledger.AsyncCallbacks;
 import org.apache.bookkeeper.mledger.AsyncCallbacks.AddEntryCallback;
 import org.apache.bookkeeper.mledger.AsyncCallbacks.DeleteCallback;
@@ -138,6 +139,7 @@
 import org.apache.pulsar.metadata.api.extended.SessionEvent;
 import org.apache.pulsar.metadata.impl.FaultInjectionMetadataStore;
 import org.awaitility.Awaitility;
+import org.mockito.ArgumentCaptor;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 import org.mockito.invocation.InvocationOnMock;
@@ -4198,6 +4200,36 @@ public void testEstimatedUnackedSizeWhenCursorCaughtUpWithLastPosition() throws
         assertEquals(cursor.getEstimatedSizeSinceMarkDeletePosition(), 0);
     }
 
+    @Test
+    public void testScheduleReadCallbackUsesManagedLedgerExecutionContext() {
+        ManagedLedgerImpl ledger = mock(ManagedLedgerImpl.class);
+        when(ledger.getConfig()).thenReturn(new ManagedLedgerConfig());
+        when(ledger.getLogger()).thenReturn(log);
+        OrderedScheduler scheduledExecutor = mock(OrderedScheduler.class);
+        ExecutorService executor = mock(ExecutorService.class);
+        when(ledger.getScheduledExecutor()).thenReturn(scheduledExecutor);
+        when(ledger.getExecutor()).thenReturn(executor);
+        ManagedCursorImpl cursor = new ManagedCursorImpl(mock(BookKeeper.class), ledger, "c1");
+        Runnable callback = mock(Runnable.class);
+        ArgumentCaptor<Runnable> scheduledTask = ArgumentCaptor.forClass(Runnable.class);
+
+        cursor.scheduleReadCallback(callback, 100, TimeUnit.MILLISECONDS);
+
+        verify(scheduledExecutor).schedule(scheduledTask.capture(), eq(100L), eq(TimeUnit.MILLISECONDS));
+        scheduledTask.getValue().run();
+        verify(executor).execute(callback);
+    }
+
+    @Test
+    public void testDefaultScheduleReadCallback() throws InterruptedException {
+        ManagedCursor cursor = mock(ManagedCursor.class, Mockito.CALLS_REAL_METHODS);
+        CountDownLatch callbackExecuted = new CountDownLatch(1);
+
+        cursor.scheduleReadCallback(callbackExecuted::countDown, 0, TimeUnit.MILLISECONDS);
+
+        assertTrue(callbackExecuted.await(5, TimeUnit.SECONDS));
+    }
+
     @Test
     public void testEstimatedUnackedSizeWhenCursorAdvancedToEmptyCurrentLedger() {
         ManagedLedgerImpl ledger = mock(ManagedLedgerImpl.class);

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentReplicator.java

@@ -53,7 +53,6 @@
 import org.apache.bookkeeper.mledger.ManagedLedgerException.CursorAlreadyClosedException;
 import org.apache.bookkeeper.mledger.ManagedLedgerException.TooManyRequestsException;
 import org.apache.bookkeeper.mledger.Position;
-import org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pulsar.broker.PulsarServerException;
@@ -402,7 +401,6 @@ public void readEntriesComplete(List<Entry> entries, Object ctx) {
         }
 
         // Retry to trigger read completes if it is not started.
-        ManagedLedgerImpl ml = (ManagedLedgerImpl) cursor.getManagedLedger();
         Runnable retryReplicateEntries = () -> {
             long estimatedTimeStampProducerConnected = this.estimatedTimeStampProducerConnected;
             long delayMillis;
@@ -411,11 +409,7 @@ public void readEntriesComplete(List<Entry> entries, Object ctx) {
             } else {
                 delayMillis = 100;
             }
-            ml.getScheduledExecutor().schedule(() -> {
-                ml.getExecutor().execute(() -> {
-                    readEntriesComplete(entries, ctx);
-                });
-            }, delayMillis, TimeUnit.MILLISECONDS);
+            cursor.scheduleReadCallback(() -> readEntriesComplete(entries, ctx), delayMillis, TimeUnit.MILLISECONDS);
         };
 
         // Retry.