Make InflightReadsLimiter asynchronous and apply it for replay queue reads

Author: lhotari

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

@@ -61,6 +61,18 @@ public class ManagedLedgerFactoryConfig {
      */
     private long managedLedgerMaxReadsInFlightSize = 0;
 
+    /**
+     * Maximum time to wait for acquiring permits for max reads in flight when managedLedgerMaxReadsInFlightSizeInMB is
+     * set (>0) and the limit is reached.
+     */
+    private long managedLedgerMaxReadsInFlightPermitsAcquireTimeoutMillis = 60000;
+
+    /**
+     * Maximum number of reads that can be queued for acquiring permits for max reads in flight when
+     * managedLedgerMaxReadsInFlightSizeInMB is set (>0) and the limit is reached.
+     */
+    private int managedLedgerMaxReadsInFlightPermitsAcquireQueueSize = 10000;
+
     /**
      * Whether trace managed ledger task execution time.
      */

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

@@ -228,7 +228,7 @@ private ManagedLedgerFactoryImpl(MetadataStoreExtended metadataStore,
                 compressionConfigForManagedCursorInfo);
         this.config = config;
         this.mbean = new ManagedLedgerFactoryMBeanImpl(this);
-        this.entryCacheManager = new RangeEntryCacheManagerImpl(this, openTelemetry);
+        this.entryCacheManager = new RangeEntryCacheManagerImpl(this, scheduledExecutor, openTelemetry);
         this.statsTask = scheduledExecutor.scheduleWithFixedDelay(catchingAndLoggingThrowables(this::refreshStats),
                 0, config.getStatsPeriodSeconds(), TimeUnit.SECONDS);
         this.flushCursorsTask = scheduledExecutor.scheduleAtFixedRate(catchingAndLoggingThrowables(this::flushCursors),

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

@@ -22,12 +22,18 @@
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.metrics.ObservableLongCounter;
 import io.prometheus.client.Gauge;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 import lombok.AllArgsConstructor;
 import lombok.ToString;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.opentelemetry.Constants;
 import org.apache.pulsar.opentelemetry.OpenTelemetryAttributes.InflightReadLimiterUtilization;
 import org.apache.pulsar.opentelemetry.annotations.PulsarDeprecatedMetric;
+import org.jctools.queues.SpscArrayQueue;
 
 @Slf4j
 public class InflightReadsLimiter implements AutoCloseable {
@@ -58,16 +64,41 @@ public class InflightReadsLimiter implements AutoCloseable {
 
     private final long maxReadsInFlightSize;
     private long remainingBytes;
+    private final long acquireTimeoutMillis;
+    private final ScheduledExecutorService timeOutExecutor;
+    private final boolean enabled;
 
-    public InflightReadsLimiter(long maxReadsInFlightSize, OpenTelemetry openTelemetry) {
-        if (maxReadsInFlightSize <= 0) {
+    @AllArgsConstructor
+    @ToString
+    static class Handle {
+        final long permits;
+        final long creationTime;
+        final boolean success;
+    }
+
+    record QueuedHandle(Handle handle, Consumer<Handle> callback) {
+    }
+
+    private final Queue<QueuedHandle> queuedHandles;
+    private boolean timeoutCheckRunning = false;
+
+    public InflightReadsLimiter(long maxReadsInFlightSize, int maxReadsInFlightAcquireQueueSize,
+                                long acquireTimeoutMillis, ScheduledExecutorService timeOutExecutor,
+                                OpenTelemetry openTelemetry) {
+        this.maxReadsInFlightSize = maxReadsInFlightSize;
+        this.remainingBytes = maxReadsInFlightSize;
+        this.acquireTimeoutMillis = acquireTimeoutMillis;
+        this.timeOutExecutor = timeOutExecutor;
+        if (maxReadsInFlightSize > 0) {
+            enabled = true;
+            this.queuedHandles = new SpscArrayQueue<>(maxReadsInFlightAcquireQueueSize);
+        } else {
+            enabled = false;
+            this.queuedHandles = null;
             // set it to -1 in order to show in the metrics that the metric is not available
             PULSAR_ML_READS_BUFFER_SIZE.set(-1);
             PULSAR_ML_READS_AVAILABLE_BUFFER_SIZE.set(-1);
         }
-        this.maxReadsInFlightSize = maxReadsInFlightSize;
-        this.remainingBytes = maxReadsInFlightSize;
-
         var meter = openTelemetry.getMeter(Constants.BROKER_INSTRUMENTATION_SCOPE_NAME);
         inflightReadsLimitCounter = meter.counterBuilder(INFLIGHT_READS_LIMITER_LIMIT_METRIC_NAME)
                 .setDescription("Maximum number of bytes that can be retained by managed ledger data read from storage "
@@ -102,71 +133,145 @@ public void close() {
         inflightReadsUsageCounter.close();
     }
 
-    @AllArgsConstructor
-    @ToString
-    static class Handle {
-        final long acquiredPermits;
-        final boolean success;
-        final int trials;
+    private static final Handle DISABLED = new Handle(0, 0, true);
+    private static final Optional<Handle> DISABLED_OPTIONAL = Optional.of(DISABLED);
 
-        final long creationTime;
+    /**
+     * Acquires permits from the limiter. If the limiter is disabled, it will immediately return a successful handle.
+     * If permits are available, it will return a handle with the acquired permits. If no permits are available,
+     * it will return an empty optional and the callback will be called when permits become available or when the
+     * acquire timeout is reached. The success field in the handle passed to the callback will be false if the acquire
+     * operation times out. The callback should be non-blocking and run on a desired executor handled within the
+     * callback itself.
+     *
+     * A successful handle will have the success field set to true, and the caller must call release with the handle
+     * when the permits are no longer needed.
+     *
+     * If an unsuccessful handle is returned immediately, it means that the queue limit has been reached and the
+     * callback will not be called. The caller should fail the read operation in this case to apply backpressure.
+     *
+     * @param permits  the number of permits to acquire
+     * @param callback the callback to be called when the permits are acquired or timed out
+     * @return an optional handle that contains the permits if acquired, otherwise an empty optional
+     */
+    public Optional<Handle> acquire(long permits, Consumer<Handle> callback) {
+        if (isDisabled()) {
+            return DISABLED_OPTIONAL;
+        }
+        return internalAcquire(permits, callback);
     }
 
-    private static final Handle DISABLED = new Handle(0, true, 0, -1);
+    private synchronized Optional<Handle> internalAcquire(long permits, Consumer<Handle> callback) {
+        Handle handle = new Handle(permits, System.currentTimeMillis(), true);
+        if (remainingBytes >= permits) {
+            remainingBytes -= permits;
+            if (log.isDebugEnabled()) {
+                log.debug("acquired permits: {}, creationTime: {}, remainingBytes:{}", permits, handle.creationTime,
+                        remainingBytes);
+            }
+            updateMetrics();
+            return Optional.of(handle);
+        } else {
+            if (queuedHandles.offer(new QueuedHandle(handle, callback))) {
+                scheduleTimeOutCheck(acquireTimeoutMillis);
+                return Optional.empty();
+            } else {
+                log.warn("Failed to queue handle for acquiring permits: {}, creationTime: {}, remainingBytes:{}",
+                        permits, handle.creationTime, remainingBytes);
+                return Optional.of(new Handle(0, handle.creationTime, false));
+            }
+        }
+    }
 
-    Handle acquire(long permits, Handle current) {
-        if (maxReadsInFlightSize <= 0) {
-            // feature is disabled
-            return DISABLED;
+    private synchronized void scheduleTimeOutCheck(long delayMillis) {
+        if (acquireTimeoutMillis <= 0) {
+            return;
         }
-        synchronized (this) {
-            try {
-                if (current == null) {
-                    if (remainingBytes == 0) {
-                        return new Handle(0, false, 1, System.currentTimeMillis());
-                    }
-                    if (remainingBytes >= permits) {
-                        remainingBytes -= permits;
-                        return new Handle(permits, true, 1, System.currentTimeMillis());
-                    } else {
-                        long possible = remainingBytes;
-                        remainingBytes = 0;
-                        return new Handle(possible, false, 1, System.currentTimeMillis());
-                    }
+        if (!timeoutCheckRunning) {
+            timeoutCheckRunning = true;
+            timeOutExecutor.schedule(this::timeoutCheck, delayMillis, TimeUnit.MILLISECONDS);
+        }
+    }
+
+    private synchronized void timeoutCheck() {
+        timeoutCheckRunning = false;
+        long delay = 0;
+        while (true) {
+            QueuedHandle queuedHandle = queuedHandles.peek();
+            if (queuedHandle != null) {
+                long age = System.currentTimeMillis() - queuedHandle.handle.creationTime;
+                if (age >= acquireTimeoutMillis) {
+                    // remove the peeked handle from the queue
+                    queuedHandles.poll();
+                    handleTimeout(queuedHandle);
                 } else {
-                    if (current.trials >= 4 && current.acquiredPermits > 0) {
-                        remainingBytes += current.acquiredPermits;
-                        return new Handle(0, false, 1, current.creationTime);
-                    }
-                    if (remainingBytes == 0) {
-                        return new Handle(current.acquiredPermits, false, current.trials + 1,
-                                current.creationTime);
-                    }
-                    long needed = permits - current.acquiredPermits;
-                    if (remainingBytes >= needed) {
-                        remainingBytes -= needed;
-                        return new Handle(permits, true, current.trials + 1, current.creationTime);
-                    } else {
-                        long possible = remainingBytes;
-                        remainingBytes = 0;
-                        return new Handle(current.acquiredPermits + possible, false,
-                                current.trials + 1, current.creationTime);
-                    }
+                    delay = acquireTimeoutMillis - age;
+                    break;
                 }
-            } finally {
-                updateMetrics();
+            } else {
+                break;
             }
         }
+        if (delay > 0) {
+            scheduleTimeOutCheck(delay);
+        }
+    }
+
+    private void handleTimeout(QueuedHandle queuedHandle) {
+        if (log.isDebugEnabled()) {
+            log.debug("timed out queued permits: {}, creationTime: {}, remainingBytes:{}",
+                    queuedHandle.handle.permits, queuedHandle.handle.creationTime, remainingBytes);
+        }
+        queuedHandle.callback.accept(new Handle(0, queuedHandle.handle.creationTime, false));
     }
 
-    void release(Handle handle) {
+    /**
+     * Releases permits back to the limiter. If the handle is disabled, this method will be a no-op.
+     *
+     * @param handle the handle containing the permits to release
+     */
+    public void release(Handle handle) {
         if (handle == DISABLED) {
             return;
         }
-        synchronized (this) {
-            remainingBytes += handle.acquiredPermits;
-            updateMetrics();
+        internalRelease(handle);
+    }
+
+    private synchronized void internalRelease(Handle handle) {
+        if (log.isDebugEnabled()) {
+            log.debug("release permits: {}, creationTime: {}, remainingBytes:{}", handle.permits,
+                    handle.creationTime, getRemainingBytes());
+        }
+        remainingBytes += handle.permits;
+        while (true) {
+            QueuedHandle queuedHandle = queuedHandles.peek();
+            if (queuedHandle != null) {
+                if (remainingBytes >= queuedHandle.handle.permits) {
+                    // remove the peeked handle from the queue
+                    queuedHandles.poll();
+                    handleQueuedHandle(queuedHandle);
+                } else if (acquireTimeoutMillis > 0
+                        && System.currentTimeMillis() - queuedHandle.handle.creationTime > acquireTimeoutMillis) {
+                    // remove the peeked handle from the queue
+                    queuedHandles.poll();
+                    handleTimeout(queuedHandle);
+                } else {
+                    break;
+                }
+            } else {
+                break;
+            }
+        }
+        updateMetrics();
+    }
+
+    private void handleQueuedHandle(QueuedHandle queuedHandle) {
+        remainingBytes -= queuedHandle.handle.permits;
+        if (log.isDebugEnabled()) {
+            log.debug("acquired queued permits: {}, creationTime: {}, remainingBytes:{}",
+                    queuedHandle.handle.permits, queuedHandle.handle.creationTime, remainingBytes);
         }
+        queuedHandle.callback.accept(queuedHandle.handle);
     }
 
     private synchronized void updateMetrics() {
@@ -175,8 +280,6 @@ private synchronized void updateMetrics() {
     }
 
     public boolean isDisabled() {
-        return maxReadsInFlightSize <= 0;
+        return !enabled;
     }
-
-
-}
+}