Fix caching of prepared databases

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

@@ -51,7 +51,6 @@
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 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.stream.Collectors;
@@ -571,6 +570,33 @@ public EmbeddedPostgres start() throws IOException {
             }
             return new EmbeddedPostgres(parentDirectory, builderDataDirectory, builderCleanDataDirectory, config, localeConfig, builderPort, connectConfig, pgBinaryResolver, errRedirector, outRedirector, pgStartupWait);
         }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            Builder builder = (Builder) o;
+            return builderCleanDataDirectory == builder.builderCleanDataDirectory &&
+                    builderPort == builder.builderPort &&
+                    Objects.equals(parentDirectory, builder.parentDirectory) &&
+                    Objects.equals(builderDataDirectory, builder.builderDataDirectory) &&
+                    Objects.equals(config, builder.config) &&
+                    Objects.equals(localeConfig, builder.localeConfig) &&
+                    Objects.equals(connectConfig, builder.connectConfig) &&
+                    Objects.equals(pgBinaryResolver, builder.pgBinaryResolver) &&
+                    Objects.equals(pgStartupWait, builder.pgStartupWait) &&
+                    Objects.equals(errRedirector, builder.errRedirector) &&
+                    Objects.equals(outRedirector, builder.outRedirector);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(parentDirectory, builderDataDirectory, config, localeConfig, builderCleanDataDirectory, builderPort, connectConfig, pgBinaryResolver, pgStartupWait, errRedirector, outRedirector);
+        }
     }
 
     private static List<String> system(String... command)
@@ -601,8 +627,8 @@ private static void mkdirs(File dir)
         }
     }
 
-    private static final AtomicReference<File> BINARY_DIR = new AtomicReference<>();
     private static final Lock PREPARE_BINARIES_LOCK = new ReentrantLock();
+    private static final Map<PgBinaryResolver, File> PREPARE_BINARIES = new HashMap<>();
 
     /**
      * Get current operating system string. The string is used in the appropriate postgres binary name.
@@ -685,8 +711,8 @@ private static File prepareBinaries(PgBinaryResolver pgBinaryResolver)
     {
         PREPARE_BINARIES_LOCK.lock();
         try {
-            if(BINARY_DIR.get() != null) {
-                return BINARY_DIR.get();
+            if (PREPARE_BINARIES.containsKey(pgBinaryResolver)) {
+                return PREPARE_BINARIES.get(pgBinaryResolver);
             }
 
             final String system = getOS();
@@ -765,7 +791,7 @@ private static File prepareBinaries(PgBinaryResolver pgBinaryResolver)
                     LOG.warn("could not delete {}", pgTbz);
                 }
             }
-            BINARY_DIR.set(pgDir);
+            PREPARE_BINARIES.put(pgBinaryResolver, pgDir);
             LOG.info("Postgres binaries at {}", pgDir);
             return pgDir;
         } finally {

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

@@ -21,6 +21,7 @@
 import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.SynchronousQueue;
@@ -41,11 +42,9 @@ public class PreparedDbProvider
      * loaded so that the databases may be cloned.
      */
     // @GuardedBy("PreparedDbProvider.class")
-    private static final Map<DatabasePreparer, PrepPipeline> CLUSTERS = new HashMap<>();
+    private static final Map<ClusterKey, PrepPipeline> CLUSTERS = new HashMap<>();
 
     private final PrepPipeline dbPreparer;
