Make reference count check atomic with release (#876)

JAVA-4490

Co-authored-by: Valentin Kovalenko <[email protected]>

Author: jyemin

driver-core/src/main/com/mongodb/internal/async/client/ClientSessionBinding.java

@@ -23,6 +23,7 @@
 import com.mongodb.connection.ClusterType;
 import com.mongodb.connection.ServerDescription;
 import com.mongodb.internal.async.SingleResultCallback;
+import com.mongodb.internal.binding.AbstractReferenceCounted;
 import com.mongodb.internal.binding.AsyncClusterAwareReadWriteBinding;
 import com.mongodb.internal.binding.AsyncConnectionSource;
 import com.mongodb.internal.binding.AsyncReadWriteBinding;
@@ -36,15 +37,15 @@
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.connection.ClusterType.LOAD_BALANCED;
 
-public class ClientSessionBinding implements AsyncReadWriteBinding {
+public class ClientSessionBinding extends AbstractReferenceCounted implements AsyncReadWriteBinding {
     private final AsyncClusterAwareReadWriteBinding wrapped;
     private final AsyncClientSession session;
     private final boolean ownsSession;
     private final ClientSessionContext sessionContext;
 
     public  ClientSessionBinding(final AsyncClientSession session, final boolean ownsSession,
                                  final AsyncClusterAwareReadWriteBinding wrapped) {
-        this.wrapped = notNull("wrapped", (wrapped));
+        this.wrapped = notNull("wrapped", wrapped).retain();
         this.ownsSession = ownsSession;
         this.session = notNull("session", session);
         this.sessionContext = new AsyncClientSessionContext(session);
@@ -113,14 +114,9 @@ private void getPinnedConnectionSource(final boolean isRead, final SingleResultC
         }
     }
 
-    @Override
-    public int getCount() {
-        return wrapped.getCount();
-    }
-
     @Override
     public AsyncReadWriteBinding retain() {
-        wrapped.retain();
+        super.retain();
         return this;
     }
 
@@ -131,15 +127,15 @@ public void getReadConnectionSource(final int minWireVersion, final ReadPreferen
     }
 
     @Override
-    public void release() {
-        wrapped.release();
-        closeSessionIfCountIsZero();
-    }
-
-    private void closeSessionIfCountIsZero() {
-        if (getCount() == 0 && ownsSession) {
-            session.close();
+    public int release() {
+        int count = super.release();
+        if (count == 0) {
+            wrapped.release();
+            if (ownsSession) {
+                session.close();
+            }
         }
+        return count;
     }
 
     private boolean isConnectionSourcePinningRequired() {
@@ -152,6 +148,7 @@ private class SessionBindingAsyncConnectionSource implements AsyncConnectionSour
 
         SessionBindingAsyncConnectionSource(final AsyncConnectionSource wrapped) {
             this.wrapped = wrapped;
+            ClientSessionBinding.this.retain();
         }
 
         @Override
@@ -214,9 +211,12 @@ public int getCount() {
         }
 
         @Override
-        public void release() {
-            wrapped.release();
-            closeSessionIfCountIsZero();
+        public int release() {
+            int count = wrapped.release();
+            if (count == 0) {
+                ClientSessionBinding.this.release();
+            }
+            return count;
         }
     }
 

driver-core/src/main/com/mongodb/internal/binding/AbstractReferenceCounted.java

@@ -34,9 +34,11 @@ public ReferenceCounted retain() {
     }
 
     @Override
-    public void release() {
-        if (referenceCount.decrementAndGet() < 0) {
+    public int release() {
+        int decrementedValue = referenceCount.decrementAndGet();
+        if (decrementedValue < 0) {
             throw new IllegalStateException("Attempted to decrement the reference count below 0");
         }
+        return decrementedValue;
     }
 }

driver-core/src/main/com/mongodb/internal/binding/AsyncClusterAwareReadWriteBinding.java

@@ -33,4 +33,7 @@ public interface AsyncClusterAwareReadWriteBinding extends AsyncReadWriteBinding
      * @param callback the to be passed the connection source
      */
     void getConnectionSource(ServerAddress serverAddress, SingleResultCallback<AsyncConnectionSource> callback);
+
+    @Override
+    AsyncClusterAwareReadWriteBinding retain();
 }

driver-core/src/main/com/mongodb/internal/binding/AsyncClusterBinding.java

@@ -73,7 +73,7 @@ public AsyncClusterBinding(final Cluster cluster, final ReadPreference readPrefe
     }
 
     @Override
-    public AsyncReadWriteBinding retain() {
+    public AsyncClusterAwareReadWriteBinding retain() {
         super.retain();
         return this;
     }
@@ -208,9 +208,10 @@ public AsyncConnectionSource retain() {
         }
 
         @Override
-        public void release() {
-            super.release();
+        public int release() {
+            int count = super.release();
             AsyncClusterBinding.this.release();
+            return count;
         }
     }
 }

driver-core/src/main/com/mongodb/internal/binding/ClusterBinding.java

@@ -182,9 +182,10 @@ public ConnectionSource retain() {
         }
 
         @Override
-        public void release() {
-            super.release();
+        public int release() {
+            int count = super.release();
             ClusterBinding.this.release();
+            return count;
         }
     }
 }

driver-core/src/main/com/mongodb/internal/binding/ReferenceCounted.java

@@ -16,6 +16,10 @@
 
 package com.mongodb.internal.binding;
 
+import com.mongodb.internal.VisibleForTesting;
+
+import static com.mongodb.internal.VisibleForTesting.AccessModifier.PRIVATE;
+
 /**
  * An interface for reference-counted objects.
  * <p>
@@ -39,9 +43,14 @@ public interface ReferenceCounted {
     /**
      * Gets the current reference count.
      *
+     * <p>
+     * This method should only be used for testing.  Production code should prefer using the count returned from {@link #release()}
+     * </p>
+     *
      * @return the current count, which must be greater than or equal to 0.
      * Returns 1 for a newly created object.
      */
+    @VisibleForTesting(otherwise = PRIVATE)
     int getCount();
 
     /**
@@ -54,6 +63,7 @@ public interface ReferenceCounted {
     /**
      * Release a reference to this object.
      * @throws java.lang.IllegalStateException if the reference count is already 0
+     * @return the reference count after the release
      */
-    void release();
+    int release();
 }

driver-core/src/main/com/mongodb/internal/binding/SingleServerBinding.java

@@ -146,11 +146,12 @@ public ConnectionSource retain() {
         }
 
         @Override
-        public void release() {
-            super.release();
-            if (super.getCount() == 0) {
+        public int release() {
+            int count = super.release();
+            if (count == 0) {
                 SingleServerBinding.this.release();
             }
+            return count;
         }
     }
 }

driver-core/src/main/com/mongodb/internal/binding/TransactionContext.java

@@ -48,13 +48,14 @@ public boolean isConnectionPinningRequired() {
     }
 
     @Override
-    public void release() {
-        super.release();
-        if (getCount() == 0) {
+    public int release() {
+        int count = super.release();
+        if (count == 0) {
             if (pinnedConnection != null) {
                 pinnedConnection.release();
             }
         }
+        return count;
     }
 
     @SuppressWarnings("unchecked")

driver-core/src/main/com/mongodb/internal/connection/AbstractReferenceCounted.java

@@ -36,9 +36,11 @@ public ReferenceCounted retain() {
     }
 
     @Override
-    public void release() {
-        if (referenceCount.decrementAndGet() < 0) {
+    public int release() {
+        int decrementedValue = referenceCount.decrementAndGet();
+        if (decrementedValue < 0) {
             throw new IllegalStateException("Attempted to decrement the reference count below 0");
         }
+        return decrementedValue;
     }
 }

driver-core/src/main/com/mongodb/internal/connection/DefaultServer.java

@@ -301,11 +301,12 @@ public int getCount() {
         }
 
         @Override
-        public void release() {
-            wrapped.release();
-            if (getCount() == 0) {
+        public int release() {
+            int count = wrapped.release();
+            if (count == 0) {
                 server.operationEnd();
             }
+            return count;
         }
 
         @Override
@@ -401,11 +402,12 @@ public int getCount() {
         }
 
         @Override
-        public void release() {
-            wrapped.release();
-            if (getCount() == 0) {
+        public int release() {
+            int count = wrapped.release();
+            if (count == 0) {
                 server.operationEnd();
             }
+            return count;
         }
 
         @Override