From 72ad353d0822534d6c76fd5de687e1e149ed7cad Mon Sep 17 00:00:00 2001
From: Amund Murstad <amund.murstad@finn.no>
Date: Wed, 7 Jul 2021 21:47:39 +0200
Subject: [PATCH] Make PreparedDbProvider AutoCloseable such that the internal
 background thread started in PrepPipeline can be shut down

---
 .../postgres/embedded/PreparedDbProvider.java | 26 +++++++++++++++----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/main/java/io/zonky/test/db/postgres/embedded/PreparedDbProvider.java b/src/main/java/io/zonky/test/db/postgres/embedded/PreparedDbProvider.java
index 0e306cc5..3fa2cd01 100644
--- a/src/main/java/io/zonky/test/db/postgres/embedded/PreparedDbProvider.java
+++ b/src/main/java/io/zonky/test/db/postgres/embedded/PreparedDbProvider.java
@@ -38,7 +38,7 @@
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.unmodifiableMap;
 
-public class PreparedDbProvider
+public class PreparedDbProvider implements AutoCloseable
 {
     private static final String JDBC_FORMAT = "jdbc:postgresql://localhost:%d/%s?user=%s";
 
@@ -50,6 +50,7 @@ public class PreparedDbProvider
     private static final Map<ClusterKey, PrepPipeline> CLUSTERS = new HashMap<>();
 
     private final PrepPipeline dbPreparer;
+    private final ClusterKey key;
 
     public static PreparedDbProvider forPreparer(DatabasePreparer preparer) {
         return forPreparer(preparer, Collections.emptyList());
@@ -61,7 +62,8 @@ public static PreparedDbProvider forPreparer(DatabasePreparer preparer, Iterable
 
     private PreparedDbProvider(DatabasePreparer preparer, Iterable<Consumer<Builder>> customizers) {
         try {
-            dbPreparer = createOrFindPreparer(preparer, customizers);
+            key = new ClusterKey(preparer, customizers);
+            dbPreparer = createOrFindPreparer(preparer, customizers, key);
         } catch (final IOException | SQLException e) {
             throw new RuntimeException(e);
         }
@@ -71,9 +73,8 @@ private PreparedDbProvider(DatabasePreparer preparer, Iterable<Consumer<Builder>
      * Each schema set has its own database cluster.  The template1 database has the schema preloaded so that
      * each test case need only create a new database and not re-invoke your preparer.
      */
-    private static synchronized PrepPipeline createOrFindPreparer(DatabasePreparer preparer, Iterable<Consumer<Builder>> customizers) throws IOException, SQLException
+    private static synchronized PrepPipeline createOrFindPreparer(DatabasePreparer preparer, Iterable<Consumer<Builder>> customizers, ClusterKey key) throws IOException, SQLException
     {
-        final ClusterKey key = new ClusterKey(preparer, customizers);
         PrepPipeline result = CLUSTERS.get(key);
         if (result != null) {
             return result;
@@ -164,14 +165,21 @@ public Map<String, String> getConfigurationTweak(String dbModuleName) throws SQL
         return result;
     }
 
+    @Override
+    public void close() throws IOException {
+        CLUSTERS.remove(key);
+        dbPreparer.close();
+    }
+
     /**
      * Spawns a background thread that prepares databases ahead of time for speed, and then uses a
      * synchronous queue to hand the prepared databases off to test cases.
      */
-    private static class PrepPipeline implements Runnable
+    private static class PrepPipeline implements Runnable, AutoCloseable
     {
         private final EmbeddedPostgres pg;
         private final SynchronousQueue<DbInfo> nextDatabase = new SynchronousQueue<DbInfo>();
+        private Thread runningThread = null;
 
         PrepPipeline(EmbeddedPostgres pg)
         {
@@ -184,6 +192,7 @@ PrepPipeline start()
                 final Thread t = new Thread(r);
                 t.setDaemon(true);
                 t.setName("cluster-" + pg + "-preparer");
+                runningThread = t;
                 return t;
             });
             service.submit(this);
@@ -228,6 +237,13 @@ public void run()
                 }
             }
         }
+
+        @Override
+        public void close() throws IOException {
+            if(runningThread != null){
+                runningThread.interrupt();
+            }
+        }
     }
 
     private static void create(final DataSource connectDb, final String dbName, final String userName) throws SQLException