From 88f34bba12e9f27c78db2ce468b4d46619a92853 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Tue, 28 May 2024 09:33:13 +0300 Subject: [PATCH 1/7] [fix][ml] Fix race conditions in RangeCache --- .../bookkeeper/mledger/impl/EntryImpl.java | 7 +- .../bookkeeper/mledger/util/RangeCache.java | 134 ++++++++++++++---- .../mledger/util/RangeCacheTest.java | 49 ++++--- 3 files changed, 143 insertions(+), 47 deletions(-) diff --git a/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/EntryImpl.java b/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/EntryImpl.java index 803979313575a..48a79a4ac529c 100644 --- a/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/EntryImpl.java +++ b/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/EntryImpl.java @@ -27,9 +27,10 @@ import org.apache.bookkeeper.client.api.LedgerEntry; import org.apache.bookkeeper.mledger.Entry; import org.apache.bookkeeper.mledger.util.AbstractCASReferenceCounted; +import org.apache.bookkeeper.mledger.util.RangeCache; public final class EntryImpl extends AbstractCASReferenceCounted implements Entry, Comparable, - ReferenceCounted { + RangeCache.ValueWithKeyValidation { private static final Recycler RECYCLER = new Recycler() { @Override @@ -205,4 +206,8 @@ protected void deallocate() { recyclerHandle.recycle(this); } + @Override + public boolean matchesKey(PositionImpl key) { + return key.compareTo(ledgerId, entryId) == 0; + } } 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 d34857e5e5177..389f7f8b64eb7 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 @@ -19,6 +19,7 @@ package org.apache.bookkeeper.mledger.util; import static com.google.common.base.Preconditions.checkArgument; +import io.netty.util.IllegalReferenceCountException; import io.netty.util.ReferenceCounted; import java.util.ArrayList; import java.util.Collection; @@ -27,6 +28,7 @@ import java.util.concurrent.ConcurrentNavigableMap; import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.atomic.AtomicLong; +import org.apache.bookkeeper.mledger.util.RangeCache.ValueWithKeyValidation; import org.apache.commons.lang3.tuple.Pair; /** @@ -37,7 +39,11 @@ * @param * Cache value */ -public class RangeCache, Value extends ReferenceCounted> { +public class RangeCache, Value extends ValueWithKeyValidation> { + public interface ValueWithKeyValidation extends ReferenceCounted { + boolean matchesKey(T key); + } + // Map from key to nodes inside the linked list private final ConcurrentNavigableMap entries; private AtomicLong size; // Total size of values stored in cache @@ -76,6 +82,9 @@ public boolean put(Key key, Value value) { // retain value so that it's not released before we put it in the cache and calculate the weight value.retain(); try { + if (!value.matchesKey(key)) { + throw new IllegalArgumentException("Value '" + value + "' does not match key '" + key + "'"); + } if (entries.putIfAbsent(key, value) == null) { size.addAndGet(weighter.getSize(value)); return true; @@ -98,9 +107,15 @@ public Value get(Key key) { } else { try { value.retain(); + } catch (IllegalReferenceCountException e) { + // Value was already deallocated + return null; + } + if (value.matchesKey(key)) { return value; - } catch (Throwable t) { - // Value was already destroyed between get() and retain() + } else { + // Value was recycled in between + value.release(); return null; } } @@ -118,12 +133,20 @@ public Collection getRange(Key first, Key last) { List values = new ArrayList(); // Return the values of the entries found in cache - for (Value value : entries.subMap(first, true, last, true).values()) { + for (Map.Entry entry : entries.subMap(first, true, last, true).entrySet()) { + Key key = entry.getKey(); + Value value = entry.getValue(); try { value.retain(); + } catch (IllegalReferenceCountException e) { + // Value was already deallocated + continue; + } + if (value.matchesKey(key)) { values.add(value); - } catch (Throwable t) { - // Value was already destroyed between get() and retain() + } else { + // Value was recycled in between + value.release(); } } @@ -143,15 +166,28 @@ public Pair removeRange(Key first, Key last, boolean lastInclusiv int removedEntries = 0; long removedSize = 0; - for (Key key : subMap.keySet()) { - Value value = entries.remove(key); + for (Map.Entry entry : subMap.entrySet()) { + Key key = entry.getKey(); + Value value = entry.getValue(); if (value == null) { continue; } - - removedSize += weighter.getSize(value); + try { + // add extra retain to avoid value being released while we are removing it + value.retain(); + } catch (IllegalReferenceCountException e) { + // Value was already released + continue; + } + // check that the value hasn't been recycled in between + if (value.matchesKey(key) && entries.remove(key, value)) { + removedSize += weighter.getSize(value); + // remove the cache reference + value.release(); + ++removedEntries; + } + // remove the extra retain value.release(); - ++removedEntries; } size.addAndGet(-removedSize); @@ -171,14 +207,27 @@ public Pair evictLeastAccessedEntries(long minSize) { int removedEntries = 0; while (removedSize < minSize) { - Map.Entry entry = entries.pollFirstEntry(); + Map.Entry entry = entries.firstEntry(); if (entry == null) { break; } - + Key key = entry.getKey(); Value value = entry.getValue(); - ++removedEntries; - removedSize += weighter.getSize(value); + try { + // add extra retain to avoid value being released while we are removing it + value.retain(); + } catch (IllegalReferenceCountException e) { + // Value was already released + continue; + } + // check that the value hasn't been recycled in between + if (value.matchesKey(key) && entries.remove(key, value)) { + ++removedEntries; + removedSize += weighter.getSize(value); + // release the cache reference + value.release(); + } + // remove the extra retain value.release(); } @@ -197,18 +246,37 @@ public Pair evictLEntriesBeforeTimestamp(long maxTimestamp) { while (true) { Map.Entry entry = entries.firstEntry(); - if (entry == null || timestampExtractor.getTimestamp(entry.getValue()) > maxTimestamp) { + if (entry == null) { break; } + Key key = entry.getKey(); Value value = entry.getValue(); - boolean removeHits = entries.remove(entry.getKey(), value); - if (!removeHits) { - break; + try { + // add extra retain to avoid value being released while we are removing it + value.retain(); + } catch (IllegalReferenceCountException e) { + // Value was already released + continue; + } + try { + // check that the value hasn't been recycled in between + if (!value.matchesKey(key)) { + continue; + } + if (timestampExtractor.getTimestamp(value) > maxTimestamp) { + break; + } + boolean removeHits = entries.remove(key, value); + if (!removeHits) { + break; + } + removedSize += weighter.getSize(value); + removedCount++; + value.release(); + } finally { + // remove the extra retain + value.release(); } - - removedSize += weighter.getSize(value); - removedCount++; - value.release(); } size.addAndGet(-removedSize); @@ -236,13 +304,27 @@ public synchronized Pair clear() { int removedCount = 0; while (true) { - Map.Entry entry = entries.pollFirstEntry(); + Map.Entry entry = entries.firstEntry(); if (entry == null) { break; } + Key key = entry.getKey(); Value value = entry.getValue(); - removedSize += weighter.getSize(value); - removedCount++; + try { + // add extra retain to avoid value being released while we are removing it + value.retain(); + } catch (IllegalReferenceCountException e) { + // Value was already released + continue; + } + // check that the value hasn't been recycled in between + if (value.matchesKey(key) && entries.remove(key, value)) { + removedSize += weighter.getSize(value); + removedCount++; + // release the cache reference + value.release(); + } + // remove the extra retain value.release(); } diff --git a/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java b/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java index 8ce0db4ac4caa..489ed55dfc200 100644 --- a/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java +++ b/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java @@ -23,25 +23,29 @@ import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; - import com.google.common.collect.Lists; import io.netty.util.AbstractReferenceCounted; import io.netty.util.ReferenceCounted; -import org.apache.commons.lang3.tuple.Pair; -import org.testng.annotations.Test; -import java.util.UUID; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import org.apache.commons.lang3.tuple.Pair; +import org.testng.annotations.Test; public class RangeCacheTest { - class RefString extends AbstractReferenceCounted implements ReferenceCounted { + class RefString extends AbstractReferenceCounted implements RangeCache.ValueWithKeyValidation { String s; + Integer matchingKey; RefString(String s) { + this(s, null); + } + + RefString(String s, Integer matchingKey) { super(); this.s = s; + this.matchingKey = matchingKey != null ? matchingKey : Integer.parseInt(s); setRefCnt(1); } @@ -65,6 +69,11 @@ public boolean equals(Object obj) { return false; } + + @Override + public boolean matchesKey(Integer key) { + return matchingKey.equals(key); + } } @Test @@ -119,8 +128,8 @@ public void simple() { public void customWeighter() { RangeCache cache = new RangeCache<>(value -> value.s.length(), x -> 0); - cache.put(0, new RefString("zero")); - cache.put(1, new RefString("one")); + cache.put(0, new RefString("zero", 0)); + cache.put(1, new RefString("one", 1)); assertEquals(cache.getSize(), 7); assertEquals(cache.getNumberOfEntries(), 2); @@ -132,9 +141,9 @@ public void customTimeExtraction() { RangeCache cache = new RangeCache<>(value -> value.s.length(), x -> x.s.length()); cache.put(1, new RefString("1")); - cache.put(2, new RefString("22")); - cache.put(3, new RefString("333")); - cache.put(4, new RefString("4444")); + cache.put(22, new RefString("22")); + cache.put(333, new RefString("333")); + cache.put(4444, new RefString("4444")); assertEquals(cache.getSize(), 10); assertEquals(cache.getNumberOfEntries(), 4); @@ -151,12 +160,12 @@ public void customTimeExtraction() { public void doubleInsert() { RangeCache cache = new RangeCache<>(); - RefString s0 = new RefString("zero"); + RefString s0 = new RefString("zero", 0); assertEquals(s0.refCnt(), 1); assertTrue(cache.put(0, s0)); assertEquals(s0.refCnt(), 1); - cache.put(1, new RefString("one")); + cache.put(1, new RefString("one", 1)); assertEquals(cache.getSize(), 2); assertEquals(cache.getNumberOfEntries(), 2); @@ -164,7 +173,7 @@ public void doubleInsert() { assertEquals(s.s, "one"); assertEquals(s.refCnt(), 2); - RefString s1 = new RefString("uno"); + RefString s1 = new RefString("uno", 1); assertEquals(s1.refCnt(), 1); assertFalse(cache.put(1, s1)); assertEquals(s1.refCnt(), 1); @@ -201,10 +210,10 @@ public void getRange() { public void eviction() { RangeCache cache = new RangeCache<>(value -> value.s.length(), x -> 0); - cache.put(0, new RefString("zero")); - cache.put(1, new RefString("one")); - cache.put(2, new RefString("two")); - cache.put(3, new RefString("three")); + cache.put(0, new RefString("zero", 0)); + cache.put(1, new RefString("one", 1)); + cache.put(2, new RefString("two", 2)); + cache.put(3, new RefString("three", 3)); // This should remove the LRU entries: 0, 1 whose combined size is 7 assertEquals(cache.evictLeastAccessedEntries(5), Pair.of(2, (long) 7)); @@ -277,11 +286,11 @@ public void evictions() { @Test public void testInParallel() { - RangeCache cache = new RangeCache<>(value -> value.s.length(), x -> 0); + RangeCache cache = new RangeCache<>(value -> value.s.length(), x -> 0); ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(); executor.scheduleWithFixedDelay(cache::clear, 10, 10, TimeUnit.MILLISECONDS); for (int i = 0; i < 1000; i++) { - cache.put(UUID.randomUUID().toString(), new RefString("zero")); + cache.put(i, new RefString(String.valueOf(i))); } executor.shutdown(); } @@ -289,7 +298,7 @@ public void testInParallel() { @Test public void testPutSameObj() { RangeCache cache = new RangeCache<>(value -> value.s.length(), x -> 0); - RefString s0 = new RefString("zero"); + RefString s0 = new RefString("zero", 0); assertEquals(s0.refCnt(), 1); assertTrue(cache.put(0, s0)); assertFalse(cache.put(0, s0)); From 172de1656a0d49f9a2b19e56fd0bbce72b57c489 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Thu, 30 May 2024 23:47:39 +0300 Subject: [PATCH 2/7] Refactor RangeCache and cover yet another race condition --- .../bookkeeper/mledger/util/RangeCache.java | 287 ++++++++++-------- 1 file changed, 160 insertions(+), 127 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 389f7f8b64eb7..95045ac9983ef 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 @@ -19,12 +19,16 @@ package org.apache.bookkeeper.mledger.util; import static com.google.common.base.Preconditions.checkArgument; +import com.google.common.base.Predicate; import io.netty.util.IllegalReferenceCountException; +import io.netty.util.Recycler; +import io.netty.util.Recycler.Handle; import io.netty.util.ReferenceCounted; import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.concurrent.ConcurrentNavigableMap; import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.atomic.AtomicLong; @@ -45,11 +49,94 @@ public interface ValueWithKeyValidation extends ReferenceCounted { } // Map from key to nodes inside the linked list - private final ConcurrentNavigableMap entries; + private final ConcurrentNavigableMap> entries; private AtomicLong size; // Total size of values stored in cache private final Weighter weighter; // Weighter object used to extract the size from values private final TimestampExtractor timestampExtractor; // Extract the timestamp associated with a value + /** + * Wrapper around the value to store in Map. This is needed to ensure that a specific instance can be removed from + * the map by calling the {@link Map#remove(Object, Object)} method. Certain race conditions could result in the + * wrong value being removed from the map. The instances of this class are recycled to avoid creating new objects. + */ + private static class IdentityWrapper { + private final Handle recyclerHandle; + private static final Recycler RECYCLER = new Recycler() { + @Override + protected IdentityWrapper newObject(Handle recyclerHandle) { + return new IdentityWrapper(recyclerHandle); + } + }; + private T value; + + private IdentityWrapper(Handle recyclerHandle) { + this.recyclerHandle = recyclerHandle; + } + + static IdentityWrapper create(T value) { + IdentityWrapper identityWrapper = RECYCLER.get(); + identityWrapper.value = value; + return identityWrapper; + } + + T get() { + return value; + } + + void recycle() { + value = null; + recyclerHandle.recycle(this); + } + + @Override + public boolean equals(Object o) { + // only match exact identity of the value + return this == o; + } + + @Override + public int hashCode() { + return Objects.hashCode(value); + } + } + + /** + * Mutable object to store the number of entries and the total size removed from the cache. The instances + * are recycled to avoid creating new instances. + */ + private static class RemovalCounters { + private final Handle recyclerHandle; + private static final Recycler RECYCLER = new Recycler() { + @Override + protected RemovalCounters newObject(Handle recyclerHandle) { + return new RemovalCounters(recyclerHandle); + } + }; + int removedEntries; + long removedSize; + private RemovalCounters(Handle recyclerHandle) { + this.recyclerHandle = recyclerHandle; + } + + static RemovalCounters create() { + RemovalCounters results = RECYCLER.get(); + results.removedEntries = 0; + results.removedSize = 0; + return results; + } + + void recycle() { + removedEntries = 0; + removedSize = 0; + recyclerHandle.recycle(this); + } + + public void entryRemoved(long size) { + removedSize += size; + removedEntries++; + } + } + /** * Construct a new RangeLruCache with default Weighter. */ @@ -74,8 +161,7 @@ public RangeCache(Weighter weighter, TimestampExtractor timestampE * Insert. * * @param key - * @param value - * ref counted value with at least 1 ref to pass on the cache + * @param value ref counted value with at least 1 ref to pass on the cache * @return whether the entry was inserted in the cache */ public boolean put(Key key, Value value) { @@ -85,10 +171,13 @@ public boolean put(Key key, Value value) { if (!value.matchesKey(key)) { throw new IllegalArgumentException("Value '" + value + "' does not match key '" + key + "'"); } - if (entries.putIfAbsent(key, value) == null) { + IdentityWrapper newWrapper = IdentityWrapper.create(value); + if (entries.putIfAbsent(key, newWrapper) == null) { size.addAndGet(weighter.getSize(value)); return true; } else { + // recycle the new wrapper as it was not used + newWrapper.recycle(); return false; } } finally { @@ -101,10 +190,14 @@ public boolean exists(Key key) { } public Value get(Key key) { - Value value = entries.get(key); - if (value == null) { + return getValue(key, entries.get(key)); + } + + private Value getValue(Key key, IdentityWrapper valueWrapper) { + if (valueWrapper == null) { return null; } else { + Value value = valueWrapper.get(); try { value.retain(); } catch (IllegalReferenceCountException e) { @@ -114,7 +207,7 @@ public Value get(Key key) { if (value.matchesKey(key)) { return value; } else { - // Value was recycled in between + // Value or IdentityWrapper was recycled and already contains another value value.release(); return null; } @@ -133,20 +226,10 @@ public Collection getRange(Key first, Key last) { List values = new ArrayList(); // Return the values of the entries found in cache - for (Map.Entry entry : entries.subMap(first, true, last, true).entrySet()) { - Key key = entry.getKey(); - Value value = entry.getValue(); - try { - value.retain(); - } catch (IllegalReferenceCountException e) { - // Value was already deallocated - continue; - } - if (value.matchesKey(key)) { + for (Map.Entry> entry : entries.subMap(first, true, last, true).entrySet()) { + Value value = getValue(entry.getKey(), entry.getValue()); + if (value != null) { values.add(value); - } else { - // Value was recycled in between - value.release(); } } @@ -161,38 +244,59 @@ public Collection getRange(Key first, Key last) { * @return an pair of ints, containing the number of removed entries and the total size */ public Pair removeRange(Key first, Key last, boolean lastInclusive) { - Map subMap = entries.subMap(first, true, last, lastInclusive); + RemovalCounters counters = RemovalCounters.create(); + Map> subMap = entries.subMap(first, true, last, lastInclusive); + for (Map.Entry> entry : subMap.entrySet()) { + removeEntry(entry, counters); + } + return handleRemovalResult(counters); + } - int removedEntries = 0; - long removedSize = 0; + enum RemoveEntryResult { + ENTRY_REMOVED, + CONTINUE_LOOP, + BREAK_LOOP; + } - for (Map.Entry entry : subMap.entrySet()) { - Key key = entry.getKey(); - Value value = entry.getValue(); - if (value == null) { - continue; - } - try { - // add extra retain to avoid value being released while we are removing it - value.retain(); - } catch (IllegalReferenceCountException e) { - // Value was already released - continue; + private RemoveEntryResult removeEntry(Map.Entry> entry, RemovalCounters counters) { + return removeEntry(entry, counters, (x) -> true); + } + + private RemoveEntryResult removeEntry(Map.Entry> entry, RemovalCounters counters, + Predicate removeCondition) { + Key key = entry.getKey(); + IdentityWrapper identityWrapper = entry.getValue(); + Value value = identityWrapper.get(); + try { + // add extra retain to avoid value being released while we are removing it + value.retain(); + } catch (IllegalReferenceCountException e) { + // Value was already released + return RemoveEntryResult.CONTINUE_LOOP; + } + try { + if (!removeCondition.test(value)) { + return RemoveEntryResult.BREAK_LOOP; } // check that the value hasn't been recycled in between - if (value.matchesKey(key) && entries.remove(key, value)) { - removedSize += weighter.getSize(value); + if (value.matchesKey(key) && entries.remove(key, identityWrapper)) { + identityWrapper.recycle(); + counters.entryRemoved(weighter.getSize(value)); // remove the cache reference value.release(); - ++removedEntries; } + } finally { // remove the extra retain value.release(); } + return RemoveEntryResult.ENTRY_REMOVED; + } - size.addAndGet(-removedSize); - - return Pair.of(removedEntries, removedSize); + private Pair handleRemovalResult(RemovalCounters counters) { + size.addAndGet(-counters.removedSize); + Pair result = Pair.of(counters.removedEntries, counters.removedSize); + counters.recycle(); + return result; } /** @@ -202,37 +306,15 @@ public Pair removeRange(Key first, Key last, boolean lastInclusiv */ public Pair evictLeastAccessedEntries(long minSize) { checkArgument(minSize > 0); - - long removedSize = 0; - int removedEntries = 0; - - while (removedSize < minSize) { - Map.Entry entry = entries.firstEntry(); + RemovalCounters counters = RemovalCounters.create(); + while (counters.removedSize < minSize) { + Map.Entry> entry = entries.firstEntry(); if (entry == null) { break; } - Key key = entry.getKey(); - Value value = entry.getValue(); - try { - // add extra retain to avoid value being released while we are removing it - value.retain(); - } catch (IllegalReferenceCountException e) { - // Value was already released - continue; - } - // check that the value hasn't been recycled in between - if (value.matchesKey(key) && entries.remove(key, value)) { - ++removedEntries; - removedSize += weighter.getSize(value); - // release the cache reference - value.release(); - } - // remove the extra retain - value.release(); + removeEntry(entry, counters); } - - size.addAndGet(-removedSize); - return Pair.of(removedEntries, removedSize); + return handleRemovalResult(counters); } /** @@ -241,46 +323,18 @@ public Pair evictLeastAccessedEntries(long minSize) { * @return the tota */ public Pair evictLEntriesBeforeTimestamp(long maxTimestamp) { - long removedSize = 0; - int removedCount = 0; - + RemovalCounters counters = RemovalCounters.create(); while (true) { - Map.Entry entry = entries.firstEntry(); + Map.Entry> entry = entries.firstEntry(); if (entry == null) { break; } - Key key = entry.getKey(); - Value value = entry.getValue(); - try { - // add extra retain to avoid value being released while we are removing it - value.retain(); - } catch (IllegalReferenceCountException e) { - // Value was already released - continue; - } - try { - // check that the value hasn't been recycled in between - if (!value.matchesKey(key)) { - continue; - } - if (timestampExtractor.getTimestamp(value) > maxTimestamp) { - break; - } - boolean removeHits = entries.remove(key, value); - if (!removeHits) { - break; - } - removedSize += weighter.getSize(value); - removedCount++; - value.release(); - } finally { - // remove the extra retain - value.release(); + if (removeEntry(entry, counters, value -> timestampExtractor.getTimestamp(value) <= maxTimestamp) + == RemoveEntryResult.BREAK_LOOP) { + break; } } - - size.addAndGet(-removedSize); - return Pair.of(removedCount, removedSize); + return handleRemovalResult(counters); } /** @@ -300,36 +354,15 @@ public long getSize() { * @return size of removed entries */ public synchronized Pair clear() { - long removedSize = 0; - int removedCount = 0; - + RemovalCounters counters = RemovalCounters.create(); while (true) { - Map.Entry entry = entries.firstEntry(); + Map.Entry> entry = entries.firstEntry(); if (entry == null) { break; } - Key key = entry.getKey(); - Value value = entry.getValue(); - try { - // add extra retain to avoid value being released while we are removing it - value.retain(); - } catch (IllegalReferenceCountException e) { - // Value was already released - continue; - } - // check that the value hasn't been recycled in between - if (value.matchesKey(key) && entries.remove(key, value)) { - removedSize += weighter.getSize(value); - removedCount++; - // release the cache reference - value.release(); - } - // remove the extra retain - value.release(); + removeEntry(entry, counters); } - - size.getAndAdd(-removedSize); - return Pair.of(removedCount, removedSize); + return handleRemovalResult(counters); } /** From 94101c77d7bdd55667526a82a0ce6329d58baa8b Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 31 May 2024 00:16:07 +0300 Subject: [PATCH 3/7] Add javadoc, add refCnt check --- .../apache/bookkeeper/mledger/util/RangeCache.java | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 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 95045ac9983ef..63a26d8ecd80f 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 @@ -189,6 +189,10 @@ public boolean exists(Key key) { return key != null ? entries.containsKey(key) : true; } + /** + * Get the value associated with the key and increment the reference count of it. + * The caller is responsible for releasing the reference. + */ public Value get(Key key) { return getValue(key, entries.get(key)); } @@ -204,10 +208,13 @@ private Value getValue(Key key, IdentityWrapper valueWrapper) { // Value was already deallocated return null; } - if (value.matchesKey(key)) { + // check that the value matches the key and that there's at least 2 references to it since + // the cache should be holding one reference and a new reference was just added in this method + if (value.refCnt() > 1 && value.matchesKey(key)) { return value; } else { // Value or IdentityWrapper was recycled and already contains another value + // release the reference added in this method value.release(); return null; } @@ -279,7 +286,9 @@ private RemoveEntryResult removeEntry(Map.Entry> ent return RemoveEntryResult.BREAK_LOOP; } // check that the value hasn't been recycled in between - if (value.matchesKey(key) && entries.remove(key, identityWrapper)) { + // there should be at least 2 references since this method adds one and the cache should have one + // it is valid that the value contains references even after the key has been removed from the cache + if (value.refCnt() > 1 && value.matchesKey(key) && entries.remove(key, identityWrapper)) { identityWrapper.recycle(); counters.entryRemoved(weighter.getSize(value)); // remove the cache reference From d28db55d6cb8578168781179edceee8932144465 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 31 May 2024 00:29:08 +0300 Subject: [PATCH 4/7] Handle the race condition where the IdentityWrapper has already been recycled --- .../bookkeeper/mledger/util/RangeCache.java | 54 ++++++++++++------- 1 file changed, 34 insertions(+), 20 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 63a26d8ecd80f..c53fd7efc9772 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 @@ -49,7 +49,7 @@ public interface ValueWithKeyValidation extends ReferenceCounted { } // Map from key to nodes inside the linked list - private final ConcurrentNavigableMap> entries; + private final ConcurrentNavigableMap> entries; private AtomicLong size; // Total size of values stored in cache private final Weighter weighter; // Weighter object used to extract the size from values private final TimestampExtractor timestampExtractor; // Extract the timestamp associated with a value @@ -59,7 +59,7 @@ public interface ValueWithKeyValidation extends ReferenceCounted { * the map by calling the {@link Map#remove(Object, Object)} method. Certain race conditions could result in the * wrong value being removed from the map. The instances of this class are recycled to avoid creating new objects. */ - private static class IdentityWrapper { + private static class IdentityWrapper { private final Handle recyclerHandle; private static final Recycler RECYCLER = new Recycler() { @Override @@ -67,19 +67,25 @@ protected IdentityWrapper newObject(Handle recyclerHandle) { return new IdentityWrapper(recyclerHandle); } }; - private T value; + private K key; + private V value; private IdentityWrapper(Handle recyclerHandle) { this.recyclerHandle = recyclerHandle; } - static IdentityWrapper create(T value) { - IdentityWrapper identityWrapper = RECYCLER.get(); + static IdentityWrapper create(K key, V value) { + IdentityWrapper identityWrapper = RECYCLER.get(); + identityWrapper.key = key; identityWrapper.value = value; return identityWrapper; } - T get() { + K getKey() { + return key; + } + + V getValue() { return value; } @@ -96,7 +102,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hashCode(value); + return Objects.hashCode(key); } } @@ -171,7 +177,7 @@ public boolean put(Key key, Value value) { if (!value.matchesKey(key)) { throw new IllegalArgumentException("Value '" + value + "' does not match key '" + key + "'"); } - IdentityWrapper newWrapper = IdentityWrapper.create(value); + IdentityWrapper newWrapper = IdentityWrapper.create(key, value); if (entries.putIfAbsent(key, newWrapper) == null) { size.addAndGet(weighter.getSize(value)); return true; @@ -197,11 +203,15 @@ public Value get(Key key) { return getValue(key, entries.get(key)); } - private Value getValue(Key key, IdentityWrapper valueWrapper) { + private Value getValue(Key key, IdentityWrapper valueWrapper) { if (valueWrapper == null) { return null; } else { - Value value = valueWrapper.get(); + if (valueWrapper.getKey() != key) { + // the wrapper has been recycled and contains another key + return null; + } + Value value = valueWrapper.getValue(); try { value.retain(); } catch (IllegalReferenceCountException e) { @@ -233,7 +243,7 @@ public Collection getRange(Key first, Key last) { List values = new ArrayList(); // Return the values of the entries found in cache - for (Map.Entry> entry : entries.subMap(first, true, last, true).entrySet()) { + for (Map.Entry> entry : entries.subMap(first, true, last, true).entrySet()) { Value value = getValue(entry.getKey(), entry.getValue()); if (value != null) { values.add(value); @@ -252,8 +262,8 @@ public Collection getRange(Key first, Key last) { */ public Pair removeRange(Key first, Key last, boolean lastInclusive) { RemovalCounters counters = RemovalCounters.create(); - Map> subMap = entries.subMap(first, true, last, lastInclusive); - for (Map.Entry> entry : subMap.entrySet()) { + Map> subMap = entries.subMap(first, true, last, lastInclusive); + for (Map.Entry> entry : subMap.entrySet()) { removeEntry(entry, counters); } return handleRemovalResult(counters); @@ -265,15 +275,19 @@ enum RemoveEntryResult { BREAK_LOOP; } - private RemoveEntryResult removeEntry(Map.Entry> entry, RemovalCounters counters) { + private RemoveEntryResult removeEntry(Map.Entry> entry, RemovalCounters counters) { return removeEntry(entry, counters, (x) -> true); } - private RemoveEntryResult removeEntry(Map.Entry> entry, RemovalCounters counters, + private RemoveEntryResult removeEntry(Map.Entry> entry, RemovalCounters counters, Predicate removeCondition) { Key key = entry.getKey(); - IdentityWrapper identityWrapper = entry.getValue(); - Value value = identityWrapper.get(); + IdentityWrapper identityWrapper = entry.getValue(); + if (identityWrapper.getKey() != key) { + // the wrapper has been recycled and contains another key + return RemoveEntryResult.CONTINUE_LOOP; + } + Value value = identityWrapper.getValue(); try { // add extra retain to avoid value being released while we are removing it value.retain(); @@ -317,7 +331,7 @@ public Pair evictLeastAccessedEntries(long minSize) { checkArgument(minSize > 0); RemovalCounters counters = RemovalCounters.create(); while (counters.removedSize < minSize) { - Map.Entry> entry = entries.firstEntry(); + Map.Entry> entry = entries.firstEntry(); if (entry == null) { break; } @@ -334,7 +348,7 @@ public Pair evictLeastAccessedEntries(long minSize) { public Pair evictLEntriesBeforeTimestamp(long maxTimestamp) { RemovalCounters counters = RemovalCounters.create(); while (true) { - Map.Entry> entry = entries.firstEntry(); + Map.Entry> entry = entries.firstEntry(); if (entry == null) { break; } @@ -365,7 +379,7 @@ public long getSize() { public synchronized Pair clear() { RemovalCounters counters = RemovalCounters.create(); while (true) { - Map.Entry> entry = entries.firstEntry(); + Map.Entry> entry = entries.firstEntry(); if (entry == null) { break; } From b0f3473eff33f566cd218a5753dbf1871870794e Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 31 May 2024 00:39:46 +0300 Subject: [PATCH 5/7] Remove synchronization from the clear method since it's not needed anymore --- .../java/org/apache/bookkeeper/mledger/util/RangeCache.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c53fd7efc9772..8de16abc38d79 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 @@ -376,7 +376,7 @@ public long getSize() { * * @return size of removed entries */ - public synchronized Pair clear() { + public Pair clear() { RemovalCounters counters = RemovalCounters.create(); while (true) { Map.Entry> entry = entries.firstEntry(); From 56c5823cfc6f43945cc0745ef44b09ce4928e89b Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 31 May 2024 00:50:58 +0300 Subject: [PATCH 6/7] Add javadoc --- .../java/org/apache/bookkeeper/mledger/util/RangeCache.java | 4 ++++ 1 file changed, 4 insertions(+) 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 8de16abc38d79..46d03bea1b5ad 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 @@ -37,6 +37,10 @@ /** * Special type of cache where get() and delete() operations can be done over a range of keys. + * The implementation avoids locks and synchronization and relies on ConcurrentSkipListMap for storing the entries. + * Since there is no locks, there is a need to have a way to ensure that a single entry in the cache is removed + * exactly once. Removing an entry multiple times would result in the entries of the cache getting released too + * while they could still be in use. * * @param * Cache key. Needs to be Comparable From 2a6c2b1b42fa1051848812307a4b215fd994939b Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 31 May 2024 00:59:42 +0300 Subject: [PATCH 7/7] Improve the current race condition test for put method --- .../bookkeeper/mledger/util/RangeCacheTest.java | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java b/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java index 489ed55dfc200..01b3c67bf1113 100644 --- a/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java +++ b/managed-ledger/src/test/java/org/apache/bookkeeper/mledger/util/RangeCacheTest.java @@ -29,6 +29,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import lombok.Cleanup; import org.apache.commons.lang3.tuple.Pair; import org.testng.annotations.Test; @@ -285,14 +286,17 @@ public void evictions() { } @Test - public void testInParallel() { + public void testPutWhileClearIsCalledConcurrently() { RangeCache cache = new RangeCache<>(value -> value.s.length(), x -> 0); - ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor(); - executor.scheduleWithFixedDelay(cache::clear, 10, 10, TimeUnit.MILLISECONDS); - for (int i = 0; i < 1000; i++) { + int numberOfThreads = 4; + @Cleanup("shutdownNow") + ScheduledExecutorService executor = Executors.newScheduledThreadPool(numberOfThreads); + for (int i = 0; i < numberOfThreads; i++) { + executor.scheduleWithFixedDelay(cache::clear, 0, 1, TimeUnit.MILLISECONDS); + } + for (int i = 0; i < 100000; i++) { cache.put(i, new RefString(String.valueOf(i))); } - executor.shutdown(); } @Test