Use coarse grained memory reporting to reduce congestion

core/trino-main/src/main/java/io/trino/operator/join/unspilled/HashBuilderOperator.java

@@ -17,7 +17,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.errorprone.annotations.ThreadSafe;
-import io.trino.memory.context.ThresholdLocalMemoryContext;
+import io.trino.memory.context.CoarseGrainLocalMemoryContext;
 import io.trino.operator.DriverContext;
 import io.trino.operator.HashArraySizeSupplier;
 import io.trino.operator.Operator;
@@ -38,7 +38,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
-import static io.trino.memory.context.ThresholdLocalMemoryContext.DEFAULT_SYNC_THRESHOLD;
+import static io.trino.memory.context.CoarseGrainLocalMemoryContext.DEFAULT_GRANULARITY;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -160,7 +160,7 @@ public enum State
     }
 
     private final OperatorContext operatorContext;
-    private final ThresholdLocalMemoryContext localUserMemoryContext;
+    private final CoarseGrainLocalMemoryContext localUserMemoryContext;
     private final PartitionedLookupSourceFactory lookupSourceFactory;
     private final ListenableFuture<Void> lookupSourceFactoryDestroyed;
     private final int partitionIndex;
@@ -194,7 +194,7 @@ public HashBuilderOperator(
             PagesIndex.Factory pagesIndexFactory,
             HashArraySizeSupplier hashArraySizeSupplier)
     {
-        this(operatorContext, lookupSourceFactory, partitionIndex, outputChannels, hashChannels, preComputedHashChannel, filterFunctionFactory, sortChannel, searchFunctionFactories, expectedPositions, pagesIndexFactory, hashArraySizeSupplier, DEFAULT_SYNC_THRESHOLD);
+        this(operatorContext, lookupSourceFactory, partitionIndex, outputChannels, hashChannels, preComputedHashChannel, filterFunctionFactory, sortChannel, searchFunctionFactories, expectedPositions, pagesIndexFactory, hashArraySizeSupplier, DEFAULT_GRANULARITY);
     }
 
     @VisibleForTesting
@@ -220,7 +220,7 @@ public HashBuilderOperator(
         this.filterFunctionFactory = filterFunctionFactory;
         this.sortChannel = sortChannel;
         this.searchFunctionFactories = searchFunctionFactories;
-        this.localUserMemoryContext = new ThresholdLocalMemoryContext(operatorContext.localUserMemoryContext(), memorySyncThreshold);
+        this.localUserMemoryContext = new CoarseGrainLocalMemoryContext(operatorContext.localUserMemoryContext(), memorySyncThreshold);
 
         this.index = pagesIndexFactory.newPagesIndex(lookupSourceFactory.getTypes(), expectedPositions);
         this.lookupSourceFactory = lookupSourceFactory;
@@ -334,7 +334,6 @@ private void finishInput()
         }
         LookupSourceSupplier partition = buildLookupSource();
         localUserMemoryContext.setBytes(partition.get().getInMemorySizeInBytes());
-        localUserMemoryContext.sync();
         lookupSourceNotNeeded = Optional.of(lookupSourceFactory.lendPartitionLookupSource(partitionIndex, partition));
 
         index = null;

lib/trino-memory-context/src/main/java/io/trino/memory/context/CoarseGrainLocalMemoryContext.java

@@ -0,0 +1,97 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.trino.memory.context;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.requireNonNull;
+
+/**
+ *  This class prevents contention in the memory tracking system by coarsening the granularity of memory tracking.
+ *  In this paradigm, the most small incremental increases at the byte granularity will not actually result
+ *  in a different coarse granularity value, so no reporting into memory tracking is necessary
+ */
+public class CoarseGrainLocalMemoryContext
+        implements LocalMemoryContext
+{
+    public static final long DEFAULT_GRANULARITY = 65536;
+
+    private final LocalMemoryContext delegate;
+    private final long granularity;
+    private final long mask;
+    private long currentBytes;
+
+    public CoarseGrainLocalMemoryContext(LocalMemoryContext delegate)
+    {
+        this(delegate, DEFAULT_GRANULARITY);
+    }
+
+    public CoarseGrainLocalMemoryContext(LocalMemoryContext delegate, long granularity)
+    {
+        this.delegate = requireNonNull(delegate, "delegate is null");
+        checkArgument(granularity > 0, "granularity must be greater than 0");
+        checkArgument((granularity & (granularity - 1)) == 0, "granularity must be a power of 2");
+        this.granularity = granularity;
+        this.mask = ~(granularity - 1);
+    }
+
+    @Override
+    public synchronized long getBytes()
+    {
+        return currentBytes;
+    }
+
+    @Override
+    public synchronized ListenableFuture<Void> setBytes(long bytes)
+    {
+        long roundedUpBytes = roundUpToNearest(bytes);
+        if (roundedUpBytes != currentBytes) {
+            currentBytes = roundedUpBytes;
+            return delegate.setBytes(currentBytes);
+        }
+        return Futures.immediateVoidFuture();
+    }
+
+    @Override
+    public synchronized boolean trySetBytes(long bytes)
+    {
+        long roundedUpBytes = roundUpToNearest(bytes);
+        if (roundedUpBytes != currentBytes) {
+            if (delegate.trySetBytes(roundedUpBytes)) {
+                currentBytes = roundedUpBytes;
+                return true;
+            }
+            return false;
+        }
+        return true;
+    }
+
+    @Override
+    public synchronized void close()
+    {
+        delegate.close();
+        currentBytes = 0;
+    }
+
+    @VisibleForTesting
+    long roundUpToNearest(long bytes)
+    {
+        long masked = bytes & mask;
+        return masked == bytes ? masked : masked + granularity;
+    }
+}

