From 5c42df06ffe2b39a9475cae8196c5fc8d3e8174c Mon Sep 17 00:00:00 2001
From: Shicong Huang <shicong@google.com>
Date: Fri, 8 Nov 2019 15:40:27 -0500
Subject: [PATCH] Resolve test flakiness caused by connection leak (#355)

---
 .../persistence/HibernateSchemaExporter.java  | 27 ++++++++++---------
 .../JpaTransactionManagerRule.java            | 27 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/core/src/main/java/google/registry/persistence/HibernateSchemaExporter.java b/core/src/main/java/google/registry/persistence/HibernateSchemaExporter.java
index 23a1e62a2..50ad9485e 100644
--- a/core/src/main/java/google/registry/persistence/HibernateSchemaExporter.java
+++ b/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,20 +63,22 @@ 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.
-    Stream.concat(entityClasses.stream(), findAllConverters().stream())
-        .forEach(metadata::addAnnotatedClass);
+      // Note that we need to also add all converters to the Hibernate context because
+      // the entity class may use the customized type.
+      Stream.concat(entityClasses.stream(), findAllConverters().stream())
+          .forEach(metadata::addAnnotatedClass);
 
-    SchemaExport export = new SchemaExport();
-    export.setHaltOnError(true);
-    export.setFormat(true);
-    export.setDelimiter(";");
-    export.setOutputFile(outputFile.getAbsolutePath());
-    export.createOnly(EnumSet.of(TargetType.SCRIPT), metadata.buildMetadata());
+      SchemaExport export = new SchemaExport();
+      export.setHaltOnError(true);
+      export.setFormat(true);
+      export.setDelimiter(";");
+      export.setOutputFile(outputFile.getAbsolutePath());
+      export.createOnly(EnumSet.of(TargetType.SCRIPT), metadata.buildMetadata());
+    }
   }
 
   private ImmutableList<Class> findAllConverters() {
diff --git a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java b/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java
index 06ec0cab8..3562adb41 100644
--- a/core/src/test/java/google/registry/model/transaction/JpaTransactionManagerRule.java
+++ b/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) {