Only sleep when necessary in SyncMongoCursor construction or close (#721)

JAVA-4176

Author: jyemin

driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/syncadapter/SyncMongoClient.java

@@ -34,6 +34,55 @@
 import static java.util.Objects.requireNonNull;
 
 public class SyncMongoClient implements MongoClient {
+
+    private static long sleepAfterCursorOpenMS;
+
+    private static long sleepAfterCursorCloseMS;
+
+    /**
+     * Unfortunately this is the only way to wait for a query to be initiated, since Reactive Streams is asynchronous
+     * and we have no way of knowing. Tests which require cursor initiation to complete before execution of the next operation
+     * can set this to a positive value.  A value of 256 ms has been shown to work well. The default value is 0.
+     */
+    public static void enableSleepAfterCursorOpen(final long sleepMS) {
+        if (sleepAfterCursorOpenMS != 0) {
+            throw new IllegalStateException("Already enabled");
+        }
+        if (sleepMS <= 0) {
+            throw new IllegalArgumentException("sleepMS must be a postive value");
+        }
+        sleepAfterCursorOpenMS = sleepMS;
+    }
+
+    /**
+     * Unfortunately this is the only way to wait for close to complete, since it's asynchronous.
+     * This is inherently racy but there are not any other good options. Tests which require cursor cancellation to complete before
+     * execution of the next operation can set this to a positive value.  A value of 256 ms has been shown to work well. The default
+     * value is 0.
+     */
+    public static void enableSleepAfterCursorClose(final long sleepMS) {
+        if (sleepAfterCursorCloseMS != 0) {
+            throw new IllegalStateException("Already enabled");
+        }
+        if (sleepMS <= 0) {
+            throw new IllegalArgumentException("sleepMS must be a postive value");
+        }
+        sleepAfterCursorCloseMS = sleepMS;
+    }
+
+    public static void disableCursorSleep() {
+        sleepAfterCursorOpenMS = 0;
+        sleepAfterCursorCloseMS = 0;
+    }
+
+    public static long getSleepAfterCursorOpen() {
+        return sleepAfterCursorOpenMS;
+    }
+
+    public static long getSleepAfterCursorClose() {
+        return sleepAfterCursorCloseMS;
+    }
+
     private final com.mongodb.reactivestreams.client.MongoClient wrapped;
 
     public SyncMongoClient(final com.mongodb.reactivestreams.client.MongoClient wrapped) {

driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/syncadapter/SyncMongoCursor.java

@@ -33,6 +33,8 @@
 import java.util.concurrent.TimeUnit;
 
 import static com.mongodb.ClusterFixture.TIMEOUT;
+import static com.mongodb.reactivestreams.client.syncadapter.SyncMongoClient.getSleepAfterCursorClose;
+import static com.mongodb.reactivestreams.client.syncadapter.SyncMongoClient.getSleepAfterCursorOpen;
 
 class SyncMongoCursor<T> implements MongoCursor<T> {
     private static final Object COMPLETED = new Object();
@@ -79,9 +81,7 @@ public void onComplete() {
             if (!latch.await(TIMEOUT, TimeUnit.SECONDS)) {
                 throw new MongoTimeoutException("Timeout waiting for subscription");
             }
-            // Unfortunately this is the only way to wait for the query to be initiated, since its asynchronous
-            // and we have no way of knowing
-            Thread.sleep(250);
+            sleep(getSleepAfterCursorOpen());
         } catch (InterruptedException e) {
             throw new MongoInterruptedException("Interrupted waiting for asynchronous cursor establishment", e);
         }
@@ -90,10 +90,12 @@ public void onComplete() {
     @Override
     public void close() {
         subscription.cancel();
-        // Unfortunately this is the only way to wait for cancellation to complete, since it's asynchronous.
-        // This is inherently racy but there are not any other good options.
+        sleep(getSleepAfterCursorClose());
+    }
+
+    private static void sleep(final long millis) {
         try {
-            Thread.sleep(250);
+            Thread.sleep(millis);
         } catch (InterruptedException e) {
             throw new MongoInterruptedException("Interrupted from nap", e);
         }

driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/unified/LoadBalancerTest.java

@@ -24,16 +24,39 @@
 import com.mongodb.reactivestreams.client.syncadapter.SyncMongoClient;
 import org.bson.BsonArray;
 import org.bson.BsonDocument;
+import org.junit.After;
 import org.junit.runners.Parameterized;
 
 import java.io.IOException;
 import java.net.URISyntaxException;
+import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 
+import static com.mongodb.reactivestreams.client.syncadapter.SyncMongoClient.disableCursorSleep;
+import static com.mongodb.reactivestreams.client.syncadapter.SyncMongoClient.enableSleepAfterCursorClose;
+import static com.mongodb.reactivestreams.client.syncadapter.SyncMongoClient.enableSleepAfterCursorOpen;
 import static org.junit.Assume.assumeFalse;
 
 public class LoadBalancerTest extends UnifiedTest {
 
+    private static final List<String> CURSOR_OPEN_TIMING_SENSITIVE_TESTS =
+            Arrays.asList(
+                    "pinned connections are returned when the cursor is drained",
+                    "only connections for a specific serviceId are closed when pools are cleared",
+                    "pinned connections are returned to the pool when the cursor is closed",
+                    "no connection is pinned if all documents are returned in the initial batch",
+                    "stale errors are ignored",
+                    "a connection can be shared by a transaction and a cursor",
+                    "wait queue timeout errors include cursor statistics");
+
+    private static final List<String> CURSOR_CLOSE_TIMING_SENSITIVE_TESTS =
+            Arrays.asList(
+                    "pinned connections are returned to the pool when the cursor is closed",
+                    "only connections for a specific serviceId are closed when pools are cleared",
+                    "pinned connections are returned after a network error during a killCursors request",
+                    "a connection can be shared by a transaction and a cursor");
+
     public LoadBalancerTest(@SuppressWarnings("unused") final String fileDescription,
                             final String testDescription,
                             final String schemaVersion, @Nullable final BsonArray runOnRequirements, final BsonArray entities,
@@ -46,6 +69,20 @@ public LoadBalancerTest(@SuppressWarnings("unused") final String fileDescription
         // Reactive streams driver can't implement this test because there is no way to tell that a change stream cursor
         // that has not yet received any results has even initiated the change stream
         assumeFalse(testDescription.equals("change streams pin to a connection"));
+
+        if (CURSOR_OPEN_TIMING_SENSITIVE_TESTS.contains(testDescription)) {
+            enableSleepAfterCursorOpen(256);
+        }
+
+        if (CURSOR_CLOSE_TIMING_SENSITIVE_TESTS.contains(testDescription)) {
+            enableSleepAfterCursorClose(256);
+        }
+    }
+
+    @After
+    public void cleanUp() {
+        super.cleanUp();
+        disableCursorSleep();
     }
 
     @Override