lib/trino-memory-context/src/test/java/io/trino/memory/context/TestMemoryContexts.java

@@ -158,6 +158,66 @@ public void testClosedLocalMemoryContext()
                 .hasMessage("SimpleLocalMemoryContext is already closed");
     }
 
+    @Test
+    public void testCoarseGrainLocalMemoryContext()
+    {
+        final long guaranteedMemory = (long) Math.pow(2, 30);
+        TestMemoryReservationHandler reservationHandler = new TestMemoryReservationHandler(2 * guaranteedMemory);
+        AggregatedMemoryContext aggregateContext = newRootAggregatedMemoryContext(reservationHandler, guaranteedMemory);
+        LocalMemoryContext delegate = aggregateContext.newLocalMemoryContext("test");
+
+        final long granularity = (long) Math.pow(2, 10);
+        CoarseGrainLocalMemoryContext coarseGrainContext = new CoarseGrainLocalMemoryContext(delegate, granularity);
+
+        assertCoarseGrainContextValues(coarseGrainContext, delegate, 1, granularity);
+
+        // boundaries
+        assertCoarseGrainContextValues(coarseGrainContext, delegate, 0, 0);
+        assertCoarseGrainContextValues(coarseGrainContext, delegate, granularity, granularity);
+
+        assertCoarseGrainContextValues(coarseGrainContext, delegate, granularity + 1, 2 * granularity);
+        assertCoarseGrainContextValues(coarseGrainContext, delegate, 2 * granularity + 1, 3 * granularity);
+        assertCoarseGrainContextValues(coarseGrainContext, delegate, 2 * granularity, 2 * granularity);
+
+        // gets set to the next coarse unit
+        // k*granularity + x in [1, granularity] leads to setting the context to (k+1) * granularity
+        assertCoarseGrainContextValues(coarseGrainContext, delegate, 0, 0);
+        for (int i = 1; i <= granularity; i++) {
+            assertCoarseGrainContextValues(coarseGrainContext, delegate, i, granularity);
+        }
+
+        // threshold not a power of 2
+        assertThatThrownBy(() -> new CoarseGrainLocalMemoryContext(delegate, 100))
+                .isInstanceOf(IllegalArgumentException.class)
+                .hasMessage("granularity must be a power of 2");
+
+        // test trySetBytes()
+        coarseGrainContext.setBytes(0);
+        assertThat(coarseGrainContext.trySetBytes(1)).isEqualTo(true);
+        assertThat(coarseGrainContext.getBytes()).isEqualTo(granularity);
+        assertThat(delegate.getBytes()).isEqualTo(granularity);
+
+        // new bytes = previously set bytes(rounded-up)
+        assertThat(coarseGrainContext.trySetBytes(2)).isEqualTo(true);
+        assertThat(coarseGrainContext.getBytes()).isEqualTo(granularity);
+        assertThat(delegate.getBytes()).isEqualTo(granularity);
+
+        // something underlying delegate cannot set
+        assertThat(coarseGrainContext.trySetBytes(guaranteedMemory * 3)).isEqualTo(false);
+        assertThat(coarseGrainContext.getBytes()).isEqualTo(granularity);
+        assertThat(delegate.getBytes()).isEqualTo(granularity);
+    }
+
+    private static void assertCoarseGrainContextValues(CoarseGrainLocalMemoryContext coarseGrainContext,
+                                                       LocalMemoryContext delegate, long valueToSet, long expectedValue)
+    {
+        assertThat(coarseGrainContext.setBytes(valueToSet)).isEqualTo(NOT_BLOCKED);
+
+        assertThat(coarseGrainContext.roundUpToNearest(valueToSet)).isEqualTo(expectedValue);
+        assertThat(coarseGrainContext.getBytes()).isEqualTo(expectedValue);
+        assertThat(delegate.getBytes()).isEqualTo(expectedValue);
+    }
+
     @Test
     public void testClosedAggregateMemoryContext()
     {