fix(s3stream): fix issues402 memory leak (#403)

Signed-off-by: Robin Han <hanxvdovehx@gmail.com>
AutoMQ · Oct 20, 2023 · 47cf183 · 47cf183
1 parent 6a0cea8
commit 47cf183
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 11 deletions.
diff --git a/s3stream/pom.xml b/s3stream/pom.xml
@@ -22,7 +22,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>com.automq.elasticstream</groupId>
     <artifactId>s3stream</artifactId>
-    <version>0.1.12-SNAPSHOT</version>
+    <version>0.1.13-SNAPSHOT</version>
     <properties>
         <mockito-core.version>5.5.0</mockito-core.version>
         <junit-jupiter.version>5.10.0</junit-jupiter.version>

diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/BlockCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/BlockCache.java
@@ -224,9 +224,13 @@ private void ensureCapacity(int size) {
                     continue;
                 }
                 CacheBlock cacheBlock = streamCache.blocks.remove(entry.getKey().startOffset);
-                cacheBlock.free();
-                if (maxSize - this.size.addAndGet(-cacheBlock.size) >= size) {
-                    return;
+                if (cacheBlock == null) {
+                    LOGGER.error("[BUG] Cannot find stream cache block: {} {}", entry.getKey().streamId, entry.getKey().startOffset);
+                } else {
+                    cacheBlock.free();
+                    if (maxSize - this.size.addAndGet(-cacheBlock.size) >= size) {
+                        return;
+                    }
                 }
             }
         }

diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/DataBlockReadAccumulator.java b/s3stream/src/main/java/com/automq/stream/s3/cache/DataBlockReadAccumulator.java
@@ -18,7 +18,10 @@
 package com.automq.stream.s3.cache;
 
 import com.automq.stream.s3.ObjectReader;
+import com.automq.stream.utils.FutureUtil;
 import org.apache.commons.lang3.tuple.Pair;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -30,6 +33,7 @@
  * Accumulate inflight data block read requests to one real read request.
  */
 public class DataBlockReadAccumulator {
+    private static final Logger LOGGER = LoggerFactory.getLogger(DataBlockReadAccumulator.class);
     private final Map<Pair<String, Integer>, DataBlockRecords> inflightDataBlockReads = new ConcurrentHashMap<>();
     private final Consumer<DataBlockRecords> dataBlockConsumer;
 
@@ -60,7 +64,10 @@ public CompletableFuture<DataBlockRecords> readDataBlock(ObjectReader reader, Ob
                             inflightDataBlockReads.remove(key, finalRecords);
                         }
                         finalRecords.complete(dataBlock, ex);
-                        dataBlockConsumer.accept(finalRecords);
+                        FutureUtil.suppress(() -> dataBlockConsumer.accept(finalRecords), LOGGER);
+                    } catch (Throwable e) {
+                        LOGGER.error("[UNEXPECTED] DataBlockRecords fail to notify listener {}", listener, e);
+                    } finally {
                         finalRecords.release();
                     }
                 });

diff --git a/s3stream/src/main/java/com/automq/stream/s3/cache/LRUCache.java b/s3stream/src/main/java/com/automq/stream/s3/cache/LRUCache.java
@@ -31,21 +31,21 @@ public LRUCache() {
         cacheEntrySet = cache.entrySet();
     }
 
-    public boolean touch(K key) {
+    public synchronized boolean touch(K key) {
         return cache.get(key) != null;
     }
 
-    public void put(K key, V value) {
+    public synchronized void put(K key, V value) {
         if (cache.put(key, value) != null) {
             touch(key);
         }
     }
 
-    public V get(K key) {
+    public synchronized V get(K key) {
         return cache.get(key);
     }
 
-    public Map.Entry<K, V> pop() {
+    public synchronized Map.Entry<K, V> pop() {
         Iterator<Map.Entry<K, V>> it = cacheEntrySet.iterator();
         if (!it.hasNext()) {
             return null;
@@ -58,11 +58,11 @@ public Map.Entry<K, V> pop() {
         return entry;
     }
 
-    public boolean remove(K key) {
+    public synchronized boolean remove(K key) {
         return cache.remove(key) != null;
     }
 
-    public int size() {
+    public synchronized int size() {
         return cache.size();
     }
 }