Merge pull request #2 from marat-gainullin/soft-locks

Lock have been softend. Waiting policy has been changed

Author: marat-gainullin

src/main/java/com/github/pgasync/PgConnectionPool.java

@@ -24,23 +24,15 @@
 import com.pgasync.Transaction;
 
 import javax.annotation.concurrent.GuardedBy;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Iterator;
-import java.util.LinkedHashMap;
-import java.util.LinkedList;
-import java.util.Map;
-import java.util.Queue;
+import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
-import java.util.function.BiConsumer;
-import java.util.function.Consumer;
-import java.util.function.Function;
-import java.util.function.Supplier;
+import java.util.function.*;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 
 /**
  * Resource pool for backend connections.
@@ -128,7 +120,7 @@ private void closeNextStatement(Iterator<PooledPgPreparedStatement> statementsSo
                             return null;
                         });
             } else {
-                onComplete.completeAsync(() -> null, futuresExecutor);
+                futuresExecutor.execute(() -> onComplete.complete(null));
             }
         }
 
@@ -147,7 +139,8 @@ CompletableFuture<Void> shutdown() {
 
         @Override
         public CompletableFuture<Void> close() {
-            return release(this);
+            release(this);
+            return CompletableFuture.completedFuture(null);
         }
 
         @Override
@@ -245,173 +238,162 @@ public CompletableFuture<Integer> fetch(BiConsumer<Map<String, PgColumn>, PgColu
 
     private final int maxConnections;
     private final int maxStatements;
-    private final Lock lock = new ReentrantLock();
 
-    @GuardedBy("lock")
+    private final Lock guard = new ReentrantLock();
+    @GuardedBy("guard")
     private int size;
-    @GuardedBy("lock")
-    private boolean closed;
-    @GuardedBy("lock")
-    private final Queue<CompletableFuture<? super Connection>> uponAvailableSubscribers = new LinkedList<>();
-    @GuardedBy("lock")
-    private final Queue<PooledPgConnection> availableConnections = new LinkedList<>();
-    @GuardedBy("lock")
-    private CompletableFuture<Void> uponFullyAvailable;
+    @GuardedBy("guard")
+    private final Queue<CompletableFuture<? super Connection>> pending = new LinkedList<>();
+    @GuardedBy("guard")
+    private final Queue<PooledPgConnection> connections = new LinkedList<>();
+    @GuardedBy("guard")
+    private CompletableFuture<Void> closing;
 
     public PgConnectionPool(ConnectibleBuilder.ConnectibleProperties properties, Supplier<CompletableFuture<ProtocolStream>> obtainStream, Executor futuresExecutor) {
         super(properties, obtainStream, futuresExecutor);
         this.maxConnections = properties.getMaxConnections();
         this.maxStatements = properties.getMaxStatements();
     }
 
-    private CompletableFuture<Void> fullyAvailable() {
-        if (uponFullyAvailable == null) {
-            if (size <= availableConnections.size()) {
-                return CompletableFuture.completedFuture(null);
-            } else {
-                uponFullyAvailable = new CompletableFuture<>();
-                return uponFullyAvailable;
-            }
-        } else {
-            return CompletableFuture.failedFuture(new IllegalStateException("Only a single 'fullyAvailable' request at a time is supported"));
+    private <T> T locked(Supplier<T> action) {
+        guard.lock();
+        try {
+            return action.get();
+        } finally {
+            guard.unlock();
         }
     }
 
-    @Override
-    public CompletableFuture<Void> close() {
-        lock.lock();
-        try {
-            closed = true;
-            while (!uponAvailableSubscribers.isEmpty()) {
-                CompletableFuture<? super Connection> queued = uponAvailableSubscribers.poll();
-                futuresExecutor.execute(() -> queued.completeExceptionally(new SqlException("Connection pool is closing")));
-            }
-            return fullyAvailable()
-                    .thenApply(v -> {
-                        lock.lock();
-                        try {
-                            uponFullyAvailable = null;
-                            Collection<CompletableFuture<Void>> shutdownTasks = new ArrayList<>();
-                            while (!availableConnections.isEmpty()) {
-                                PooledPgConnection connection = availableConnections.poll();
-                                shutdownTasks.add(connection.shutdown());
-                                size--;
-                            }
-                            return CompletableFuture.allOf(shutdownTasks.toArray(CompletableFuture<?>[]::new));
-                        } finally {
-                            lock.unlock();
-                        }
-                    })
-                    .thenCompose(Function.identity());
-        } finally {
-            lock.unlock();
+    private void release(PooledPgConnection connection) {
+        if (connection == null) {
+            throw new IllegalArgumentException("'connection' should be not null");
         }
+        Runnable lucky = locked(() -> {
+            CompletableFuture<? super Connection> nextUser = pending.poll();
+            if (nextUser != null) {
+                return () -> nextUser.complete(connection);
+            } else {
+                connections.add(connection);
+                return checkClosed();
+            }
+        });
+        futuresExecutor.execute(lucky);
     }
 
     @Override
     public CompletableFuture<Connection> getConnection() {
-        CompletableFuture<Connection> uponAvailable = new CompletableFuture<>();
-        lock.lock();
-        try {
-            if (closed) {
-                futuresExecutor.execute(() -> uponAvailable.completeExceptionally(new SqlException("Connection pool is closed")));
+        if (locked(() -> closing != null)) {
+            return CompletableFuture.failedFuture(new SqlException("Connection pool is closed"));
+        } else {
+            Connection cached = locked(this::firstAliveConnection);
+            if (cached != null) {
+                return CompletableFuture.completedFuture(cached);
             } else {
-                Connection connection = firstAliveConnection();
-                if (connection != null) {
-                    uponAvailable.completeAsync(() -> connection, futuresExecutor);
-                } else {
-                    if (tryIncreaseSize()) {
-                        obtainStream.get()
-                                .thenApply(stream -> new PooledPgConnection(new PgConnection(stream, dataConverter))
-                                        .connect(username, password, database))
-                                .thenCompose(Function.identity())
-                                .thenApply(pooledConnection -> {
-                                    if (validationQuery != null && !validationQuery.isBlank()) {
-                                        return pooledConnection.completeScript(validationQuery)
-                                                .handle((rss, th) -> {
-                                                    if (th != null) {
-                                                        return ((PooledPgConnection) pooledConnection).delegate.close()
-                                                                .thenApply(v -> CompletableFuture.<Connection>failedFuture(th))
-                                                                .thenCompose(Function.identity());
-                                                    } else {
-                                                        return CompletableFuture.completedFuture(pooledConnection);
-                                                    }
-                                                })
-                                                .thenCompose(Function.identity());
-                                    } else {
-                                        return CompletableFuture.completedFuture(pooledConnection);
-                                    }
-                                })
-                                .thenCompose(Function.identity())
-                                .thenAccept(pooledConnection -> uponAvailable.completeAsync(() -> pooledConnection, futuresExecutor))
-                                .exceptionally(th -> {
-                                    lock.lock();
-                                    try {
-                                        size--;
-                                        futuresExecutor.execute(() -> uponAvailable.completeExceptionally(th));
-                                        if (uponFullyAvailable != null && size <= availableConnections.size()) {
-                                            uponFullyAvailable.completeAsync(null, futuresExecutor);
-                                        }
-                                        return null;
-                                    } finally {
-                                        lock.unlock();
-                                    }
-                                });
+                CompletableFuture<Connection> deferred = new CompletableFuture<>();
+                boolean makeNewConnection = locked(() -> {
+                    pending.add(deferred);
+                    if (size < maxConnections) {
+                        size++;
+                        return true;
                     } else {
-                        // Pool is full now and all connections are busy
-                        uponAvailableSubscribers.offer(uponAvailable);
+                        return false;
                     }
+                });
+                if (makeNewConnection) {
+                    obtainStream.get()
+                            .thenApply(stream -> new PooledPgConnection(new PgConnection(stream, dataConverter))
+                                    .connect(username, password, database))
+                            .thenCompose(Function.identity())
+                            .thenApply(pooledConnection -> {
+                                if (validationQuery != null && !validationQuery.isBlank()) {
+                                    return pooledConnection.completeScript(validationQuery)
+                                            .handle((rss, th) -> {
+                                                if (th != null) {
+                                                    return ((PooledPgConnection) pooledConnection).delegate.close()
+                                                            .thenApply(v -> CompletableFuture.<Connection>failedFuture(th))
+                                                            .thenCompose(Function.identity());
+                                                } else {
+                                                    return CompletableFuture.completedFuture(pooledConnection);
+                                                }
+                                            })
+                                            .thenCompose(Function.identity());
+                                } else {
+                                    return CompletableFuture.completedFuture(pooledConnection);
+                                }
+                            })
+                            .thenCompose(Function.identity())
+                            .whenComplete((connected, th) -> {
+                                if (th == null) {
+                                    release((PooledPgConnection) connected);
+                                } else {
+                                    Collection<Runnable> actions = locked(() -> {
+                                        size--;
+                                        List<Runnable> unlucky = pending.stream()
+                                                .<Runnable>map(item -> () ->
+                                                        item.completeExceptionally(th))
+                                                .collect(Collectors.toList());
+                                        unlucky.add(checkClosed());
+                                        pending.clear();
+                                        return unlucky;
+                                    });
+                                    actions.forEach(futuresExecutor::execute);
+                                }
+                            });
                 }
+                return deferred;
             }
-        } finally {
-            lock.unlock();
         }
-
-        return uponAvailable;
     }
 
-    private Connection firstAliveConnection() {
-        Connection connection = availableConnections.poll();
-        while (connection != null && !connection.isConnected()) {
-            size--;
-            connection = availableConnections.poll();
+    private static class CloseTuple {
+        private final CompletableFuture<Void> closing;
+        private final Runnable immediate;
+
+        public CloseTuple(CompletableFuture<Void> closing, Runnable immediate) {
+            this.closing = closing;
+            this.immediate = immediate;
         }
-        return connection;
     }
 
-    private boolean tryIncreaseSize() {
-        if (size < maxConnections) {
-            size++;
-            return true;
+    @Override
+    public CompletableFuture<Void> close() {
+        CloseTuple tuple = locked(() -> {
+            if (closing == null) {
+                closing = new CompletableFuture<>()
+                        .thenApply(v -> locked(() ->
+                                CompletableFuture.allOf(connections.stream()
+                                        .map(PooledPgConnection::shutdown)
+                                        .toArray(CompletableFuture[]::new)
+                                )
+                        ))
+                        .thenCompose(Function.identity());
+                return new CloseTuple(closing, checkClosed());
+            } else {
+                return new CloseTuple(CompletableFuture.failedFuture(new IllegalStateException("PG pool is already shutting down")), NO_OP);
+            }
+        });
+        futuresExecutor.execute(tuple.immediate);
+        return tuple.closing;
+    }
+
+    private static final Runnable NO_OP = () -> {
+    };
+
+    private Runnable checkClosed() {
+        if (closing != null && size <= connections.size()) {
+            assert pending.isEmpty();
+            return () -> closing.complete(null);
         } else {
-            return false;
+            return NO_OP;
         }
     }
 
-    private CompletableFuture<Void> release(PooledPgConnection connection) {
-        if (connection == null) {
-            throw new IllegalArgumentException("'connection' should be not null");
-        }
-        lock.lock();
-        try {
-            if (connection.isConnected()) {
-                if (!uponAvailableSubscribers.isEmpty()) {
-                    uponAvailableSubscribers.poll().completeAsync(() -> connection, futuresExecutor);
-                } else {
-                    availableConnections.offer(connection);
-                    if (uponFullyAvailable != null && size <= availableConnections.size()) {
-                        uponFullyAvailable.complete(null);
-                    }
-                }
-            } else {
-                size--;
-                if (uponFullyAvailable != null && size <= availableConnections.size()) {
-                    uponFullyAvailable.complete(null);
-                }
-            }
-        } finally {
-            lock.unlock();
+    private Connection firstAliveConnection() {
+        Connection connection = connections.poll();
+        while (connection != null && !connection.isConnected()) {
+            size--;
+            connection = connections.poll();
         }
-        return CompletableFuture.completedFuture(null);
+        return connection;
     }
 }

src/test/java/com/github/pgasync/PerformanceTest.java

@@ -117,7 +117,7 @@ private void performBatches(SortedMap<Integer, SortedMap<Integer, Long>> results
 
     private class Batch {
 
-        private long batchSize;
+        private final long batchSize;
         private long performed;
         private long startedAt;
         private CompletableFuture<Long> onBatch;