Added property registerShutdownHook to EmbeddedPostgres. Default the shutdownHook is registered. When this property is set to false no shutdown hook is registered. This makes it possible to create a spring bean of EmbeddedPostgres and let spring handle the shutdown of the EmbeddedPostgres instance in the right order and let the application shutdown gracefully

Author: JelmerBouma1985

src/main/java/io/zonky/test/db/postgres/embedded/EmbeddedPostgres.java

@@ -106,26 +106,28 @@ public class EmbeddedPostgres implements Closeable
     private volatile FileOutputStream lockStream;
     private volatile FileLock lock;
     private final boolean cleanDataDirectory;
+    private final boolean registerShutdownHook;
 
     private final ProcessBuilder.Redirect errorRedirector;
     private final ProcessBuilder.Redirect outputRedirector;
     private final Process process;
 
-    EmbeddedPostgres(File parentDirectory, File dataDirectory, boolean cleanDataDirectory,
+    EmbeddedPostgres(File parentDirectory, File dataDirectory, boolean cleanDataDirectory, boolean registerShutdownHook,
         Map<String, String> postgresConfig, Map<String, String> localeConfig, int port, Map<String, String> connectConfig,
         PgBinaryResolver pgBinaryResolver, ProcessBuilder.Redirect errorRedirector, ProcessBuilder.Redirect outputRedirector) throws IOException
     {
-        this(parentDirectory, dataDirectory, cleanDataDirectory, postgresConfig, localeConfig, port, connectConfig,
+        this(parentDirectory, dataDirectory, cleanDataDirectory, registerShutdownHook, postgresConfig, localeConfig, port, connectConfig,
                 pgBinaryResolver, errorRedirector, outputRedirector, DEFAULT_PG_STARTUP_WAIT, null, null);
     }
 
-    EmbeddedPostgres(File parentDirectory, File dataDirectory, boolean cleanDataDirectory,
+    EmbeddedPostgres(File parentDirectory, File dataDirectory, boolean cleanDataDirectory, boolean registerShutdownHook,
                      Map<String, String> postgresConfig, Map<String, String> localeConfig, int port, Map<String, String> connectConfig,
                      PgBinaryResolver pgBinaryResolver, ProcessBuilder.Redirect errorRedirector,
                      ProcessBuilder.Redirect outputRedirector, Duration pgStartupWait,
                      File overrideWorkingDirectory, Consumer<File> dataDirectoryCustomizer) throws IOException
     {
         this.cleanDataDirectory = cleanDataDirectory;
+        this.registerShutdownHook = registerShutdownHook;
         this.postgresConfig = new HashMap<>(postgresConfig);
         this.localeConfig = new HashMap<>(localeConfig);
         this.connectConfig = new HashMap<>(connectConfig);
@@ -291,7 +293,9 @@ private Process startPostmaster() throws IOException
 
         LOG.info("{} postmaster started as {} on port {}.  Waiting up to {} for server startup to finish.", instanceId, postmaster.toString(), port, pgStartupWait);
 
-        Runtime.getRuntime().addShutdownHook(newCloserThread());
+        if (registerShutdownHook) {
+            Runtime.getRuntime().addShutdownHook(newCloserThread());
+        }
 
         waitForServerStartup(watch);
         return postmaster;
@@ -503,6 +507,7 @@ public static class Builder
         private final Map<String, String> config = new HashMap<>();
         private final Map<String, String> localeConfig = new HashMap<>();
         private boolean builderCleanDataDirectory = true;
+        private boolean builderRegisterShutdownHook = true;
         private int builderPort = 0;
         private final Map<String, String> connectConfig = new HashMap<>();
         private PgBinaryResolver pgBinaryResolver = DefaultPostgresBinaryResolver.INSTANCE;
@@ -532,6 +537,11 @@ public Builder setCleanDataDirectory(boolean cleanDataDirectory) {
             return this;
         }
 
+        public Builder setRegisterShutdownHook(boolean registerShutdownHook) {
+            builderRegisterShutdownHook = registerShutdownHook;
+            return this;
+        }
+
         public Builder setDataDirectory(Path path) {
             return setDataDirectory(path.toFile());
         }
@@ -599,8 +609,8 @@ public EmbeddedPostgres start() throws IOException {
             if (builderDataDirectory == null) {
                 builderDataDirectory = Files.createTempDirectory("epg").toFile();
             }
-            return new EmbeddedPostgres(parentDirectory, builderDataDirectory, builderCleanDataDirectory, config,
-                    localeConfig, builderPort, connectConfig, pgBinaryResolver, errRedirector, outRedirector,
+            return new EmbeddedPostgres(parentDirectory, builderDataDirectory, builderCleanDataDirectory, builderRegisterShutdownHook,
+                    config, localeConfig, builderPort, connectConfig, pgBinaryResolver, errRedirector, outRedirector,
                     pgStartupWait, overrideWorkingDirectory, dataDirectoryCustomizer);
         }
 
@@ -614,6 +624,7 @@ public boolean equals(Object o) {
             }
             Builder builder = (Builder) o;
             return builderCleanDataDirectory == builder.builderCleanDataDirectory &&
+                    builderRegisterShutdownHook == builder.builderRegisterShutdownHook &&
                     builderPort == builder.builderPort &&
                     Objects.equals(parentDirectory, builder.parentDirectory) &&
                     Objects.equals(builderDataDirectory, builder.builderDataDirectory) &&
@@ -630,7 +641,7 @@ public boolean equals(Object o) {
 
         @Override
         public int hashCode() {
-            return Objects.hash(parentDirectory, builderDataDirectory, config, localeConfig, builderCleanDataDirectory, builderPort, connectConfig, pgBinaryResolver, pgStartupWait, errRedirector, outRedirector);
+            return Objects.hash(parentDirectory, builderDataDirectory, config, localeConfig, builderCleanDataDirectory, builderRegisterShutdownHook, builderPort, connectConfig, pgBinaryResolver, pgStartupWait, errRedirector, outRedirector);
         }
     }