+ * 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;
+
+/**
+ * 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);
+
+}
diff --git a/bookkeeper-common/src/main/java/org/apache/bookkeeper/common/util/WriteMemoryCounter.java b/bookkeeper-common/src/main/java/org/apache/bookkeeper/common/util/WriteMemoryCounter.java
new file mode 100644
index 00000000000..aa199ced51b
--- /dev/null
+++ b/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
+ * 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.client;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.bookkeeper.client.api.BKException;
+import org.apache.bookkeeper.common.util.WritableListener;
+import org.apache.bookkeeper.test.BookKeeperClusterTestCase;
+import org.junit.Test;
+
+
+@Slf4j
+public class BookieClientMemoryCounterTest extends BookKeeperClusterTestCase {
+
+ static final int MESSAGE_SIZE = 1024;
+ static final long LOW_WATER_MARK = 10 * 1024;
+ static final long HIGH_WATER_MARK = 20 * 1024;
+
+ public BookieClientMemoryCounterTest() {
+ super(1);
+ baseClientConf.setWriteMemoryHighWaterMark(HIGH_WATER_MARK);
+ baseClientConf.setWriteMemoryLowWaterMark(LOW_WATER_MARK);
+ }
+
+ @Test
+ public void testPendingAddEntryMemory() throws Exception {
+ // listen to the write state change events
+ AtomicBoolean writeState = new AtomicBoolean(true);
+ bkc.getWriteMemoryCounter().register(new WritableListener() {
+ @Override
+ public void onWriteStateChanged(boolean writable) {
+ long usage = bkc.getWriteMemoryCounter().getSize();
+ log.info("Write state changed to {}, current memory usage is {}", writeState, usage);
+ // when the writable change to ture, the usage should under the LowWaterMark.
+ // when the writable change to false, the usage should over than the HighWaterMark.
+ if (writable) {
+ assertEquals(LOW_WATER_MARK - MESSAGE_SIZE, usage);
+ } else {
+ assertEquals(HIGH_WATER_MARK + MESSAGE_SIZE, usage);
+ }
+ writeState.set(writable);
+ }
+ });
+
+ LedgerHandle lh = bkc.createLedger(1, 1, BookKeeper.DigestType.CRC32, "".getBytes());
+ byte[] msg = new byte[MESSAGE_SIZE];
+
+ int testMessagesNum = 1000;
+
+ // start a thread to send message
+ AtomicInteger addCount = new AtomicInteger(testMessagesNum);
+ new Thread(() -> {
+ for (int i = 0; i < testMessagesNum; i++) {
+ while (!writeState.get()) {
+ log.info("wait for the memory released");
+ try {
+ TimeUnit.MILLISECONDS.sleep(10);
+ } catch (InterruptedException e) {
+ // ignore
+ }
+ }
+ lh.asyncAddEntry(msg, new AsyncCallback.AddCallback() {
+ @Override
+ public void addComplete(int rc, LedgerHandle lh, long entryId, Object ctx) {
+ if (rc == BKException.Code.OK) {
+ log.info("Add complete with rc {}", rc);
+ addCount.getAndDecrement();
+ }
+ }
+ }, null);
+ }
+ }).start();
+
+ // while sending messages, we listen on the memory counter size. The size should never over than the
+ // (highWaterMark + 1 message) bytes.
+ while (addCount.get() != 0) {
+ long size = bkc.getWriteMemoryCounter().getSize();
+ assertTrue(size >= 0 && size < baseClientConf.getWriteMemoryHighWaterMark() + msg.length + 1);
+ TimeUnit.MILLISECONDS.sleep(10);
+ }
+
+ lh.close();
+ }
+}
diff --git a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockBookKeeperTestCase.java b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockBookKeeperTestCase.java
index 2ee6feb3c8e..28a32e125e1 100644
--- a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockBookKeeperTestCase.java
+++ b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockBookKeeperTestCase.java
@@ -58,6 +58,7 @@
import org.apache.bookkeeper.client.api.OpenBuilder;
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.conf.ClientConfiguration;
import org.apache.bookkeeper.meta.LedgerIdGenerator;
import org.apache.bookkeeper.meta.LedgerManager;
@@ -170,6 +171,7 @@ public void setup() throws Exception {
when(bk.getMainWorkerPool()).thenReturn(executor);
when(bk.getBookieClient()).thenReturn(bookieClient);
when(bk.getScheduler()).thenReturn(scheduler);
+ when(bk.getWriteMemoryCounter()).thenCallRealMethod();
setBookKeeperConfig(new ClientConfiguration());
when(bk.getStatsLogger()).thenReturn(NullStatsLogger.INSTANCE);
@@ -224,7 +226,12 @@ public boolean isClientClosed() {
public ByteBufAllocator getByteBufAllocator() {
return UnpooledByteBufAllocator.DEFAULT;
}
- };
+
+ @Override
+ public WriteMemoryCounter getWriteMemoryCounter() {
+ return bk.getWriteMemoryCounter();
+ }
+ };
when(bk.getClientCtx()).thenReturn(clientCtx);
when(bk.getLedgerManager()).thenReturn(ledgerManager);
when(bk.getLedgerIdGenerator()).thenReturn(ledgerIdGenerator);
diff --git a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockClientContext.java b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockClientContext.java
index 93078a05129..5078eb79094 100644
--- a/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockClientContext.java
+++ b/bookkeeper-server/src/test/java/org/apache/bookkeeper/client/MockClientContext.java
@@ -27,6 +27,7 @@
import java.util.function.BooleanSupplier;
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.conf.ClientConfiguration;
import org.apache.bookkeeper.discover.MockRegistrationClient;
import org.apache.bookkeeper.meta.LedgerManager;
@@ -53,6 +54,7 @@ public class MockClientContext implements ClientContext {
private BooleanSupplier isClientClosed;
private MockRegistrationClient regClient;
private ByteBufAllocator allocator;
+ private WriteMemoryCounter writeMemoryCounter;
static MockClientContext create(MockBookies mockBookies) throws Exception {
ClientConfiguration conf = new ClientConfiguration();
@@ -64,7 +66,7 @@ static MockClientContext create(MockBookies mockBookies) throws Exception {
new DefaultBookieAddressResolver(regClient),
NullStatsLogger.INSTANCE);
bookieWatcherImpl.initialBlockingBookieRead();
-
+ WriteMemoryCounter memoryCounter = new WriteMemoryCounter();
return new MockClientContext()
.setConf(ClientInternalConf.fromConfig(conf))
.setLedgerManager(new MockLedgerManager())
@@ -76,7 +78,8 @@ static MockClientContext create(MockBookies mockBookies) throws Exception {
.setMainWorkerPool(scheduler)
.setScheduler(scheduler)
.setClientStats(BookKeeperClientStats.newInstance(NullStatsLogger.INSTANCE))
- .setIsClientClosed(() -> false);
+ .setIsClientClosed(() -> false)
+ .setWriteMemoryCounter(memoryCounter);
}
static MockClientContext create() throws Exception {
@@ -95,7 +98,8 @@ static MockClientContext copyOf(ClientContext other) {
.setScheduler(other.getScheduler())
.setClientStats(other.getClientStats())
.setByteBufAllocator(other.getByteBufAllocator())
- .setIsClientClosed(other::isClientClosed);
+ .setIsClientClosed(other::isClientClosed)
+ .setWriteMemoryCounter(other.getWriteMemoryCounter());
}
public MockRegistrationClient getMockRegistrationClient() {
@@ -168,6 +172,11 @@ public MockClientContext setByteBufAllocator(ByteBufAllocator allocator) {
return this;
}
+ public MockClientContext setWriteMemoryCounter(WriteMemoryCounter writeMemoryCounter) {
+ this.writeMemoryCounter = writeMemoryCounter;
+ return this;
+ }
+
private static