-    private final Iterable<Consumer<Builder>> customizers;
-
 
     public static PreparedDbProvider forPreparer(DatabasePreparer preparer) {
         return forPreparer(preparer, Collections.emptyList());
@@ -55,11 +54,9 @@ public static PreparedDbProvider forPreparer(DatabasePreparer preparer, Iterable
         return new PreparedDbProvider(preparer, customizers);
     }
 
-    private PreparedDbProvider(DatabasePreparer preparer, Iterable<Consumer<Builder>> customizers)
-    {
-        this.customizers = customizers;
+    private PreparedDbProvider(DatabasePreparer preparer, Iterable<Consumer<Builder>> customizers) {
         try {
-            dbPreparer = createOrFindPreparer(preparer);
+            dbPreparer = createOrFindPreparer(preparer, customizers);
         } catch (final IOException | SQLException e) {
             throw new RuntimeException(e);
         }
@@ -69,9 +66,10 @@ 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 synchronized PrepPipeline createOrFindPreparer(DatabasePreparer preparer) throws IOException, SQLException
+    private static synchronized PrepPipeline createOrFindPreparer(DatabasePreparer preparer, Iterable<Consumer<Builder>> customizers) throws IOException, SQLException
     {
-        PrepPipeline result = CLUSTERS.get(preparer);
+        final ClusterKey key = new ClusterKey(preparer, customizers);
+        PrepPipeline result = CLUSTERS.get(key);
         if (result != null) {
             return result;
         }
@@ -82,7 +80,7 @@ private synchronized PrepPipeline createOrFindPreparer(DatabasePreparer preparer
         preparer.prepare(pg.getTemplateDatabase());
 
         result = new PrepPipeline(pg).start();
-        CLUSTERS.put(preparer, result);
+        CLUSTERS.put(key, result);
         return result;
     }
 
@@ -233,6 +231,36 @@ private static void create(final DataSource connectDb, final String dbName, fina
         }
     }
 
+    private static class ClusterKey {
+
+        private final DatabasePreparer preparer;
+        private final Builder builder;
+
+        ClusterKey(DatabasePreparer preparer, Iterable<Consumer<Builder>> customizers) {
+            this.preparer = preparer;
+            this.builder = EmbeddedPostgres.builder();
+            customizers.forEach(c -> c.accept(this.builder));
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            ClusterKey that = (ClusterKey) o;
+            return Objects.equals(preparer, that.preparer) &&
+                    Objects.equals(builder, that.builder);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(preparer, builder);
+        }
+    }
+
     public static class DbInfo
     {
         public static DbInfo ok(final String dbName, final int port, final String user) {

src/test/java/io/zonky/test/db/postgres/embedded/PreparedDbCustomizerTest.java

@@ -0,0 +1,57 @@
+/*
+ * 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
+ *
+ * http://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 io.zonky.test.db.postgres.embedded;
+
+import io.zonky.test.db.postgres.junit.EmbeddedPostgresRules;
+import io.zonky.test.db.postgres.junit.PreparedDbRule;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.time.Duration;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+public class PreparedDbCustomizerTest {
+
+    private static final DatabasePreparer EMPTY_PREPARER = ds -> {};
+
+    @Rule
+    public PreparedDbRule dbA1 = EmbeddedPostgresRules.preparedDatabase(EMPTY_PREPARER);
+    @Rule
+    public PreparedDbRule dbA2 = EmbeddedPostgresRules.preparedDatabase(EMPTY_PREPARER).customize(builder -> {});
+    @Rule
+    public PreparedDbRule dbA3 = EmbeddedPostgresRules.preparedDatabase(EMPTY_PREPARER).customize(builder -> builder.setPGStartupWait(Duration.ofSeconds(10)));
+    @Rule
+    public PreparedDbRule dbB1 = EmbeddedPostgresRules.preparedDatabase(EMPTY_PREPARER).customize(builder -> builder.setPGStartupWait(Duration.ofSeconds(11)));
+    @Rule
+    public PreparedDbRule dbB2 = EmbeddedPostgresRules.preparedDatabase(EMPTY_PREPARER).customize(builder -> builder.setPGStartupWait(Duration.ofSeconds(11)));
+
+    @Test
+    public void testCustomizers() {
+        int dbA1Port = dbA1.getConnectionInfo().getPort();
+        int dbA2Port = dbA2.getConnectionInfo().getPort();
+        int dbA3Port = dbA3.getConnectionInfo().getPort();
+
+        assertEquals(dbA1Port, dbA2Port);
+        assertEquals(dbA1Port, dbA3Port);
+
+        int dbB1Port = dbB1.getConnectionInfo().getPort();
+        int dbB2Port = dbB2.getConnectionInfo().getPort();
+
+        assertEquals(dbB1Port, dbB2Port);
+
+        assertNotEquals(dbA1Port, dbB2Port);
+    }
+}