Add memory limiter for the add entry request to avoid OOM

bookkeeper-common/src/main/java/org/apache/bookkeeper/common/util/WritableListener.java

@@ -0,0 +1,31 @@
+/**
+ * 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
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.common.util;
+
+/**
+ * WritableListener used to listen the writable status changes.
+ *
+ * We use {@link WriteMemoryCounter} to listen on the memory usage when the client adds entries. The listener can
+ * take actions if they have been notified.
+ */
+public interface WritableListener {
+
+    void onWriteStateChanged(boolean writable);
+
+}

bookkeeper-common/src/main/java/org/apache/bookkeeper/common/util/WriteMemoryCounter.java

@@ -0,0 +1,80 @@
+/**
+ * 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.common.util;
+
+import java.util.LinkedList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * {@link WriteMemoryCounter} counts the memory usage on Adds request.
+ * When there has an Add request created, the {@link WriteMemoryCounter} will record the request content size.
+ * When the request is finished, the {@link WriteMemoryCounter} will decrease the record count.
+ * The range of the counter should in the range of {@link WriteWaterMark}'s high watermark and low watermark.
+ *
+ * If the record size is over to the high watermark, the registered listeners will receive writable state change
+ * to false and take actions. The listeners will receive writable state change to true until the record size is
+ * down to the low watermark.
+ */
+@Slf4j
+public class WriteMemoryCounter {
+    private final WriteWaterMark writeWaterMark;
+    private AtomicLong sizeCounter = new AtomicLong(0);
+    private AtomicBoolean writeState = new AtomicBoolean(true);
+    private final List<WritableListener> listeners = new LinkedList<>();
+
+    public WriteMemoryCounter(WriteWaterMark writeWaterMark) {
+        this.writeWaterMark = writeWaterMark;
+    }
+
+    public WriteMemoryCounter() {
+        this.writeWaterMark = new WriteWaterMark();
+    }
+
+    public void register(WritableListener listener) {
+        listeners.add(listener);
+    }
+
+    public void incrementPendingWriteBytes(long size) {
+        long usage = sizeCounter.addAndGet(size);
+        log.info("increment the size to {}", usage);
+        if (usage > writeWaterMark.high() && writeState.get()) {
+            setWritable(false);
+        }
+    }
+
+    public void decrementPendingWriteBytes(long size) {
+        long usage = sizeCounter.addAndGet(-size);
+        log.info("decrement the size to {}", usage);
+        if (usage < writeWaterMark.low() && !writeState.get()) {
+            setWritable(true);
+        }
+    }
+
+    public void setWritable(boolean state) {
+        writeState.set(state);
+        listeners.forEach(l -> l.onWriteStateChanged(state));
+    }
+
+    public long getSize() {
+        return sizeCounter.get();
+    }
+}

bookkeeper-common/src/main/java/org/apache/bookkeeper/common/util/WriteWaterMark.java

@@ -0,0 +1,48 @@
+/**
+ * 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.common.util;
+
+/**
+ * {@link WriteWaterMark} is used to configure the max value and the min value of the memory usage.
+ */
+public class WriteWaterMark {
+    private static final long DEFAULT_LOW_WATER_MARK = 1610612736L; // 1.5GB
+    private static final long DEFAULT_HIGH_WATER_MARK = 2147483648L; // 2GB
+
+    private final long low;
+    private final long high;
+
+    public WriteWaterMark(long low, long high) {
+        this.low = low;
+        this.high = high;
+    }
+
+    public WriteWaterMark() {
+        this.low = DEFAULT_LOW_WATER_MARK;
+        this.high = DEFAULT_HIGH_WATER_MARK;
+    }
+
+    public long low() {
+        return low;
+    }
+
+    public long high() {
+        return high;
+    }
+}

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java

@@ -67,6 +67,8 @@
 import org.apache.bookkeeper.common.util.OrderedExecutor;
 import org.apache.bookkeeper.common.util.OrderedScheduler;
 import org.apache.bookkeeper.common.util.ReflectionUtils;
+import org.apache.bookkeeper.common.util.WriteMemoryCounter;
+import org.apache.bookkeeper.common.util.WriteWaterMark;
 import org.apache.bookkeeper.conf.AbstractConfiguration;
 import org.apache.bookkeeper.conf.ClientConfiguration;
 import org.apache.bookkeeper.feature.FeatureProvider;
