implement automatic releasing of prepared databases in prefetching provider

Author: tomix26

embedded-database-spring-test/src/main/java/io/zonky/test/db/provider/common/PrefetchingDatabaseProvider.java

@@ -26,13 +26,15 @@
 import io.zonky.test.db.provider.EmbeddedDatabase;
 import io.zonky.test.db.provider.ProviderException;
 import org.apache.commons.lang3.RandomStringUtils;
+import org.apache.commons.lang3.tuple.Pair;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.core.env.Environment;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.util.concurrent.ListenableFutureCallback;
 import org.springframework.util.concurrent.ListenableFutureTask;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -64,6 +66,7 @@ public class PrefetchingDatabaseProvider implements DatabaseProvider {
 
     protected static final ThreadPoolTaskExecutor taskExecutor = new PriorityThreadPoolTaskExecutor();
     protected static final ConcurrentMap<PipelineKey, DatabasePipeline> pipelines = new ConcurrentHashMap<>();
+    protected static final AtomicLong databaseCount = new AtomicLong();
 
     protected final int pipelineMaxCacheSize;
 
@@ -92,6 +95,7 @@ public PrefetchingDatabaseProvider(DatabaseProvider provider, Environment enviro
     public EmbeddedDatabase createDatabase(DatabasePreparer preparer) throws ProviderException {
         Stopwatch stopwatch = Stopwatch.createStarted();
         logger.trace("Prefetching pipelines: {}", pipelines.values());
+        databaseCount.decrementAndGet();
 
         PipelineKey key = new PipelineKey(provider, preparer);
         DatabasePipeline pipeline = pipelines.computeIfAbsent(key, k -> new DatabasePipeline());
@@ -108,12 +112,13 @@ public EmbeddedDatabase createDatabase(DatabasePreparer preparer) throws Provide
         }
 
         long invocationCount = pipeline.requests.incrementAndGet();
-        if (invocationCount > 1) {
-            if (invocationCount - 1 <= pipelineMaxCacheSize) {
-                prepareDatabase(key, -1);
-            }
-            reschedulePipeline(key);
+        long databasesCount = pipeline.tasks.size() + pipeline.results.size();
+        if (result == null) databasesCount--;
+
+        if (databasesCount < invocationCount - 1 && databasesCount < pipelineMaxCacheSize) {
+            prepareDatabase(key, -1);
         }
+        reschedulePipeline(key);
 
         if (result == null) {
             try {
@@ -141,6 +146,21 @@ protected PrefetchingTask prepareDatabase(PipelineKey key, int priority) {
     }
 
     protected PrefetchingTask prepareNewDatabase(PipelineKey key, int priority) {
+        databaseCount.incrementAndGet();
+
+        Pair<PipelineKey, EmbeddedDatabase> databaseToRemove = findDatabaseToRemove().orElse(null);
+        if (databaseToRemove != null) {
+            databaseCount.decrementAndGet();
+
+            if (databaseToRemove.getKey().equals(key)) {
+                return executeTask(key, PrefetchingTask.withDatabase(databaseToRemove.getValue(), priority));
+            } else {
+                databaseToRemove.getValue().close();
+                DatabasePipeline pipeline = pipelines.get(databaseToRemove.getKey());
+                logger.trace("Prepared database has been cleaned: {}", pipeline.key);
+            }
+        }
+
         return executeTask(key, PrefetchingTask.forPreparer(key.provider, key.preparer, priority));
     }
 
@@ -221,6 +241,35 @@ public void onFailure(Throwable error) {
         return task;
     }
 
+    protected Optional<Pair<PipelineKey, EmbeddedDatabase>> findDatabaseToRemove() {
+        while (databaseCount.get() > 35) {
+            long timestampThreshold = System.currentTimeMillis() - 10_000;
+
+            PipelineKey key = pipelines.entrySet().stream()
+                    .map(e -> Pair.of(e.getKey(), e.getValue().results.peek()))
+                    .filter(e -> e.getValue() != null && e.getValue().getTimestamp() < timestampThreshold)
+                    .min(Comparator.comparing(e -> e.getValue().getTimestamp()))
+                    .map(Pair::getKey).orElse(null);
+
+            if (key == null) {
+                return Optional.empty();
+            }
+
+            DatabasePipeline pipeline = pipelines.get(key);
+            if (pipeline != null) {
+                PreparedResult result = pipeline.results.poll();
+                if (result != null) {
+                    if (result.hasResult()) {
+                        return Optional.of(Pair.of(key, result.get()));
+                    } else {
+                        databaseCount.decrementAndGet();
+                    }
+                }
+            }
+        }
+        return Optional.empty();
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -288,6 +337,7 @@ protected enum State {
 
     protected static class PreparedResult {
 
+        private final long timestamp = System.currentTimeMillis();
         private final EmbeddedDatabase result;
         private final Throwable error;
 
@@ -304,6 +354,14 @@ protected PreparedResult(EmbeddedDatabase result, Throwable error) {
             this.error = error;
         }
 
+        public long getTimestamp() {
+            return timestamp;
+        }
+
+        public boolean hasResult() {
+            return result != null;
+        }
+
         public EmbeddedDatabase get() throws ProviderException {
             if (result != null) {
                 return result;
@@ -340,6 +398,10 @@ public static PrefetchingTask withDatabase(EmbeddedDatabase database, DatabasePr
             });
         }
 
+        public static PrefetchingTask withDatabase(EmbeddedDatabase database, int priority) {
+            return new PrefetchingTask(priority, EXISTING_DATABASE, () -> database);
+        }
+
         public static PrefetchingTask fromTask(PrefetchingTask task, int priority) {
             return new PrefetchingTask(priority, task.type, task.action);
         }

embedded-database-spring-test/src/main/java/io/zonky/test/db/provider/common/TemplatingDatabaseProvider.java

@@ -358,7 +358,7 @@ public int hashCode() {
         public static class Builder {
 
             private long durationThreshold = 0;
-            private int maxTemplateCount = 50;
+            private int maxTemplateCount = 15;
 
             private Builder() {}