From eaf8dee1b8c90a2334102f4312dcf0b1bf167ed0 Mon Sep 17 00:00:00 2001
From: dao-jun <daojun@apache.org>
Date: Thu, 20 Feb 2025 18:19:49 +0800
Subject: [PATCH] Fix client can OOM when there are some bookies slow

---
 .../java/org/apache/bookkeeper/client/BookKeeper.java    | 1 +
 .../java/org/apache/bookkeeper/proto/BookieClient.java   | 9 +++++++++
 .../org/apache/bookkeeper/proto/BookieClientImpl.java    | 9 +++++++++
 3 files changed, 19 insertions(+)

diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java
index 40d6691390e..6a4d1232543 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java
@@ -640,6 +640,7 @@ void checkForFaultyBookies() {
         for (BookieId faultyBookie : faultyBookies) {
             if (Math.random() <= bookieQuarantineRatio) {
                 bookieWatcher.quarantineBookie(faultyBookie);
+                bookieClient.closeQuarantinedBookieChannels(faultyBookie);
                 statsLogger.getCounter(BookKeeperServerStats.BOOKIE_QUARANTINE).inc();
             } else {
                 statsLogger.getCounter(BookKeeperServerStats.BOOKIE_QUARANTINE_SKIP).inc();
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClient.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClient.java
index 5d20e1b22d1..ba435a8fd7e 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClient.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClient.java
@@ -281,4 +281,13 @@ CompletableFuture<AvailabilityOfEntriesOfLedger> getListOfEntriesOfLedger(Bookie
      * Close the bookie client object.
      */
     void close();
+
+    /**
+     * Close quarantined bookie channels to release byte buffers.
+     *
+     * @param bookieId bookie id
+     */
+    default void closeQuarantinedBookieChannels(BookieId bookieId) {
+        // no-op
+    }
 }
diff --git a/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClientImpl.java b/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClientImpl.java
index a12d9fd64d5..e7ec4ef8666 100644
--- a/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClientImpl.java
+++ b/bookkeeper-server/src/main/java/org/apache/bookkeeper/proto/BookieClientImpl.java
@@ -608,6 +608,15 @@ public void close() {
         }
     }
 
+
+    @Override
+    public void closeQuarantinedBookieChannels(BookieId bookieId) {
+        PerChannelBookieClientPool connectionPool = this.channels.remove(bookieId);
+        if (connectionPool != null) {
+            connectionPool.close(false);
+        }
+    }
+
     private static class Counter {
         int i;
         int total;