Check all queries returned memory on query runner teardown

trinodb · Feb 23, 2022 · 173402e · 173402e
1 parent bb60e2d
commit 173402e
Showing 1 changed file with 41 additions and 5 deletions.
diff --git a/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestQueryFramework.java b/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestQueryFramework.java
@@ -27,6 +27,7 @@
 import io.trino.metadata.TableMetadata;
 import io.trino.operator.OperatorStats;
 import io.trino.server.DynamicFilterService.DynamicFiltersStats;
+import io.trino.server.testing.TestingTrinoServer;
 import io.trino.spi.QueryId;
 import io.trino.spi.type.Type;
 import io.trino.sql.analyzer.QueryExplainer;
@@ -62,8 +63,11 @@
 import static io.trino.SystemSessionProperties.JOIN_REORDERING_STRATEGY;
 import static io.trino.sql.ParsingUtil.createParsingOptions;
 import static io.trino.sql.SqlFormatter.formatSql;
+import static io.trino.testing.assertions.Assert.assertEventually;
 import static io.trino.transaction.TransactionBuilder.transaction;
+import static java.lang.String.format;
 import static java.util.Collections.emptyList;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.testng.Assert.assertEquals;
@@ -93,11 +97,43 @@ protected abstract QueryRunner createQueryRunner()
     public final void close()
             throws Exception
     {
-        afterClassCloser.close();
-        queryRunner = null;
-        h2QueryRunner = null;
-        sqlParser = null;
-        queryAssertions = null;
+        try (afterClassCloser) {
+            checkQueryMemoryReleased();
+        }
+        finally {
+            queryRunner = null;
+            h2QueryRunner = null;
+            sqlParser = null;
+            queryAssertions = null;
+        }
+    }
+
+    private void checkQueryMemoryReleased()
+    {
+        if (queryRunner == null) {
+            return;
+        }
+        if (!(queryRunner instanceof DistributedQueryRunner)) {
+            return;
+        }
+        DistributedQueryRunner distributedQueryRunner = (DistributedQueryRunner) queryRunner;
+        assertEventually(
+                new Duration(10, SECONDS),
+                new Duration(1, SECONDS),
+                () -> {
+                    List<TestingTrinoServer> servers = distributedQueryRunner.getServers();
+                    for (int serverId = 0; serverId < servers.size(); ++serverId) {
+                        TestingTrinoServer server = servers.get(serverId);
+                        String serverName = format("server_%d(%s)", serverId, server.isCoordinator() ? "coordinator" : "worker");
+                        assertThat(server.getLocalMemoryManager().getMemoryPool().getReservedBytes())
+                                .describedAs("memory reservation on " + serverName)
+                                .isZero();
+                    }
+
+                    assertThat(distributedQueryRunner.getCoordinator().getClusterMemoryManager().getClusterTotalMemoryReservation())
+                            .describedAs("cluster memory reservation")
+                            .isZero();
+                });
     }
 
     @Test