Clean Up + Tests

Author: SmushyTaco

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

@@ -42,6 +42,7 @@
 import java.util.Map.Entry;
 import java.util.concurrent.Phaser;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Consumer;
@@ -92,6 +93,8 @@ public class EmbeddedPostgres implements Closeable {
     private final ProcessBuilder.Redirect errorRedirector;
     private final ProcessBuilder.Redirect outputRedirector;
 
+    private final AtomicReference<Thread> shutdownHook = new AtomicReference<>();
+
     @SuppressWarnings("unused")
     EmbeddedPostgres(final Path parentDirectory, final Path dataDirectory, final boolean cleanDataDirectory, final boolean registerShutdownHook,
                      final Map<String, String> postgresConfig, final Map<String, String> localeConfig, final int port, final Map<String, String> connectConfig,
@@ -285,7 +288,11 @@ private void startPostmaster() {
                 "-w", "start"
         ));
 
-        if (registerShutdownHook) Runtime.getRuntime().addShutdownHook(newCloserThread());
+        if (registerShutdownHook) {
+            final Thread hook = newCloserThread();
+            Runtime.getRuntime().addShutdownHook(hook);
+            shutdownHook.set(hook);
+        }
 
         system(pgCtl, args, pgStartupWait);
 
@@ -346,6 +353,7 @@ private static void deleteDirectoryRecursively(final Path dir) throws IOExceptio
     @Override
     public void close() throws IOException {
         if (closed.getAndSet(true)) return;
+        unregisterShutdownHook();
         final Instant start = Instant.now();
         try {
             pgCtl(dataDirectory, "stop");
@@ -375,6 +383,19 @@ public void close() throws IOException {
         }
     }
 
+    private void unregisterShutdownHook() {
+        final Thread currentShutdownHook = shutdownHook.getAndSet(null);
+        if (currentShutdownHook == null || Thread.currentThread() == currentShutdownHook) return;
+
+        try {
+            Runtime.getRuntime().removeShutdownHook(currentShutdownHook);
+        } catch (final IllegalStateException e) {
+            LOG.trace("JVM is already shutting down; could not remove shutdown hook {}", currentShutdownHook.getName(), e);
+        } catch (final Exception e) {
+            LOG.warn("Failed to remove shutdown hook {}", currentShutdownHook.getName(), e);
+        }
+    }
+
     @SuppressWarnings("SameParameterValue")
     private void pgCtl(final Path dir, final String action) {
         final List<String> args = new ArrayList<>(Arrays.asList(

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

@@ -240,6 +240,7 @@ public static void closeAll() throws IOException {
     }
 
     private void ensureOpen() {
+        if (sharedCluster.closed.get()) closed.set(true);
         if (closed.get()) throw new IllegalStateException("PreparedDbProvider has been closed");
     }
 

src/test/java/com/smushytaco/postgres/embedded/EmbeddedPostgresShutdownHookTest.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2025 Nikan Radan
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.smushytaco.postgres.embedded;
+
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+class EmbeddedPostgresShutdownHookTest {
+    @Test
+    void testExplicitCloseUnregistersShutdownHookAndRemainsIdempotent() throws IOException, NoSuchFieldException, IllegalAccessException {
+        try (final EmbeddedPostgres pg = EmbeddedPostgres.builder()
+                .setRegisterShutdownHook(true)
+                .start()) {
+            final AtomicReference<?> shutdownHook = getShutdownHook(pg);
+            assertNotNull(shutdownHook.get());
+
+            pg.close();
+            assertNull(shutdownHook.get());
+            assertDoesNotThrow(pg::close);
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private static AtomicReference<Thread> getShutdownHook(final EmbeddedPostgres pg) throws NoSuchFieldException, IllegalAccessException {
+        final Field shutdownHookField = EmbeddedPostgres.class.getDeclaredField("shutdownHook");
+        shutdownHookField.setAccessible(true);
+        return (AtomicReference<Thread>) shutdownHookField.get(pg);
+    }
+}

src/test/java/com/smushytaco/postgres/embedded/PreparedDbProviderFailureCleanupTest.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2025 Nikan Radan
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.smushytaco.postgres.embedded;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.net.ServerSocket;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class PreparedDbProviderFailureCleanupTest {
+    @AfterEach
+    void tearDown() throws IOException {
+        PreparedDbProvider.closeAll();
+    }
+
+    @SuppressWarnings({"SqlNoDataSourceInspection", "java:S5778"})
+    @Test
+    void testFailedPreparationDoesNotLeakStartedPostgresInstance() throws IOException, SQLException {
+        final int port = detectFreePort();
+
+        @SuppressWarnings("resource") final RuntimeException exception = assertThrows(RuntimeException.class, () -> PreparedDbProvider.forPreparer(_ -> {
+            throw new SQLException("boom");
+        }, List.of(builder -> {
+            builder.setPort(port);
+            builder.setRegisterShutdownHook(false);
+        }), ClusterRetentionPolicy.CLOSE_ON_LAST_RELEASE));
+        assertInstanceOf(SQLException.class, exception.getCause());
+
+        try (final EmbeddedPostgres pg = EmbeddedPostgres.builder()
+                .setPort(port)
+                .setRegisterShutdownHook(false)
+                .start();
+                final Connection connection = pg.getPostgresDatabase().getConnection();
+                final Statement statement = connection.createStatement();
+                final ResultSet resultSet = statement.executeQuery("SELECT 1")) {
+            assertTrue(resultSet.next());
+            assertEquals(1, resultSet.getInt(1));
+        }
+    }
+
+    private static int detectFreePort() throws IOException {
+        try (final ServerSocket socket = new ServerSocket(0)) {
+            socket.setReuseAddress(true);
+            return socket.getLocalPort();
+        }
+    }
+}

src/test/java/com/smushytaco/postgres/embedded/PreparedDbProviderLifecycleTest.java

@@ -0,0 +1,98 @@
+/*
+ * Copyright 2025 Nikan Radan
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.smushytaco.postgres.embedded;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import javax.sql.DataSource;
+import java.io.IOException;
+import java.sql.SQLException;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+class PreparedDbProviderLifecycleTest {
+    @AfterEach
+    void tearDown() throws IOException {
+        PreparedDbProvider.closeAll();
+    }
+
+    @Test
+    void testKeepUntilCloseAllRetainsClusterUntilExplicitGlobalClose() throws SQLException, IOException {
+        final CountingPreparer preparer = new CountingPreparer();
+
+        final PreparedDbProvider provider1 = PreparedDbProvider.forPreparer(preparer, ClusterRetentionPolicy.KEEP_UNTIL_CLOSE_ALL);
+        provider1.createNewDatabase();
+        assertEquals(1, preparer.getInvocationCount());
+
+        provider1.close();
+        assertThrows(IllegalStateException.class, provider1::createNewDatabase);
+
+        final PreparedDbProvider provider2 = PreparedDbProvider.forPreparer(preparer, ClusterRetentionPolicy.KEEP_UNTIL_CLOSE_ALL);
+        provider2.createNewDatabase();
+        assertEquals(1, preparer.getInvocationCount());
+
+        PreparedDbProvider.closeAll();
+        assertThrows(IllegalStateException.class, provider2::createNewDatabase);
+
+        final PreparedDbProvider provider3 = PreparedDbProvider.forPreparer(preparer, ClusterRetentionPolicy.KEEP_UNTIL_CLOSE_ALL);
+        provider3.createNewDatabase();
+        assertEquals(2, preparer.getInvocationCount());
+        provider3.close();
+    }
+
+    @Test
+    void testCloseOnLastReleaseRecreatesClusterAfterFinalClose() throws SQLException, IOException {
+        final CountingPreparer preparer = new CountingPreparer();
+
+        final PreparedDbProvider provider1 = PreparedDbProvider.forPreparer(preparer, ClusterRetentionPolicy.CLOSE_ON_LAST_RELEASE);
+        final PreparedDbProvider provider2 = PreparedDbProvider.forPreparer(preparer, ClusterRetentionPolicy.CLOSE_ON_LAST_RELEASE);
+
+        provider1.createNewDatabase();
+        provider2.createNewDatabase();
+        assertEquals(1, preparer.getInvocationCount());
+
+        provider1.close();
+        assertThrows(IllegalStateException.class, provider1::createNewDatabase);
+
+        provider2.createNewDatabase();
+        assertEquals(1, preparer.getInvocationCount());
+
+        provider2.close();
+        assertThrows(IllegalStateException.class, provider2::createNewDatabase);
+
+        final PreparedDbProvider provider3 = PreparedDbProvider.forPreparer(preparer, ClusterRetentionPolicy.CLOSE_ON_LAST_RELEASE);
+        provider3.createNewDatabase();
+        assertEquals(2, preparer.getInvocationCount());
+        provider3.close();
+    }
+
+    private static final class CountingPreparer implements DatabasePreparer {
+        private final AtomicInteger invocationCount = new AtomicInteger();
+
+        @Override
+        public void prepare(final DataSource ds) {
+            invocationCount.incrementAndGet();
+        }
+
+        private int getInvocationCount() {
+            return invocationCount.get();
+        }
+    }
+}