From 70d086f8f35d36800059d0d68e13d0ca017bf233 Mon Sep 17 00:00:00 2001
From: Lari Hotari <lhotari@users.noreply.github.com>
Date: Sat, 7 Oct 2023 05:26:48 +0300
Subject: [PATCH] [fix][ml] Fix thread safe issue with RangeCache.put and
 RangeCache.clear (#21302)

---
 .../bookkeeper/mledger/util/RangeCache.java   | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/util/RangeCache.java b/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/util/RangeCache.java
index 7599e2cc1874f..d34857e5e5177 100644
--- a/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/util/RangeCache.java
+++ b/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/util/RangeCache.java
@@ -27,7 +27,6 @@
 import java.util.concurrent.ConcurrentNavigableMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.atomic.AtomicLong;
-import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.apache.commons.lang3.tuple.Pair;
 
 /**
@@ -74,13 +73,18 @@ public RangeCache(Weighter<Value> weighter, TimestampExtractor<Value> timestampE
      * @return whether the entry was inserted in the cache
      */
     public boolean put(Key key, Value value) {
-        MutableBoolean flag = new MutableBoolean();
-        entries.computeIfAbsent(key, (k) -> {
-            size.addAndGet(weighter.getSize(value));
-            flag.setValue(true);
-            return value;
-        });
-        return flag.booleanValue();
+        // retain value so that it's not released before we put it in the cache and calculate the weight
+        value.retain();
+        try {
+            if (entries.putIfAbsent(key, value) == null) {
+                size.addAndGet(weighter.getSize(value));
+                return true;
+            } else {
+                return false;
+            }
+        } finally {
+            value.release();
+        }
     }
 
     public boolean exists(Key key) {
@@ -242,7 +246,6 @@ public synchronized Pair<Integer, Long> clear() {
             value.release();
         }
 
-        entries.clear();
         size.getAndAdd(-removedSize);
         return Pair.of(removedCount, removedSize);
     }