Add more memory assertions in TestHashAggregationOperator

trinodb · Jan 13, 2022 · 8ec9c08 · 8ec9c08
1 parent ab9962c
commit 8ec9c08
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/core/trino-main/src/test/java/io/trino/operator/TestHashAggregationOperator.java b/core/trino-main/src/test/java/io/trino/operator/TestHashAggregationOperator.java
@@ -300,6 +300,9 @@ public void testHashAggregationMemoryReservation(boolean hashEnabled, boolean sp
         Operator operator = operatorFactory.createOperator(driverContext);
         toPages(operator, input.iterator(), revokeMemoryWhenAddingPages);
         assertEquals(getOnlyElement(operator.getOperatorContext().getNestedOperatorStats()).getUserMemoryReservation().toBytes(), 0);
+        assertEquals(getOnlyElement(operator.getOperatorContext().getNestedOperatorStats()).getRevocableMemoryReservation().toBytes(), 0);
+        // TODO (https://github.com/trinodb/trino/issues/10596): it should be 0, since operator is finished
+        assertEquals(getOnlyElement(operator.getOperatorContext().getNestedOperatorStats()).getSystemMemoryReservation().toBytes(), spillEnabled && revokeMemoryWhenAddingPages ? 5_350_968 : 0);
     }
 
     @Test(dataProvider = "hashEnabled", expectedExceptions = ExceededMemoryLimitException.class, expectedExceptionsMessageRegExp = "Query exceeded per-node user memory limit of 10B.*")
@@ -574,6 +577,7 @@ public void testMultiplePartialFlushes(boolean hashEnabled)
 
         assertEquals(driverContext.getSystemMemoryUsage(), 0);
         assertEquals(driverContext.getMemoryUsage(), 0);
+        assertEquals(driverContext.getRevocableMemoryUsage(), 0);
     }
 
     @Test
@@ -721,6 +725,7 @@ private void testMemoryTracking(boolean useSystemMemory)
 
         assertEquals(driverContext.getSystemMemoryUsage(), 0);
         assertEquals(driverContext.getMemoryUsage(), 0);
+        assertEquals(driverContext.getRevocableMemoryUsage(), 0);
     }
 
     private DriverContext createDriverContext()