Improve channel tracking. (neo4j#741)

This uses a reentrant read/write lock to avoid synchronizing and a mixture of concurrent hashmaps and atomic integers.

Co-authored-by: Michael Simons <[email protected]>

Author: 2 people

driver/src/main/java/org/neo4j/driver/internal/async/pool/NettyChannelTracker.java

@@ -25,9 +25,11 @@
 import io.netty.channel.pool.ChannelPoolHandler;
 import io.netty.util.concurrent.EventExecutor;
 
+import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Supplier;
 
 import org.neo4j.driver.internal.BoltServerAddress;
 import org.neo4j.driver.internal.messaging.BoltProtocol;
@@ -40,8 +42,11 @@
 
 public class NettyChannelTracker implements ChannelPoolHandler
 {
-    private final Map<BoltServerAddress,AtomicInteger> addressToInUseChannelCount = new ConcurrentHashMap<>();
-    private final Map<BoltServerAddress,AtomicInteger> addressToIdleChannelCount = new ConcurrentHashMap<>();
+    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+    private final Lock read = lock.readLock();
+    private final Lock write = lock.writeLock();
+    private final Map<BoltServerAddress,Integer> addressToInUseChannelCount = new HashMap<>();
+    private final Map<BoltServerAddress,Integer> addressToIdleChannelCount = new HashMap<>();
     private final Logger log;
     private final MetricsListener metricsListener;
     private final ChannelFutureListener closeListener = future -> channelClosed( future.channel() );
@@ -62,21 +67,28 @@ public NettyChannelTracker( MetricsListener metricsListener, ChannelGroup channe
     @Override
     public void channelReleased( Channel channel )
     {
-        log.debug( "Channel [0x%s] released back to the pool", channel.id() );
-        decrementInUse( channel );
-        incrementIdle( channel );
+        doInWriteLock( () ->
+                       {
+                           decrementInUse( channel );
+                           incrementIdle( channel );
+                       } );
+
         channel.closeFuture().addListener( closeListener );
+        log.debug( "Channel [0x%s] released back to the pool", channel.id() );
     }
 
     @Override
     public void channelAcquired( Channel channel )
     {
-        log.debug( "Channel [%s] acquired from the pool. Local address: %s, remote address: %s",
-                channel.id(), channel.localAddress(), channel.remoteAddress() );
+        doInWriteLock( () ->
+                       {
+                           incrementInUse( channel );
+                           decrementIdle( channel );
+                       } );
 
-        incrementInUse( channel );
-        decrementIdle( channel );
         channel.closeFuture().removeListener( closeListener );
+        log.debug( "Channel [0x%s] acquired from the pool. Local address: %s, remote address: %s", channel.id(), channel.localAddress(),
+                   channel.remoteAddress() );
     }
 
     @Override
@@ -87,13 +99,12 @@ public void channelCreated( Channel channel )
 
     public void channelCreated( Channel channel, ListenerEvent creatingEvent )
     {
-        log.debug( "Channel [%s] created. Local address: %s, remote address: %s",
-                channel.id(), channel.localAddress(), channel.remoteAddress() );
+        doInWriteLock( () -> incrementInUse( channel ) );
 
-        incrementInUse( channel );
         metricsListener.afterCreated( serverAddress( channel ), creatingEvent );
 
         allChannels.add( channel );
+        log.debug( "Channel [0x%s] created. Local address: %s, remote address: %s", channel.id(), channel.localAddress(), channel.remoteAddress() );
     }
 
     public ListenerEvent channelCreating( BoltServerAddress address )
@@ -110,20 +121,18 @@ public void channelFailedToCreate( BoltServerAddress address )
 
     public void channelClosed( Channel channel )
     {
-        decrementIdle( channel );
+        doInWriteLock( () -> decrementIdle( channel ) );
         metricsListener.afterClosed( serverAddress( channel ) );
     }
 
     public int inUseChannelCount( BoltServerAddress address )
     {
-        AtomicInteger count = addressToInUseChannelCount.get( address );
-        return count == null ? 0 : count.get();
+        return retrieveInReadLock( () -> addressToInUseChannelCount.getOrDefault( address, 0 ) );
     }
 
     public int idleChannelCount( BoltServerAddress address )
     {
-        AtomicInteger count = addressToIdleChannelCount.get( address );
-        return count == null ? 0 : count.get();
+        return retrieveInReadLock( () -> addressToIdleChannelCount.getOrDefault( address, 0 ) );
     }
 
     public void prepareToCloseChannels()
@@ -139,7 +148,8 @@ public void prepareToCloseChannels()
             {
                 // only logging it
                 log.debug( "Failed to prepare to close Channel %s due to error %s. " +
-                        "It is safe to ignore this error as the channel will be closed despite if it is successfully prepared to close or not.", channel, e.getMessage() );
+                           "It is safe to ignore this error as the channel will be closed despite if it is successfully prepared to close or not.", channel,
+                           e.getMessage() );
             }
         }
     }
@@ -164,22 +174,47 @@ private void decrementIdle( Channel channel )
         decrement( channel, addressToIdleChannelCount );
     }
 
-    private void increment( Channel channel, Map<BoltServerAddress,AtomicInteger> countMap )
+    private void increment( Channel channel, Map<BoltServerAddress,Integer> countMap )
     {
         BoltServerAddress address = serverAddress( channel );
-        AtomicInteger count = countMap.computeIfAbsent( address, k -> new AtomicInteger() );
-        count.incrementAndGet();
+        Integer count = countMap.computeIfAbsent( address, k -> 0 );
+        countMap.put( address, count + 1 );
     }
 
-    private void decrement( Channel channel, Map<BoltServerAddress,AtomicInteger> countMap )
+    private void decrement( Channel channel, Map<BoltServerAddress,Integer> countMap )
     {
         BoltServerAddress address = serverAddress( channel );
-        AtomicInteger count = countMap.get( address );
-        if ( count == null )
+        if ( !countMap.containsKey( address ) )
         {
-            System.out.println( countMap );
             throw new IllegalStateException( "No count exist for address '" + address + "'" );
         }
-        count.decrementAndGet();
+        Integer count = countMap.get( address );
+        countMap.put( address, count - 1 );
+    }
+
+    private void doInWriteLock( Runnable work )
+    {
+        try
+        {
+            write.lock();
+            work.run();
+        }
+        finally
+        {
+            write.unlock();
+        }
+    }
+
+    private <T> T retrieveInReadLock( Supplier<T> work )
+    {
+        try
+        {
+            read.lock();
+            return work.get();
+        }
+        finally
+        {
+            read.unlock();
+        }
     }
 }