[fix][ml] Track all pending read callbacks for timeouts

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

@@ -133,6 +133,7 @@ public class ManagedLedgerFactoryImpl implements ManagedLedgerFactory {
     private long lastStatTimestamp = System.nanoTime();
     private final ScheduledFuture<?> statsTask;
     private final ScheduledFuture<?> flushCursorsTask;
+    private final ReadEntryTimeoutTracker readEntryTimeoutTracker;
 
     private volatile long cacheEvictionTimeThresholdNanos;
     private final MetadataStore metadataStore;
@@ -245,6 +246,7 @@ private ManagedLedgerFactoryImpl(MetadataStoreExtended metadataStore,
         this.config = config;
         this.mbean = new ManagedLedgerFactoryMBeanImpl(this);
         this.entryCacheManager = new RangeEntryCacheManagerImpl(this, scheduledExecutor, openTelemetry);
+        this.readEntryTimeoutTracker = new ReadEntryTimeoutTracker(scheduledExecutor);
         this.statsTask = scheduledExecutor.scheduleWithFixedDelay(catchingAndLoggingThrowables(this::refreshStats),
                 0, config.getStatsPeriodSeconds(), TimeUnit.SECONDS);
         this.flushCursorsTask = scheduledExecutor.scheduleAtFixedRate(catchingAndLoggingThrowables(this::flushCursors),
@@ -324,6 +326,10 @@ public synchronized void doCacheEviction() {
         entryCacheManager.doCacheEviction();
     }
 
+    ReadEntryTimeoutTracker getReadEntryTimeoutTracker() {
+        return readEntryTimeoutTracker;
+    }
+
     /**
      * Waits for all pending cache evictions based on total cache size or entry TTL to complete.
      * This is for testing purposes only, so that we can ensure all cache evictions are done before proceeding with
@@ -662,6 +668,7 @@ public CompletableFuture<Void> shutdownAsync() throws ManagedLedgerException {
 
         statsTask.cancel(true);
         flushCursorsTask.cancel(true);
+        readEntryTimeoutTracker.close();
         cacheEvictionExecutor.shutdownNow();
 
         List<String> ledgerNames = new ArrayList<>(this.ledgers.keySet());

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

@@ -348,12 +348,6 @@ public boolean isFenced() {
             .newUpdater(ManagedLedgerImpl.class, "addOpCount");
     private volatile long addOpCount = 0;
 
-    // last read-operation's callback to check read-timeout on it.
-    private volatile ReadEntryCallbackWrapper lastReadCallback = null;
-    private static final AtomicReferenceFieldUpdater<ManagedLedgerImpl, ReadEntryCallbackWrapper>
-            LAST_READ_CALLBACK_UPDATER = AtomicReferenceFieldUpdater
-            .newUpdater(ManagedLedgerImpl.class, ReadEntryCallbackWrapper.class, "lastReadCallback");
-
     /**
      * Queue of pending entries to be added to the managed ledger. Typically, entries are queued when a new ledger is.
      * created asynchronously and hence there is no ready ledger to write into.
@@ -529,7 +523,7 @@ public void operationFailed(MetaStoreException e) {
             }
         });
 
-        scheduleTimeoutTask();
+        scheduleAddEntryTimeoutTask();
     }
 
     protected ManagedLedgerInterceptor.LastEntryHandle createLastEntryHandle(LedgerHandle lh) {
@@ -2437,10 +2431,10 @@ protected void asyncReadEntry(ReadHandle ledger, Position position, ReadEntryCal
         if (config.getReadEntryTimeoutSeconds() > 0) {
             // set readOpCount to uniquely validate if ReadEntryCallbackWrapper is already recycled
             long readOpCount = READ_OP_COUNT_UPDATER.incrementAndGet(this);
-            long createdTime = System.nanoTime();
-            ReadEntryCallbackWrapper readCallback = ReadEntryCallbackWrapper.create(name, position.getLedgerId(),
-                    position.getEntryId(), callback, readOpCount, createdTime, ctx);
-            lastReadCallback = readCallback;
+            ReadEntryCallbackWrapper readCallback = ReadEntryCallbackWrapper.create(this, position.getLedgerId(),
+                    position.getEntryId(), callback, readOpCount, ctx);
+            readCallback.readTimeout = factory.getReadEntryTimeoutTracker().add(readCallback, readOpCount,
+                    timeoutAtNanos(config.getReadEntryTimeoutSeconds()));
             entryCache.asyncReadEntry(ledger, position, readCallback, readOpCount);
         } else {
             entryCache.asyncReadEntry(ledger, position, callback, ctx);
@@ -2453,58 +2447,60 @@ protected void asyncReadEntry(ReadHandle ledger, long firstEntry, long lastEntry
         if (config.getReadEntryTimeoutSeconds() > 0) {
             // set readOpCount to uniquely validate if ReadEntryCallbackWrapper is already recycled
             long readOpCount = READ_OP_COUNT_UPDATER.incrementAndGet(this);
-            long createdTime = System.nanoTime();
-            ReadEntryCallbackWrapper readCallback = ReadEntryCallbackWrapper.create(name, ledger.getId(), firstEntry,
-                    opReadEntry, readOpCount, createdTime, ctx);
-            lastReadCallback = readCallback;
+            ReadEntryCallbackWrapper readCallback = ReadEntryCallbackWrapper.create(this, ledger.getId(), firstEntry,
+                    opReadEntry, readOpCount, ctx);
+            readCallback.readTimeout = factory.getReadEntryTimeoutTracker().add(readCallback, readOpCount,
+                    timeoutAtNanos(config.getReadEntryTimeoutSeconds()));
             entryCache.asyncReadEntry(ledger, firstEntry, lastEntry, expectedReadCount, readCallback, readOpCount);
         } else {
             entryCache.asyncReadEntry(ledger, firstEntry, lastEntry, expectedReadCount, opReadEntry, ctx);
         }
     }
 
+    private static long timeoutAtNanos(long timeoutSec) {
+        return System.nanoTime() + TimeUnit.SECONDS.toNanos(timeoutSec);
+    }
+
     static final class ReadEntryCallbackWrapper implements ReadEntryCallback, ReadEntriesCallback {
 
         volatile ReadEntryCallback readEntryCallback;
         volatile ReadEntriesCallback readEntriesCallback;
-        String name;
+        volatile ManagedLedgerImpl managedLedger;
         long ledgerId;
         long entryId;
         volatile long readOpCount = -1;
         private static final AtomicLongFieldUpdater<ReadEntryCallbackWrapper> READ_OP_COUNT_UPDATER =
                 AtomicLongFieldUpdater.newUpdater(ReadEntryCallbackWrapper.class, "readOpCount");
-        volatile long createdTime = -1;
         volatile Object cntx;
+        volatile ReadEntryTimeoutTracker.ReadTimeoutWrapper readTimeout;
 
         final Handle<ReadEntryCallbackWrapper> recyclerHandle;
 
         private ReadEntryCallbackWrapper(Handle<ReadEntryCallbackWrapper> recyclerHandle) {
             this.recyclerHandle = recyclerHandle;
         }
 
-        static ReadEntryCallbackWrapper create(String name, long ledgerId, long entryId, ReadEntryCallback callback,
-                long readOpCount, long createdTime, Object ctx) {
+        static ReadEntryCallbackWrapper create(ManagedLedgerImpl managedLedger, long ledgerId, long entryId,
+                                               ReadEntryCallback callback, long readOpCount, Object ctx) {
             ReadEntryCallbackWrapper readCallback = RECYCLER.get();
-            readCallback.name = name;
+            readCallback.managedLedger = managedLedger;
             readCallback.ledgerId = ledgerId;
             readCallback.entryId = entryId;
             readCallback.readEntryCallback = callback;
             readCallback.cntx = ctx;
             readCallback.readOpCount = readOpCount;
-            readCallback.createdTime = createdTime;
             return readCallback;
         }
 
-        static ReadEntryCallbackWrapper create(String name, long ledgerId, long entryId, ReadEntriesCallback callback,
-                long readOpCount, long createdTime, Object ctx) {
+        static ReadEntryCallbackWrapper create(ManagedLedgerImpl managedLedger, long ledgerId, long entryId,
+                                               ReadEntriesCallback callback, long readOpCount, Object ctx) {
             ReadEntryCallbackWrapper readCallback = RECYCLER.get();
-            readCallback.name = name;
+            readCallback.managedLedger = managedLedger;
             readCallback.ledgerId = ledgerId;
             readCallback.entryId = entryId;
             readCallback.readEntriesCallback = callback;
             readCallback.cntx = ctx;
             readCallback.readOpCount = readOpCount;
-            readCallback.createdTime = createdTime;
             return readCallback;
         }
 
@@ -2562,6 +2558,15 @@ private long reOpCount(Object ctx) {
             return (ctx instanceof Long) ? (long) ctx : -1;
         }
 
+        boolean shouldTriggerReadTimeout() {
+            ManagedLedgerImpl ledger = managedLedger;
+            if (ledger == null) {
+                return false;
+            }
+            State state = STATE_UPDATER.get(ledger);
+            return state != State.Closed && !state.isFenced();
+        }
+
         public void readFailed(ManagedLedgerException exception, Object ctx) {
             if (readEntryCallback != null) {
                 readEntryFailed(exception, ctx);
@@ -2575,12 +2580,17 @@ public void readFailed(ManagedLedgerException exception, Object ctx) {
         private boolean recycle(long readOpCount) {
             if (readOpCount != -1
                     && READ_OP_COUNT_UPDATER.compareAndSet(ReadEntryCallbackWrapper.this, readOpCount, -1)) {
-                createdTime = -1;
+                ReadEntryTimeoutTracker.ReadTimeoutWrapper timeout = readTimeout;
+                if (timeout != null) {
+                    timeout.clearCallback(this);
+                }
                 readEntryCallback = null;
                 readEntriesCallback = null;
                 ledgerId = -1;
                 entryId = -1;
-                name = null;
+                managedLedger = null;
+                cntx = null;
+                readTimeout = null;
                 recyclerHandle.recycle(this);
                 return true;
             }
@@ -4724,29 +4734,20 @@ protected boolean checkAndCompleteLedgerOpTask(int rc, LedgerHandle lh, Object c
         return false;
     }
 
-    private void scheduleTimeoutTask() {
-        // disable timeout task checker if timeout <= 0
-        if (config.getAddEntryTimeoutSeconds() > 0 || config.getReadEntryTimeoutSeconds() > 0) {
-            long timeoutSec = Math.min(config.getAddEntryTimeoutSeconds(), config.getReadEntryTimeoutSeconds());
-            timeoutSec = timeoutSec <= 0
-                    ? Math.max(config.getAddEntryTimeoutSeconds(), config.getReadEntryTimeoutSeconds())
-                    : timeoutSec;
+    private void scheduleAddEntryTimeoutTask() {
+        if (config.getAddEntryTimeoutSeconds() > 0) {
+            long timeoutSec = config.getAddEntryTimeoutSeconds();
             this.timeoutTask = this.scheduledExecutor.scheduleAtFixedRate(
-                    this::checkTimeouts, timeoutSec, timeoutSec, TimeUnit.SECONDS);
+                    this::checkAddTimeout, timeoutSec, timeoutSec, TimeUnit.SECONDS);
         }
     }
 
-    private void checkTimeouts() {
+    private void checkAddTimeout() {
         final State state = STATE_UPDATER.get(this);
         if (state == State.Closed
                 || state.isFenced()) {
             return;
         }
-        checkAddTimeout();
-        checkReadTimeout();
-    }
-
-    private void checkAddTimeout() {
         long timeoutSec = config.getAddEntryTimeoutSeconds();
         if (timeoutSec < 1) {
             return;
@@ -4766,24 +4767,6 @@ private void checkAddTimeout() {
         }
     }
 
-    private void checkReadTimeout() {
-        long timeoutSec = config.getReadEntryTimeoutSeconds();
-        if (timeoutSec < 1) {
-            return;
-        }
-        ReadEntryCallbackWrapper callback = this.lastReadCallback;
-        long readOpCount = callback != null ? callback.readOpCount : 0;
-        boolean timeout = callback != null && (TimeUnit.NANOSECONDS
-                .toSeconds(System.nanoTime() - callback.createdTime) >= timeoutSec);
-        if (readOpCount > 0 && timeout) {
-            log.warn().attr("ledgerId", this.lastReadCallback.ledgerId)
-                    .attr("entryId", this.lastReadCallback.entryId)
-                    .attr("timeoutSec", timeoutSec).log("Read entry timeout");
-            callback.readFailed(createManagedLedgerException(BKException.Code.TimeoutException), readOpCount);
-            LAST_READ_CALLBACK_UPDATER.compareAndSet(this, callback, null);
-        }
-    }
-
     @Override
     public long getOffloadedSize() {
         long offloadedSize = 0;

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

@@ -0,0 +1,132 @@
+/*
+ * 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.apache.bookkeeper.mledger.impl.ManagedLedgerImpl.createManagedLedgerException;
+import static org.apache.pulsar.common.util.Runnables.catchingAndLoggingThrowables;
+import com.google.common.annotations.VisibleForTesting;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
+import lombok.CustomLog;
+import org.apache.bookkeeper.client.BKException;
+import org.apache.bookkeeper.common.util.OrderedScheduler;
+import org.jctools.queues.MpscUnboundedArrayQueue;
+
+@CustomLog
+class ReadEntryTimeoutTracker implements AutoCloseable {
+    private static final int READ_TIMEOUT_QUEUE_CHUNK_SIZE = 128 * 1024;
+    private static final int CHECK_INTERVAL_SECONDS = 1;
+
+    private final MpscUnboundedArrayQueue<ReadTimeoutWrapper> timeoutQueue = new MpscUnboundedArrayQueue<>(
+            READ_TIMEOUT_QUEUE_CHUNK_SIZE);
+    private final AtomicInteger timeoutQueueSize = new AtomicInteger();
+    private final ScheduledFuture<?> timeoutTask;
+
+    ReadEntryTimeoutTracker(OrderedScheduler scheduledExecutor) {
+        this.timeoutTask = scheduledExecutor.scheduleAtFixedRate(catchingAndLoggingThrowables(this::checkTimeouts),
+                CHECK_INTERVAL_SECONDS, CHECK_INTERVAL_SECONDS, TimeUnit.SECONDS);
+    }
+
+    ReadTimeoutWrapper add(ManagedLedgerImpl.ReadEntryCallbackWrapper callback, long readOpCount,
+                           long timeoutAtNanos) {
+        ReadTimeoutWrapper timeout = new ReadTimeoutWrapper(readOpCount, timeoutAtNanos, callback);
+        timeoutQueue.offer(timeout);
+        timeoutQueueSize.incrementAndGet();
+        return timeout;
+    }
+
+    @VisibleForTesting
+    synchronized void checkTimeouts() {
+        long now = System.nanoTime();
+        int entriesToProcess = timeoutQueueSize.get();
+        for (int i = 0; i < entriesToProcess; i++) {
+            ReadTimeoutWrapper timeout = timeoutQueue.poll();
+            if (timeout == null) {
+                return;
+            }
+            timeoutQueueSize.decrementAndGet();
+            ManagedLedgerImpl.ReadEntryCallbackWrapper callback = timeout.getCallback();
+            if (callback == null) {
+                continue;
+            }
+            if (!callback.shouldTriggerReadTimeout()) {
+                timeout.clearCallback(callback);
+                continue;
+            }
+            if (timeout.timeoutAtNanos > now) {
+                requeue(timeout);
+                continue;
+            }
+            callback = timeout.clearCallback();
+            if (callback != null) {
+                log.warn()
+                        .attr("overdueNanos", now - timeout.timeoutAtNanos)
+                        .log("Read entry timeout");
+                callback.readFailed(createManagedLedgerException(BKException.Code.TimeoutException),
+                        timeout.readOpCount);
+            }
+        }
+    }
+
+    private void requeue(ReadTimeoutWrapper timeout) {
+        timeoutQueue.offer(timeout);
+        timeoutQueueSize.incrementAndGet();
+    }
+
+    @Override
+    public void close() {
+        timeoutTask.cancel(false);
+    }
+
+    @VisibleForTesting
+    int pendingTimeoutCount() {
+        return timeoutQueueSize.get();
+    }
+
+    static final class ReadTimeoutWrapper {
+        private static final AtomicReferenceFieldUpdater<ReadTimeoutWrapper,
+                ManagedLedgerImpl.ReadEntryCallbackWrapper> CALLBACK_UPDATER = AtomicReferenceFieldUpdater
+                .newUpdater(ReadTimeoutWrapper.class, ManagedLedgerImpl.ReadEntryCallbackWrapper.class, "callback");
+
+        final long readOpCount;
+        final long timeoutAtNanos;
+        volatile ManagedLedgerImpl.ReadEntryCallbackWrapper callback;
+
+        ReadTimeoutWrapper(long readOpCount, long timeoutAtNanos,
+                           ManagedLedgerImpl.ReadEntryCallbackWrapper callback) {
+            this.readOpCount = readOpCount;
+            this.timeoutAtNanos = timeoutAtNanos;
+            this.callback = callback;
+        }
+
+        ManagedLedgerImpl.ReadEntryCallbackWrapper getCallback() {
+            return callback;
+        }
+
+        ManagedLedgerImpl.ReadEntryCallbackWrapper clearCallback() {
+            return CALLBACK_UPDATER.getAndSet(this, null);
+        }
+
+        void clearCallback(ManagedLedgerImpl.ReadEntryCallbackWrapper callback) {
+            CALLBACK_UPDATER.compareAndSet(this, callback, null);
+        }
+    }
+}