From 4e0fccda7f4ac63da0cc0d99cf0be3625234a6db Mon Sep 17 00:00:00 2001 From: gbrodman Date: Fri, 25 Sep 2026 17:46:41 +0000 Subject: [PATCH] Speed up test database resets and Postgres config (#3238) - Mount the test Postgres data directory on an in-memory tmpfs and disable fsync, synchronous_commit, and full_page_writes for ephemeral test containers. - Override the Nomulus SQL dialect to emit 'deferrable initially deferred' on foreign keys exported by HibernateSchemaExporter. This matches prod. - Replace per-test TRUNCATES (etc) with DELETEs -- the vast majority of tests only change a couple tables and DELETEing a few rows is much faster than having to reallocate all the disk space for each and every table every time - Cache the connection to Postgres that we use for resetting data, instead of having to re-create it each time - Add a test for child entity foreign key generation This drastically speeds up tests, cutting out like 30-40% of total test time (not including compilation) (we also update a couple queries to be deterministically ordered to avoid test flakes) --- .../registry/model/registrar/Registrar.java | 6 +- .../java/google/registry/model/tld/Tlds.java | 3 +- .../persistence/NomulusPostgreSQLDialect.java | 12 +++ .../HibernateSchemaExporterTest.java | 29 +++++- .../JpaTransactionManagerExtension.java | 98 +++++++++++++------ .../sql/schema/db-schema.sql.generated | 20 ++-- .../registry/sql/flyway/SchemaTest.java | 5 +- 7 files changed, 127 insertions(+), 46 deletions(-) diff --git a/core/src/main/java/google/registry/model/registrar/Registrar.java b/core/src/main/java/google/registry/model/registrar/Registrar.java index 3ecaacdf6..738d9dcc5 100644 --- a/core/src/main/java/google/registry/model/registrar/Registrar.java +++ b/core/src/main/java/google/registry/model/registrar/Registrar.java @@ -211,7 +211,7 @@ public class Registrar extends UpdateAutoTimestampEntity implements Buildable, J memoizeWithShortExpiration( () -> Maps.uniqueIndex( - tm().reTransact(() -> tm().loadAllOf(Registrar.class)), + tm().reTransact(() -> tm().loadAllOfSorted(Registrar.class, "registrarId")), Registrar::getRegistrarId)); /** @@ -1050,9 +1050,9 @@ public class Registrar extends UpdateAutoTimestampEntity implements Buildable, J return email; } - /** Loads all registrar entities directly from the database. */ + /** Loads all registrar entities directly from the database, sorted by {@code registrarId}. */ public static Iterable loadAll() { - return tm().transact(() -> tm().loadAllOf(Registrar.class)); + return loadAllSorted("registrarId"); } /** Loads all registrar entities directly from the database, sorted by the given field names. */ diff --git a/core/src/main/java/google/registry/model/tld/Tlds.java b/core/src/main/java/google/registry/model/tld/Tlds.java index d80fab7a0..33917608d 100644 --- a/core/src/main/java/google/registry/model/tld/Tlds.java +++ b/core/src/main/java/google/registry/model/tld/Tlds.java @@ -63,7 +63,8 @@ public final class Tlds { EntityManager entityManager = tm().getEntityManager(); Stream resultStream = entityManager - .createQuery("SELECT tldStr, tldType FROM Tld", Object[].class) + .createQuery( + "SELECT tldStr, tldType FROM Tld ORDER BY tldStr", Object[].class) .getResultStream(); return resultStream .map(e -> Maps.immutableEntry((String) e[0], ((TldType) e[1]))) diff --git a/core/src/main/java/google/registry/persistence/NomulusPostgreSQLDialect.java b/core/src/main/java/google/registry/persistence/NomulusPostgreSQLDialect.java index ba2717e26..3d155d979 100644 --- a/core/src/main/java/google/registry/persistence/NomulusPostgreSQLDialect.java +++ b/core/src/main/java/google/registry/persistence/NomulusPostgreSQLDialect.java @@ -41,6 +41,18 @@ public class NomulusPostgreSQLDialect extends PostgreSQLDialect { return elementTypeName + "[]"; } + @Override + public String getAddForeignKeyConstraintString( + String constraintName, + String[] foreignKey, + String referencedTable, + String[] primaryKey, + boolean referencesPrimaryKey) { + return super.getAddForeignKeyConstraintString( + constraintName, foreignKey, referencedTable, primaryKey, referencesPrimaryKey) + + " deferrable initially deferred"; + } + @Override public void contributeTypes( TypeContributions typeContributions, ServiceRegistry serviceRegistry) { diff --git a/core/src/test/java/google/registry/persistence/HibernateSchemaExporterTest.java b/core/src/test/java/google/registry/persistence/HibernateSchemaExporterTest.java index 9ba25f3ca..4e5d6ffc3 100644 --- a/core/src/test/java/google/registry/persistence/HibernateSchemaExporterTest.java +++ b/core/src/test/java/google/registry/persistence/HibernateSchemaExporterTest.java @@ -18,7 +18,10 @@ import static com.google.common.truth.Truth.assertThat; import com.google.common.collect.ImmutableList; import jakarta.persistence.Entity; +import jakarta.persistence.ForeignKey; import jakarta.persistence.Id; +import jakarta.persistence.JoinColumn; +import jakarta.persistence.ManyToOne; import java.io.File; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -56,16 +59,29 @@ class HibernateSchemaExporterTest { @Test void export_succeeds() throws IOException { File sqlFile = Files.createFile(tmpDir.resolve("tempfile.dat")).toFile(); - exporter.export(ImmutableList.of(HibernateSchemaTestEntity.class), sqlFile); + exporter.export( + ImmutableList.of(HibernateSchemaTestEntity.class, HibernateSchemaChildEntity.class), + sqlFile); assertThat(Files.readAllBytes(sqlFile.toPath())) .isEqualTo( - """ +""" + + create table "ChildEntity" ( + id text not null, + parent_name text, + primary key (id) + ); create table "TestEntity" ( name text not null, cu text, primary key (name) ); + + alter table if exists "ChildEntity"\s + add constraint FK_child_parent\s + foreign key (parent_name)\s + references "TestEntity" deferrable initially deferred; """ .getBytes(StandardCharsets.UTF_8)); } @@ -76,4 +92,13 @@ class HibernateSchemaExporterTest { CurrencyUnit cu; } + + @Entity(name = "ChildEntity") + private static class HibernateSchemaChildEntity { + @Id String id; + + @ManyToOne + @JoinColumn(name = "parent_name", foreignKey = @ForeignKey(name = "FK_child_parent")) + HibernateSchemaTestEntity parent; + } } diff --git a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java index 18b5b9216..a71ae439b 100644 --- a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java +++ b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java @@ -21,7 +21,6 @@ import static google.registry.testing.DatabaseHelper.persistResources; import static java.nio.charset.StandardCharsets.UTF_8; import static org.testcontainers.containers.PostgreSQLContainer.POSTGRESQL_PORT; -import com.google.common.base.Joiner; import com.google.common.base.Suppliers; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -104,17 +103,20 @@ public abstract class JpaTransactionManagerExtension // reused between test methods if the requested schema remains the same. private static EntityManagerFactory emf; // Hash of the ORM entity names in the current schema in the test db. - private static int emfEntityHash; - - private JpaTransactionManager cachedTm; // Hash of the ORM entity names requested by this extension instance. private final int entityHash; + private JpaTransactionManager cachedTm; // Whether to create nomulus tables in the test db. Right now, only the JpaTestExtensions set this // to false. private boolean includeNomulusSchema = true; + // Cached re-usable connection used for resetting tables / schema + private static Connection reusableConnection; + // Cached "reset" statement so we don't need to query SQL each time to find out what tables exist + private static String cachedResetSql; + // Whether to pre-populate some registrars for ease of testing. private final boolean withCannedData; @@ -144,7 +146,19 @@ public abstract class JpaTransactionManagerExtension "POSTGRES_INITDB_ARGS", "--encoding=UTF8 --lc-collate=en_US.UTF8 --lc-ctype=en_US.UTF8" + " --locale-provider=libc --no-locale") - .withDatabaseName(POSTGRES_DB_NAME); + .withDatabaseName(POSTGRES_DB_NAME) + // Use tmpFs so that we keep everything in memory. No point in writing to disk + .withTmpFs(ImmutableMap.of("/var/lib/postgresql/data", "rw")) + // Testcontainers defaults to fsync=off but we want the additional options turned off as + // well, for speed purposes + .withCommand( + "postgres", + "-c", + "fsync=off", + "-c", + "synchronous_commit=off", + "-c", + "full_page_writes=off"); container.start(); return container; } @@ -159,11 +173,31 @@ public abstract class JpaTransactionManagerExtension .hashCode(); } + private static void closeReusableConnection() { + if (reusableConnection != null) { + try { + reusableConnection.close(); + } catch (SQLException e) { + // Ignore close errors during cleanup + } + reusableConnection = null; + } + } + + private static Connection getReusableConnection() throws SQLException { + if (reusableConnection == null || reusableConnection.isClosed()) { + reusableConnection = createConnection(); + } + return reusableConnection; + } + /** * Drops and recreates the 'public' schema and all tables, then creates a new {@link * EntityManagerFactory} and save it in {@link #emf}. */ private void recreateSchema() throws Exception { + closeReusableConnection(); + cachedResetSql = null; if (emf != null) { emf.close(); emf = null; @@ -178,6 +212,10 @@ public abstract class JpaTransactionManagerExtension exporter.export(extraEntityClasses, tempSqlFile); executeSql(Files.readString(tempSqlFile.toPath(), UTF_8)); } + // Lower project_wide_unique_id_seq's non-1 MINVALUE from the golden schema so deterministic + // sequence resets to 1 succeed. + executeSql( + "ALTER SEQUENCE IF EXISTS project_wide_unique_id_seq START 1 MINVALUE 1 RESTART WITH 1;"); assertReasonableNumDbConnections(); emf = createEntityManagerFactory(getJpaProperties()); emfEntityHash = entityHash; @@ -228,16 +266,6 @@ public abstract class JpaTransactionManagerExtension cachedTm = TransactionManagerFactory.tm(); TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(txnManager)); TransactionManagerFactory.setReplicaJpaTm(Suppliers.ofInstance(readOnlyTxnManager)); - // Reset SQL Sequence based id allocation so that ids are deterministic in tests. - TransactionManagerFactory.tm() - .transact( - () -> - TransactionManagerFactory.tm() - .getEntityManager() - .createNativeQuery( - "alter sequence if exists project_wide_unique_id_seq start 1 minvalue 1" - + " restart with 1") - .executeUpdate()); if (withCannedData) { loadInitialData(); } @@ -256,22 +284,34 @@ public abstract class JpaTransactionManagerExtension } private void resetTablesAndSequences() { - try (Connection conn = createConnection(); - Statement statement = conn.createStatement()) { - ResultSet rs = - statement.executeQuery( - "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public';"); - ImmutableList.Builder tableNamesBuilder = new ImmutableList.Builder<>(); - while (rs.next()) { - tableNamesBuilder.add('"' + rs.getString(1) + '"'); - } - ImmutableList tableNames = tableNamesBuilder.build(); - if (!tableNames.isEmpty()) { - String sql = - String.format("TRUNCATE %s RESTART IDENTITY CASCADE", Joiner.on(',').join(tableNames)); - executeSql(sql); + try (Statement statement = getReusableConnection().createStatement()) { + if (cachedResetSql == null) { + // We build up a list of tables/sequences to reset by querying SQL. We only need to do this + // once per schema. + StringBuilder sqlBuilder = new StringBuilder("BEGIN; SET CONSTRAINTS ALL DEFERRED; "); + try (ResultSet rs = + statement.executeQuery( + "SELECT table_name FROM information_schema.tables" + + " WHERE table_schema = 'public' AND table_type = 'BASE TABLE';")) { + while (rs.next()) { + sqlBuilder.append(String.format("DELETE FROM \"%s\";", rs.getString(1))); + } + } + try (ResultSet rs = + statement.executeQuery( + "SELECT sequence_name FROM information_schema.sequences" + + " WHERE sequence_schema = 'public';")) { + while (rs.next()) { + sqlBuilder.append( + String.format("ALTER SEQUENCE \"%s\" RESTART WITH 1;", rs.getString(1))); + } + } + sqlBuilder.append("COMMIT;"); + cachedResetSql = sqlBuilder.toString(); } + statement.execute(cachedResetSql); } catch (Exception e) { + closeReusableConnection(); throw new RuntimeException(e); } } diff --git a/db/src/main/resources/sql/schema/db-schema.sql.generated b/db/src/main/resources/sql/schema/db-schema.sql.generated index 7e484566f..c9d74594e 100644 --- a/db/src/main/resources/sql/schema/db-schema.sql.generated +++ b/db/src/main/resources/sql/schema/db-schema.sql.generated @@ -951,49 +951,49 @@ alter table if exists "ConsoleUpdateHistory" add constraint FKnhl1eolgix64u90xv3pj6xa3x foreign key (acting_user) - references "User"; + references "User" deferrable initially deferred; alter table if exists "DelegationSignerData" add constraint FKtr24j9v14ph2mfuw2gsmt12kq foreign key (domain_repo_id) - references "Domain"; + references "Domain" deferrable initially deferred; alter table if exists "DomainDsDataHistory" add constraint FKo4ilgyyfnvppbpuivus565i0j foreign key (domain_repo_id, domain_history_revision_id) - references "DomainHistory"; + references "DomainHistory" deferrable initially deferred; alter table if exists "DomainHistoryHost" add constraint FKa9woh3hu8gx5x0vly6bai327n foreign key (domain_history_domain_repo_id, domain_history_history_revision_id) - references "DomainHistory"; + references "DomainHistory" deferrable initially deferred; alter table if exists "DomainHost" add constraint FKeq1guccbre1yk3oosgp2io554 foreign key (domain_repo_id) - references "Domain"; + references "Domain" deferrable initially deferred; alter table if exists "DomainTransactionRecord" add constraint FKcjqe54u72kha71vkibvxhjye7 foreign key (domain_repo_id, history_revision_id) - references "DomainHistory"; + references "DomainHistory" deferrable initially deferred; alter table if exists "GracePeriod" add constraint FK2mys4hojm6ev2g9tmy5aq6m7g foreign key (domain_repo_id) - references "Domain"; + references "Domain" deferrable initially deferred; alter table if exists "GracePeriodHistory" add constraint FK7w3cx8d55q8bln80e716tr7b8 foreign key (domain_repo_id, domain_history_revision_id) - references "DomainHistory"; + references "DomainHistory" deferrable initially deferred; alter table if exists "RegistryLock" add constraint FK2lhcwpxlnqijr96irylrh1707 foreign key (relock_revision_id) - references "RegistryLock"; + references "RegistryLock" deferrable initially deferred; alter table if exists "SignedMarkRevocationEntry" add constraint FK5ivlhvs3121yx2li5tqh54u4 foreign key (revision_id) - references "SignedMarkRevocationList"; + references "SignedMarkRevocationList" deferrable initially deferred; diff --git a/db/src/test/java/google/registry/sql/flyway/SchemaTest.java b/db/src/test/java/google/registry/sql/flyway/SchemaTest.java index 90ac0b8ce..afbeabcab 100644 --- a/db/src/test/java/google/registry/sql/flyway/SchemaTest.java +++ b/db/src/test/java/google/registry/sql/flyway/SchemaTest.java @@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static google.registry.testing.truth.TextDiffSubject.assertThat; import com.google.common.base.Joiner; +import com.google.common.collect.ImmutableMap; import com.google.common.flogger.FluentLogger; import com.google.common.io.Resources; import google.registry.persistence.NomulusPostgreSql; @@ -81,7 +82,9 @@ class SchemaTest { private final PostgreSQLContainer sqlContainer = new PostgreSQLContainer<>(NomulusPostgreSql.getDockerImageName()) .withClasspathResourceMapping( - MOUNTED_RESOURCE_PATH, CONTAINER_MOUNT_POINT, BindMode.READ_WRITE); + MOUNTED_RESOURCE_PATH, CONTAINER_MOUNT_POINT, BindMode.READ_WRITE) + // Use tmpFs so that we keep everything in memory. No point in writing to disk + .withTmpFs(ImmutableMap.of("/var/lib/postgresql/data", "rw")); @Test @DisabledIfSystemProperty(named = "deploy_to_existing_db", matches = ".*")