Resolve test flakiness caused by connection leak (#355)

This commit is contained in:
Shicong Huang 2019-11-08 15:40:27 -05:00 committed by GitHub
parent fe8ed4784d
commit 5c42df06ff
2 changed files with 41 additions and 13 deletions

View file

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

View file

@ -14,6 +14,7 @@
package google.registry.model.transaction; package google.registry.model.transaction;
import static com.google.common.truth.Truth.assertThat;
import static org.joda.time.DateTimeZone.UTC; import static org.joda.time.DateTimeZone.UTC;
import static org.testcontainers.containers.PostgreSQLContainer.POSTGRESQL_PORT; import static org.testcontainers.containers.PostgreSQLContainer.POSTGRESQL_PORT;
@ -33,6 +34,7 @@ import java.nio.charset.StandardCharsets;
import java.nio.file.Files; import java.nio.file.Files;
import java.sql.Connection; import java.sql.Connection;
import java.sql.Driver; import java.sql.Driver;
import java.sql.ResultSet;
import java.sql.SQLException; import java.sql.SQLException;
import java.sql.Statement; import java.sql.Statement;
import java.util.ArrayList; import java.util.ArrayList;
@ -114,7 +116,7 @@ public class JpaTransactionManagerRule extends ExternalResource {
builder.putAll(userProperties); builder.putAll(userProperties);
properties = builder.build(); properties = builder.build();
} }
assertNormalActiveConnection();
emf = emf =
createEntityManagerFactory( createEntityManagerFactory(
getJdbcUrlFor(POSTGRES_DB_NAME), getJdbcUrlFor(POSTGRES_DB_NAME),
@ -132,8 +134,31 @@ public class JpaTransactionManagerRule extends ExternalResource {
TransactionManagerFactory.jpaTm = cachedTm; TransactionManagerFactory.jpaTm = cachedTm;
if (emf != null) { if (emf != null) {
emf.close(); emf.close();
emf = null;
} }
cachedTm = 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) { private static String readSqlInClassPath(String sqlScriptPath) {