diff --git a/core/trino-main/src/main/java/io/trino/memory/MemoryPool.java b/core/trino-main/src/main/java/io/trino/memory/MemoryPool.java index b330931af945..191179a8ae16 100644 --- a/core/trino-main/src/main/java/io/trino/memory/MemoryPool.java +++ b/core/trino-main/src/main/java/io/trino/memory/MemoryPool.java @@ -31,6 +31,7 @@ import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import static com.google.common.base.MoreObjects.toStringHelper; @@ -53,9 +54,9 @@ public class MemoryPool @GuardedBy("this") private NonCancellableMemoryFuture future; - @GuardedBy("this") // TODO: It would be better if we just tracked QueryContexts, but their lifecycle is managed by a weak reference, so we can't do that - private final Map queryMemoryReservations = new HashMap<>(); + // It is guarded for updates by this, but can be read without holding a lock + private final Map queryMemoryReservations = new ConcurrentHashMap<>(); // This map keeps track of all the tagged allocations, e.g., query-1 -> ['TableScanOperator': 10MB, 'LazyOutputBuffer': 5MB, ...] @GuardedBy("this") @@ -347,7 +348,7 @@ public synchronized long getReservedRevocableBytes() return reservedRevocableBytes; } - synchronized long getQueryMemoryReservation(QueryId queryId) + long getQueryMemoryReservation(QueryId queryId) { return queryMemoryReservations.getOrDefault(queryId, 0L); } diff --git a/core/trino-main/src/main/java/io/trino/memory/QueryContext.java b/core/trino-main/src/main/java/io/trino/memory/QueryContext.java index c30e3ad282f7..0c7711e8d139 100644 --- a/core/trino-main/src/main/java/io/trino/memory/QueryContext.java +++ b/core/trino-main/src/main/java/io/trino/memory/QueryContext.java @@ -225,7 +225,7 @@ public synchronized MemoryPool getMemoryPool() return memoryPool; } - public synchronized long getUserMemoryReservation() + public long getUserMemoryReservation() { return memoryPool.getQueryMemoryReservation(queryId); }