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 index 0c17008e2010..569b441264c7 100644 --- 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 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