Resolve test flakiness caused by connection leak (#355)

core/src/main/java/google/registry/persistence/HibernateSchemaExporter.java

@@ -25,6 +25,7 @@ import java.util.Properties;
 import java.util.stream.Stream;
 import javax.persistence.AttributeConverter;
 import org.hibernate.boot.MetadataSources;
+import org.hibernate.boot.registry.StandardServiceRegistry;
 import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
 import org.hibernate.cfg.Environment;
 import org.hibernate.jpa.boot.internal.ParsedPersistenceXmlDescriptor;
@@ -62,8 +63,9 @@ public class HibernateSchemaExporter {
     settings.put(
         Environment.PHYSICAL_NAMING_STRATEGY, NomulusNamingStrategy.class.getCanonicalName());
 
-    MetadataSources metadata =
-        new MetadataSources(new StandardServiceRegistryBuilder().applySettings(settings).build());
+    try (StandardServiceRegistry registry =
+        new StandardServiceRegistryBuilder().applySettings(settings).build()) {
+      MetadataSources metadata = new MetadataSources(registry);
 
       // Note that we need to also add all converters to the Hibernate context because
       // the entity class may use the customized type.
@@ -77,6 +79,7 @@ public class HibernateSchemaExporter {
       export.setOutputFile(outputFile.getAbsolutePath());
       export.createOnly(EnumSet.of(TargetType.SCRIPT), metadata.buildMetadata());
     }
+  }
 
   private ImmutableList<Class> findAllConverters() {
     ParsedPersistenceXmlDescriptor descriptor =

core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java

@@ -14,6 +14,7 @@
 
 package google.registry.model.transaction;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.joda.time.DateTimeZone.UTC;
 import static org.testcontainers.containers.PostgreSQLContainer.POSTGRESQL_PORT;
 
@@ -33,6 +34,7 @@ import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.sql.Connection;
 import java.sql.Driver;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
@@ -114,7 +116,7 @@ public class JpaTransactionManagerRule extends ExternalResource {
       builder.putAll(userProperties);
       properties = builder.build();
     }
-
+    assertNormalActiveConnection();
     emf =
         createEntityManagerFactory(
             getJdbcUrlFor(POSTGRES_DB_NAME),
@@ -132,8 +134,31 @@ public class JpaTransactionManagerRule extends ExternalResource {
     TransactionManagerFactory.jpaTm = cachedTm;
     if (emf != null) {
       emf.close();
+      emf = null;
     }
     cachedTm = null;
+    assertNormalActiveConnection();
+  }
+
+  /**
+   * This function throws exception if it detects connection leak by checking the metadata table
+   * pg_stat_activity.
+   */
+  private void assertNormalActiveConnection() {
+    try (Connection conn = createConnection(POSTGRES_DB_NAME);
+        Statement statement = conn.createStatement()) {
+      ResultSet rs =
+          statement.executeQuery(
+              "SELECT COUNT(1) FROM pg_stat_activity WHERE usename = '"
+                  + database.getUsername()
+                  + "'");
+      rs.next();
+      long activeConns = rs.getLong(1);
+      // There should be only 1 active connection which is executing this query
+      assertThat(activeConns).isEqualTo(1L);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
   }
 
   private static String readSqlInClassPath(String sqlScriptPath) {