Thread Safe JUnit Extensions

JUnit extensions are now fully thread safe through the use of `ExtensionContext.Store`. `SingleInstancePostgresExtension` has had its functionality extended to be able to run for the entire class instead of just per method.

Author: SmushyTaco

README.md

@@ -34,7 +34,7 @@ dependencies {
 And the following to your `gradle.properties`:
 ```properties
 # Check this on https://central.sonatype.com/artifact/com.smushytaco/embedded-postgres/
-embedded_postgres_version = 3.0.1
+embedded_postgres_version = 3.0.2
 ```
 
 The default version of the embedded postgres is `PostgreSQL 18.0.0`, but you can change it by following the instructions described in [Postgres version](#postgres-version).
@@ -45,7 +45,7 @@ In your JUnit test just add:
 
 ```java
 @RegisterExtension
-SingleInstancePostgresExtension pg = EmbeddedPostgresExtension.singleInstance();
+final SingleInstancePostgresExtension pg = EmbeddedPostgresExtension.singleInstance();
 ```
 
 This simply has JUnit manage an instance of EmbeddedPostgres (start, stop). You can then use this to get a DataSource with: `pg.getEmbeddedPostgres().getPostgresDatabase();`  
@@ -60,13 +60,13 @@ You can easily integrate Flyway or Liquibase database schema migration:
 ##### Flyway
 ```java
 @RegisterExtension
-PreparedDbExtension db = EmbeddedPostgresExtension.preparedDatabase(FlywayPreparer.forClasspathLocation("db/my-db-schema"));
+final PreparedDbExtension db = EmbeddedPostgresExtension.preparedDatabase(FlywayPreparer.forClasspathLocation("db/my-db-schema"));
 ```
 
 ##### Liquibase
 ```java
 @RegisterExtension
-PreparedDbExtension db = EmbeddedPostgresExtension.preparedDatabase(LiquibasePreparer.forClasspathLocation("liqui/master.xml"));
+final PreparedDbExtension db = EmbeddedPostgresExtension.preparedDatabase(LiquibasePreparer.forClasspathLocation("liqui/master.xml"));
 ```
 
 This will create an independent database for every test with the given schema loaded from the classpath.

gradle.properties

@@ -7,7 +7,7 @@ org.gradle.configuration-cache = true
 # Maven Properties
 group = com.smushytaco
 name = embedded-postgres
-version = 3.0.1
+version = 3.0.2
 ##########################################################################
 # Java Version
 java_version = 25

src/main/java/com/smushytaco/postgres/embedded/DefaultPostgresBinaryResolver.java

@@ -41,6 +41,7 @@
  * (e.g., {@code postgres-Linux-x86_64.txz}) and provides them as an {@link InputStream}.
  * It also includes fallback behavior for unsupported or emulated architectures.
  */
+@SuppressWarnings("java:S6548")
 public class DefaultPostgresBinaryResolver implements PgBinaryResolver {
     private static final Logger logger = LoggerFactory.getLogger(DefaultPostgresBinaryResolver.class);
 

src/main/java/com/smushytaco/postgres/embedded/EmbeddedPostgres.java

@@ -383,6 +383,7 @@ private void pgCtl(final Path dir, final String action) {
         system(pgCtl, args, null);
     }
 
+    @SuppressWarnings("java:S1141")
     private void cleanOldDataDirectories(final Path parentDirectory) {
         try (final Stream<Path> children = Files.list(parentDirectory)) {
             for (final Path dir : children.toList()) {
@@ -844,6 +845,7 @@ private void closeChannel(final Channel channel) {
         }
     }
 
+    @SuppressWarnings("java:S1141")
     private static Path prepareBinaries(final PgBinaryResolver pgBinaryResolver, final Path overrideWorkingDirectory) {
         PREPARE_BINARIES_LOCK.lock();
         try {

src/main/java/com/smushytaco/postgres/embedded/PreparedDbProvider.java

@@ -259,6 +259,7 @@ private static String randomAlphabetic(final int length) {
                     .toLowerCase(Locale.ENGLISH);
         }
 
+        @SuppressWarnings("java:S2189")
         @Override
         public void run() {
             while (true) {
@@ -401,6 +402,7 @@ public Map<String, String> getProperties() {
          *
          * @return the exception that caused creation to fail, or {@code null} on success
          */
+        @SuppressWarnings("unused")
         public SQLException getException() {
             return ex;
         }

src/main/java/com/smushytaco/postgres/junit/PreparedDbExtension.java

@@ -27,6 +27,7 @@
 import java.sql.SQLException;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 
 /**
@@ -37,14 +38,15 @@
  * per test class or each test can get its own fresh database instance.
  */
 public class PreparedDbExtension implements BeforeAllCallback, AfterAllCallback, BeforeEachCallback, AfterEachCallback {
-    private final DatabasePreparer preparer;
-    private boolean perClass = false;
-    private volatile DataSource dataSource;
-    private volatile PreparedDbProvider provider;
-    private volatile ConnectionInfo connectionInfo;
+    private final Object stateKey = new Object();
 
+    private final DatabasePreparer preparer;
     private final List<Consumer<EmbeddedPostgres.Builder>> builderCustomizers = new CopyOnWriteArrayList<>();
 
+    private final AtomicBoolean started = new AtomicBoolean(false);
+
+    private final ThreadLocal<State> current = new ThreadLocal<>();
+
     PreparedDbExtension(final DatabasePreparer preparer) {
         if (preparer == null) throw new IllegalStateException("null preparer");
         this.preparer = preparer;
@@ -62,41 +64,49 @@ public class PreparedDbExtension implements BeforeAllCallback, AfterAllCallback,
      * @throws AssertionError if the extension has already been started
      */
     public PreparedDbExtension customize(final Consumer<EmbeddedPostgres.Builder> customizer) {
-        if (dataSource != null) throw new AssertionError("already started");
+        if (started.get()) throw new AssertionError("already started");
         builderCustomizers.add(customizer);
         return this;
     }
 
+    @SuppressWarnings("DuplicatedCode")
     @Override
-    public void beforeAll(@NonNull final ExtensionContext extensionContext) throws SQLException {
-        provider = PreparedDbProvider.forPreparer(preparer, builderCustomizers);
-        connectionInfo = provider.createNewDatabase();
-        dataSource = provider.createDataSourceFromConnectionInfo(connectionInfo);
-        perClass = true;
+    public void beforeAll(@NonNull final ExtensionContext ctx) throws SQLException {
+        final ExtensionContext.Store classStore = classStore(ctx);
+        State state = classStore.get(stateKey, State.class);
+        if (state == null) {
+            state = createState();
+            classStore.put(stateKey, state);
+            started.set(true);
+        }
+        current.set(state);
     }
 
     @Override
-    public void afterAll(@NonNull final ExtensionContext extensionContext) {
-        dataSource = null;
-        connectionInfo = null;
-        provider = null;
-        perClass = false;
+    public void afterAll(@NonNull final ExtensionContext ctx) {
+        current.remove();
     }
 
+    @SuppressWarnings("DuplicatedCode")
     @Override
-    public void beforeEach(@NonNull final ExtensionContext extensionContext) throws SQLException {
-        if (perClass) return;
-        provider = PreparedDbProvider.forPreparer(preparer, builderCustomizers);
-        connectionInfo = provider.createNewDatabase();
-        dataSource = provider.createDataSourceFromConnectionInfo(connectionInfo);
+    public void beforeEach(@NonNull final ExtensionContext ctx) throws SQLException {
+        final ExtensionContext.Store classStore = classStore(ctx);
+        State state = classStore.get(stateKey, State.class);
+        if (state == null) {
+            final ExtensionContext.Store methodStore = methodStore(ctx);
+            state = methodStore.get(stateKey, State.class);
+            if (state == null) {
+                state = createState();
+                methodStore.put(stateKey, state);
+                started.set(true);
+            }
+        }
+        current.set(state);
     }
 
     @Override
-    public void afterEach(@NonNull final ExtensionContext extensionContext) {
-        if (perClass) return;
-        dataSource = null;
-        connectionInfo = null;
-        provider = null;
+    public void afterEach(@NonNull final ExtensionContext ctx) {
+        current.remove();
     }
 
     /**
@@ -108,8 +118,9 @@ public void afterEach(@NonNull final ExtensionContext extensionContext) {
      * @throws AssertionError if the extension has not been initialized yet
      */
     public DataSource getTestDatabase() {
-        if (dataSource == null) throw new AssertionError("not initialized");
-        return dataSource;
+        final State s = current.get();
+        if (s == null) throw new AssertionError("not initialized");
+        return s.dataSource;
     }
 
     /**
@@ -120,8 +131,9 @@ public DataSource getTestDatabase() {
      * @throws AssertionError if the extension has not been initialized yet
      */
     public ConnectionInfo getConnectionInfo() {
-        if (connectionInfo == null) throw new AssertionError("not initialized");
-        return connectionInfo;
+        final State s = current.get();
+        if (s == null) throw new AssertionError("not initialized");
+        return s.connectionInfo;
     }
 
     /**
@@ -132,7 +144,34 @@ public ConnectionInfo getConnectionInfo() {
      * @throws AssertionError if the extension has not been initialized yet
      */
     public PreparedDbProvider getDbProvider() {
-        if (provider == null) throw new AssertionError("not initialized");
-        return provider;
+        final State s = current.get();
+        if (s == null) throw new AssertionError("not initialized");
+        return s.provider;
+    }
+
+    private State createState() throws SQLException {
+        final PreparedDbProvider provider = PreparedDbProvider.forPreparer(preparer, builderCustomizers);
+        final ConnectionInfo connectionInfo = provider.createNewDatabase();
+        return new State(provider, connectionInfo, provider.createDataSourceFromConnectionInfo(connectionInfo));
+    }
+
+    private static ExtensionContext.Store classStore(final ExtensionContext ctx) {
+        return ctx.getStore(ExtensionContext.Namespace.create(PreparedDbExtension.class, ctx.getRequiredTestClass()));
+    }
+
+    private static ExtensionContext.Store methodStore(final ExtensionContext ctx) {
+        return ctx.getStore(ExtensionContext.Namespace.create(PreparedDbExtension.class, ctx.getRequiredTestMethod()));
+    }
+
+    private record State(PreparedDbProvider provider, ConnectionInfo connectionInfo, DataSource dataSource) implements AutoCloseable {
+        @Override
+        public void close() {
+            try {
+                if (provider instanceof final AutoCloseable autoCloseable) autoCloseable.close();
+            } catch (final Exception _) { /* noop */ }
+            try {
+                if (dataSource instanceof final AutoCloseable autoCloseable) autoCloseable.close();
+            } catch (final Exception _) { /* noop */ }
+        }
     }
 }