Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix evictable cache invalidation race condition #22302

Merged
merged 1 commit into from
Jun 13, 2024
Merged

Fix evictable cache invalidation race condition #22302

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 51 additions & 32 deletions lib/trino-cache/src/main/java/io/trino/cache/EvictableCache.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,9 @@
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiFunction;

import static com.google.common.base.MoreObjects.firstNonNull;
import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Verify.verify;
import static java.lang.String.format;
Expand Down Expand Up @@ -68,7 +68,8 @@ class EvictableCache<K, V>
// The dataCache must be bounded.
private final LoadingCache<Token<K>, V> dataCache;

private final AtomicInteger invalidations = new AtomicInteger();
// Logically a concurrent Multiset
private final ConcurrentHashMap<Token<K>, Long> ongoingLoads = new ConcurrentHashMap<>();

EvictableCache(CacheBuilder<? super Token<K>, ? super V> cacheBuilder, CacheLoader<? super K, V> cacheLoader)
{
Expand All @@ -77,9 +78,13 @@ class EvictableCache<K, V>
.<Token<K>, V>removalListener(removal -> {
Token<K> token = removal.getKey();
verify(token != null, "token is null");
if (removal.getCause() != RemovalCause.REPLACED) {
tokens.remove(token.getKey(), token);
if (removal.getCause() == RemovalCause.REPLACED) {
return;
}
if (removal.getCause() == RemovalCause.EXPIRED && ongoingLoads.containsKey(token)) {
return;
}
tokens.remove(token.getKey(), token);
}),
new TokenCacheLoader<>(cacheLoader));
}
Expand All @@ -106,22 +111,15 @@ public V get(K key, Callable<? extends V> valueLoader)
throws ExecutionException
{
Token<K> newToken = new Token<>(key);
int invalidations = this.invalidations.get();
Token<K> token = tokens.computeIfAbsent(key, _ -> newToken);
try {
V value = dataCache.get(token, valueLoader);
if (invalidations == this.invalidations.get()) {
// Revive token if it got expired before reloading
if (tokens.putIfAbsent(key, token) == null) {
// Revived
if (!dataCache.asMap().containsKey(token)) {
// We revived, but the token does not correspond to a live entry anymore.
// It would stay in tokens forever, so let's remove it.
tokens.remove(key, token);
}
}
startLoading(token);
try {
return dataCache.get(token, valueLoader);
}
finally {
endLoading(token);
}
return value;
}
catch (Throwable e) {
if (newToken == token) {
Expand All @@ -132,29 +130,25 @@ public V get(K key, Callable<? extends V> valueLoader)
}
throw e;
}
finally {
removeDangling(token);
}
}

@Override
public V get(K key)
throws ExecutionException
{
Token<K> newToken = new Token<>(key);
int invalidations = this.invalidations.get();
Token<K> token = tokens.computeIfAbsent(key, _ -> newToken);
try {
V value = dataCache.get(token);
if (invalidations == this.invalidations.get()) {
// Revive token if it got expired before reloading
if (tokens.putIfAbsent(key, token) == null) {
// Revived
if (!dataCache.asMap().containsKey(token)) {
// We revived, but the token does not correspond to a live entry anymore.
// It would stay in tokens forever, so let's remove it.
tokens.remove(key, token);
}
}
startLoading(token);
try {
return dataCache.get(token);
}
finally {
endLoading(token);
}
return value;
}
catch (Throwable e) {
if (newToken == token) {
Expand All @@ -165,6 +159,9 @@ public V get(K key)
}
throw e;
}
finally {
removeDangling(token);
}
}

@Override
Expand Down Expand Up @@ -218,6 +215,30 @@ public ImmutableMap<K, V> getAll(Iterable<? extends K> keys)
}
}

private void startLoading(Token<K> token)
{
ongoingLoads.compute(token, (_, count) -> firstNonNull(count, 0L) + 1);
}

private void endLoading(Token<K> token)
{
ongoingLoads.compute(token, (_, count) -> {
verify(count != null && count > 0, "Incorrect count for token %s: %s", token, count);
if (count == 1) {
return null;
}
return count - 1;
});
}

// Token eviction via removalListener is blocked during loading, so we may need to do manual cleanup
private void removeDangling(Token<K> token)
{
if (!dataCache.asMap().containsKey(token)) {
tokens.remove(token.getKey(), token);
}
}

@Override
public void refresh(K key)
{
Expand Down Expand Up @@ -248,7 +269,6 @@ int tokensCount()
@Override
public void invalidate(Object key)
{
invalidations.incrementAndGet();
@SuppressWarnings("SuspiciousMethodCalls") // Object passed to map as key K
Token<K> token = tokens.remove(key);
if (token != null) {
Expand All @@ -259,7 +279,6 @@ public void invalidate(Object key)
@Override
public void invalidateAll()
{
invalidations.incrementAndGet();
dataCache.invalidateAll();
tokens.clear();
}
Expand Down