limit the memory used by reads end-to-end

conf/broker.conf

@@ -1064,7 +1064,15 @@ managedLedgerCursorRolloverTimeInSeconds=14400
 # crashes.
 managedLedgerMaxUnackedRangesToPersist=10000
 
-# Max number of "acknowledgment holes" that can be stored in Zookeeper. If number of unack message range is higher
+# Maximum amount of memory used hold data read from storage (or from the cache).
+# This mechanism prevents the broker to have too many concurrent
+# reads from storage and fall into Out of Memory errors in case
+# of multiple concurrent reads to multiple concurrent consumers.
+# Set 0 in order to disable the feature.
+#
+managedLedgerMaxReadsInFlightSizeInMB=0
+
+# Max number of "acknowledgment holes" that can be stored in MetadataStore. If number of unack message range is higher
 # than this limit then broker will persist unacked ranges into bookkeeper to avoid additional data overhead into
 # zookeeper.
 managedLedgerMaxUnackedRangesToPersistInZooKeeper=1000

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

@@ -57,6 +57,11 @@ public class ManagedLedgerFactoryConfig {
      */
     private boolean copyEntriesInCache = false;
 
+    /**
+     * Maximum number of (estimated) data in-flight reading from storage and the cache.
+     */
+    private long managedLedgerMaxReadsInFlightSize = 0;
+
     /**
      * Whether trace managed ledger task execution time.
      */

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

@@ -45,6 +45,8 @@ protected EntryImpl newObject(Handle<EntryImpl> handle) {
     private long entryId;
     ByteBuf data;
 
+    private Runnable onDeallocate;
+
     public static EntryImpl create(LedgerEntry ledgerEntry) {
         EntryImpl entry = RECYCLER.get();
         entry.timestamp = System.nanoTime();
@@ -103,6 +105,22 @@ private EntryImpl(Recycler.Handle<EntryImpl> recyclerHandle) {
         this.recyclerHandle = recyclerHandle;
     }
 
+    public void onDeallocate(Runnable r) {
+        if (this.onDeallocate == null) {
+            this.onDeallocate = r;
+        } else {
+            // this is not expected to happen
+            Runnable previous = this.onDeallocate;
+            this.onDeallocate = () -> {
+                try {
+                    previous.run();
+                } finally {
+                    r.run();
+                }
+            };
+        }
+    }
+
     public long getTimestamp() {
         return timestamp;
     }
@@ -160,6 +178,13 @@ public ReferenceCounted touch(Object hint) {
     @Override
     protected void deallocate() {
         // This method is called whenever the ref-count of the EntryImpl reaches 0, so that now we can recycle it
+        if (onDeallocate != null) {
+            try {
+                onDeallocate.run();
+            } finally {
+                onDeallocate = null;
+            }
+        }
         data.release();
         data = null;
         timestamp = -1;

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

@@ -0,0 +1,137 @@
+/**
+ * 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.cache;
+
+import com.google.common.annotations.VisibleForTesting;
+import io.prometheus.client.Gauge;
+import lombok.AllArgsConstructor;
+import lombok.ToString;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class InflightReadsLimiter {
+
+    private static final Gauge PULSAR_ML_READS_BUFFER_SIZE = Gauge
+            .build()
+            .name("pulsar_ml_reads_inflight_bytes")
+            .help("Estimated number of bytes retained by data read from storage or cache")
+            .register();
+
+    private static final Gauge PULSAR_ML_READS_AVAILABLE_BUFFER_SIZE = Gauge
+            .build()
+            .name("pulsar_ml_reads_available_inflight_bytes")
+            .help("Available space for inflight data read from storage or cache")
+            .register();
+
+    private final long maxReadsInFlightSize;
+    private long remainingBytes;
+
+    public InflightReadsLimiter(long maxReadsInFlightSize) {
+        if (maxReadsInFlightSize <= 0) {
+            // 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;
+    }
+
+    @VisibleForTesting
+    public synchronized long getRemainingBytes() {
+        return remainingBytes;
+    }
+
+    @AllArgsConstructor
+    @ToString
+    static class Handle {
+        final long acquiredPermits;
+        final boolean success;
+        final int trials;
+
+        final long creationTime;
+    }
+
+    private static final Handle DISABLED = new Handle(0, true, 0, -1);
+
+    Handle acquire(long permits, Handle current) {
+        if (maxReadsInFlightSize <= 0) {
+            // feature is disabled
+            return DISABLED;
+        }
+        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());
+                    }
+                } 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);
+                    }
+                }
+            } finally {
+                updateMetrics();
+            }
+        }
+    }
+
+    void release(Handle handle) {
+        if (handle == DISABLED) {
+            return;
+        }
+        synchronized (this) {
+            remainingBytes += handle.acquiredPermits;
+            updateMetrics();
+        }
+    }
+
+    private synchronized void updateMetrics() {
+        PULSAR_ML_READS_BUFFER_SIZE.set(maxReadsInFlightSize - remainingBytes);
+        PULSAR_ML_READS_AVAILABLE_BUFFER_SIZE.set(remainingBytes);
+    }
+
+    public boolean isDisabled() {
+        return maxReadsInFlightSize <= 0;
+    }
+
+
+}

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

@@ -21,18 +21,23 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl.createManagedLedgerException;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Lists;
 import com.google.common.primitives.Longs;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.PooledByteBufAllocator;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.bookkeeper.client.api.BKException;
 import org.apache.bookkeeper.client.api.LedgerEntry;
 import org.apache.bookkeeper.client.api.ReadHandle;
+import org.apache.bookkeeper.mledger.AsyncCallbacks;
 import org.apache.bookkeeper.mledger.AsyncCallbacks.ReadEntriesCallback;
 import org.apache.bookkeeper.mledger.AsyncCallbacks.ReadEntryCallback;
+import org.apache.bookkeeper.mledger.Entry;
+import org.apache.bookkeeper.mledger.ManagedLedgerConfig;
 import org.apache.bookkeeper.mledger.ManagedLedgerException;
 import org.apache.bookkeeper.mledger.impl.EntryImpl;
 import org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl;
@@ -48,18 +53,26 @@
  */
 public class RangeEntryCacheImpl implements EntryCache {
 
+    /**
+     * Overhead per-entry to take into account the envelope.
+     */
+    private static final long BOOKKEEPER_READ_OVERHEAD_PER_ENTRY = 64;
+
     private final RangeEntryCacheManagerImpl manager;
     private final ManagedLedgerImpl ml;
     private ManagedLedgerInterceptor interceptor;
     private final RangeCache<PositionImpl, EntryImpl> entries;
     private final boolean copyEntries;
+    private volatile long estimatedEntrySize = 10 * 1024;
 
+    private final long readEntryTimeoutMillis;
     private static final double MB = 1024 * 1024;
 
     public RangeEntryCacheImpl(RangeEntryCacheManagerImpl manager, ManagedLedgerImpl ml, boolean copyEntries) {
         this.manager = manager;
         this.ml = ml;
         this.interceptor = ml.getManagedLedgerInterceptor();
+        this.readEntryTimeoutMillis = getManagedLedgerConfig().getReadEntryTimeoutSeconds();
         this.entries = new RangeCache<>(EntryImpl::getLength, EntryImpl::getTimestamp);
         this.copyEntries = copyEntries;
 
@@ -68,11 +81,21 @@ public RangeEntryCacheImpl(RangeEntryCacheManagerImpl manager, ManagedLedgerImpl
         }
     }
 
+    @VisibleForTesting
+    ManagedLedgerConfig getManagedLedgerConfig() {
+        return ml.getConfig();
+    }
+
     @Override
     public String getName() {
         return ml.getName();
     }
 
+    @VisibleForTesting
+    InflightReadsLimiter getPendingReadsLimiter() {
+        return manager.getInflightReadsLimiter();
+    }
+
     public static final PooledByteBufAllocator ALLOCATOR = new PooledByteBufAllocator(true, // preferDirect
             0, // nHeapArenas,
             PooledByteBufAllocator.defaultNumDirectArena(), // nDirectArena
@@ -256,6 +279,19 @@ public void asyncReadEntry(ReadHandle lh, long firstEntry, long lastEntry, boole
     @SuppressWarnings({ "unchecked", "rawtypes" })
     private void asyncReadEntry0(ReadHandle lh, long firstEntry, long lastEntry, boolean isSlowestReader,
             final ReadEntriesCallback callback, Object ctx) {
+        asyncReadEntry0WithLimits(lh, firstEntry, lastEntry, isSlowestReader, callback, ctx, null);
+    }
+
+    void asyncReadEntry0WithLimits(ReadHandle lh, long firstEntry, long lastEntry, boolean isSlowestReader,
+        final ReadEntriesCallback originalCallback, Object ctx, InflightReadsLimiter.Handle handle) {
+
+        final AsyncCallbacks.ReadEntriesCallback callback =
+                handlePendingReadsLimits(lh, firstEntry, lastEntry, isSlowestReader,
+                        originalCallback, ctx, handle);
+        if (callback == null) {
+            return;
+        }
+
         final long ledgerId = lh.getId();
         final int entriesToRead = (int) (lastEntry - firstEntry) + 1;
         final PositionImpl firstPosition = PositionImpl.get(lh.getId(), firstEntry);
@@ -329,6 +365,77 @@ private void asyncReadEntry0(ReadHandle lh, long firstEntry, long lastEntry, boo
         }
     }
 
+
+    private AsyncCallbacks.ReadEntriesCallback handlePendingReadsLimits(ReadHandle lh,
+                                                                        long firstEntry, long lastEntry,
+                                                                        boolean shouldCacheEntry,
+                                                                        AsyncCallbacks.ReadEntriesCallback
+                                                                                originalCallback,
+                                                                        Object ctx,
+                                                                        InflightReadsLimiter.Handle handle) {
+        InflightReadsLimiter pendingReadsLimiter = getPendingReadsLimiter();
+        if (pendingReadsLimiter.isDisabled()) {
+            return originalCallback;
+        }
+        long estimatedReadSize = (1 + lastEntry - firstEntry)
+                * (estimatedEntrySize + BOOKKEEPER_READ_OVERHEAD_PER_ENTRY);
+        final AsyncCallbacks.ReadEntriesCallback callback;
+        InflightReadsLimiter.Handle newHandle = pendingReadsLimiter.acquire(estimatedReadSize, handle);
+        if (!newHandle.success) {
+            long now = System.currentTimeMillis();
+            if (now - newHandle.creationTime > readEntryTimeoutMillis) {
+                String message = "Time-out elapsed while acquiring enough permits "
+                        + "on the memory limiter to read from ledger "
+                        + lh.getId()
+                        + ", " + getName()
+                        + ", estimated read size " + estimatedReadSize + " bytes"
+                        + " for " + (1 + lastEntry - firstEntry)
+                        + " entries (check managedLedgerMaxReadsInFlightSizeInMB)";
+                log.error(message);
+                pendingReadsLimiter.release(newHandle);
+                originalCallback.readEntriesFailed(
+                        new ManagedLedgerException.TooManyRequestsException(message), ctx);
+                return null;
+            }
+            ml.getExecutor().submitOrdered(lh.getId(), () -> {
+                asyncReadEntry0WithLimits(lh, firstEntry, lastEntry, shouldCacheEntry,
+                        originalCallback, ctx, newHandle);
+                return null;
+            });
+            return null;
+        } else {
+            callback = new AsyncCallbacks.ReadEntriesCallback() {
+
+                @Override
+                public void readEntriesComplete(List<Entry> entries, Object ctx) {
+                    if (!entries.isEmpty()) {
+                        long size = entries.get(0).getLength();
+                        estimatedEntrySize = size;
+
+                        AtomicInteger remainingCount = new AtomicInteger(entries.size());
+                        for (Entry entry : entries) {
+                            ((EntryImpl) entry).onDeallocate(() -> {
+                                if (remainingCount.decrementAndGet() <= 0) {
+                                    pendingReadsLimiter.release(newHandle);
+                                }
+                            });
+                        }
+                    } else {
+                        pendingReadsLimiter.release(newHandle);
+                    }
+                    originalCallback.readEntriesComplete(entries, ctx);
+                }
+
+                @Override
+                public void readEntriesFailed(ManagedLedgerException exception, Object ctx) {
+                    pendingReadsLimiter.release(newHandle);
+                    originalCallback.readEntriesFailed(exception, ctx);
+                }
+            };
+        }
+        return callback;
+    }
+
     @Override
     public void clear() {
         Pair<Integer, Long> removedPair = entries.clear();