@@ -151,6 +153,8 @@ public class BookKeeper implements org.apache.bookkeeper.client.api.BookKeeper {
     boolean closed = false;
     final ReentrantReadWriteLock closeLock = new ReentrantReadWriteLock();
 
+    final WriteMemoryCounter writeMemoryCounter;
+
     /**
      * BookKeeper Client Builder to build client instances.
      *
@@ -538,6 +542,10 @@ public BookKeeper(ClientConfiguration conf, ZooKeeper zk, EventLoopGroup eventLo
         this.ledgerIdGenerator = ledgerManagerFactory.newLedgerIdGenerator();
 
         this.bookieQuarantineRatio = conf.getBookieQuarantineRatio();
+
+        this.writeMemoryCounter = new WriteMemoryCounter(
+            new WriteWaterMark(conf.getWriteMemoryLowWaterMark(), conf.getWriteMemoryHighWaterMark()));
+
         scheduleBookieHealthCheckIfEnabled(conf);
     }
 
@@ -566,6 +574,7 @@ public BookKeeper(ClientConfiguration conf, ZooKeeper zk, EventLoopGroup eventLo
         bookieClient = null;
         allocator = UnpooledByteBufAllocator.DEFAULT;
         bookieQuarantineRatio = 1.0;
+        writeMemoryCounter = null;
     }
 
     protected EnsemblePlacementPolicy initializeEnsemblePlacementPolicy(ClientConfiguration conf,
@@ -705,6 +714,10 @@ public MetadataClientDriver getMetadataClientDriver() {
         return metadataDriver;
     }
 
+    WriteMemoryCounter getWriteMemoryCounter() {
+        return writeMemoryCounter;
+    }
+
     /**
      * There are 3 digest types that can be used for verification. The CRC32 is
      * cheap to compute but does not protect against byzantine bookies (i.e., a
@@ -1662,6 +1675,11 @@ public boolean isClientClosed() {
             public ByteBufAllocator getByteBufAllocator() {
                 return allocator;
             }
+
+            @Override
+            public WriteMemoryCounter getWriteMemoryCounter() {
+                return writeMemoryCounter;
+            }
         };
 
     public ClientContext getClientCtx() {

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/ClientContext.java

@@ -23,6 +23,7 @@
 import io.netty.buffer.ByteBufAllocator;
 import org.apache.bookkeeper.common.util.OrderedExecutor;
 import org.apache.bookkeeper.common.util.OrderedScheduler;
+import org.apache.bookkeeper.common.util.WriteMemoryCounter;
 import org.apache.bookkeeper.meta.LedgerManager;
 import org.apache.bookkeeper.proto.BookieClient;
 
@@ -43,4 +44,6 @@ public interface ClientContext {
     OrderedScheduler getScheduler();
     BookKeeperClientStats getClientStats();
     boolean isClientClosed();
+
+    WriteMemoryCounter getWriteMemoryCounter();
 }

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/PendingAddOp.java

@@ -96,7 +96,7 @@ static PendingAddOp create(LedgerHandle lh, ClientContext clientCtx,
         op.currentLedgerLength = -1;
         op.payload = payload;
         op.entryLength = payload.readableBytes();
-
+        op.clientCtx.getWriteMemoryCounter().incrementPendingWriteBytes(op.entryLength);
         op.completed = false;
         op.ensemble = ensemble;
         op.ackSet = lh.getDistributionSchedule().getAckSet();
@@ -493,6 +493,7 @@ private void maybeRecycle() {
         }
         // only recycle a pending add op after it has been run.
         if (hasRun && toSend == null && pendingWriteRequests == 0) {
+            clientCtx.getWriteMemoryCounter().decrementPendingWriteBytes(entryLength);
             recyclePendAddOpObject();
         }
     }

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/api/WriteHandle.java

@@ -141,6 +141,15 @@ default long append(byte[] data, int offset, int length) throws BKException, Int
      */
     long getLastAddPushed();
 
+    /**
+     *
+     *
+     * @return
+     */
+    default boolean isWritable() {
+        throw new UnsupportedOperationException("This operation is not supported for the current handler");
+    }
+
     /**
      * Asynchronous close the write handle, any adds in flight will return errors.
      *

bookkeeper-server/src/main/java/org/apache/bookkeeper/conf/ClientConfiguration.java

@@ -199,6 +199,9 @@ public class ClientConfiguration extends AbstractConfiguration<ClientConfigurati
     protected static final String CLIENT_CONNECT_BOOKIE_UNAVAILABLE_LOG_THROTTLING =
             "clientConnectBookieUnavailableLogThrottling";
 
+    protected static final String WRITE_MEMORY_LOW_WATER_MARK = "writeMemoryLowWaterMark";
+    protected static final String WRITE_MEMORY_HIGH_WATER_MARK = "writeMemoryHighWaterMark";
+
     /**
      * Construct a default client-side configuration.
      */
@@ -2058,6 +2061,24 @@ public long getClientConnectBookieUnavailableLogThrottlingMs() {
         return getLong(CLIENT_CONNECT_BOOKIE_UNAVAILABLE_LOG_THROTTLING, 5_000L);
     }
 
+    public ClientConfiguration setWriteMemoryLowWaterMark(long bytes) {
+        setProperty(WRITE_MEMORY_LOW_WATER_MARK, bytes);
+        return this;
+    }
+
+    public long getWriteMemoryLowWaterMark() {
+        return getInt(WRITE_MEMORY_LOW_WATER_MARK, 64 * 1024 * 1024);
+    }
+
+    public ClientConfiguration setWriteMemoryHighWaterMark(long bytes) {
+        setProperty(WRITE_MEMORY_HIGH_WATER_MARK, bytes);
+        return this;
+    }
+
+    public long getWriteMemoryHighWaterMark() {
+        return getInt(WRITE_MEMORY_HIGH_WATER_MARK, 256 * 1024 * 1024);
+    }
+
     @Override
     protected ClientConfiguration getThis() {
         return this;