Add BatchCursor interceptor.

Author: vbabanin

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

@@ -21,16 +21,24 @@
 import com.mongodb.ServerCursor;
 import com.mongodb.client.MongoCursor;
 import com.mongodb.lang.Nullable;
+import com.mongodb.reactivestreams.client.internal.BatchCursor;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Operators;
+import reactor.util.context.Context;
 
 import java.util.NoSuchElementException;
 import java.util.concurrent.BlockingDeque;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import static com.mongodb.ClusterFixture.TIMEOUT;
 import static com.mongodb.internal.thread.InterruptionUtil.interruptAndCreateMongoInterruptedException;
@@ -41,6 +49,7 @@
 class SyncMongoCursor<T> implements MongoCursor<T> {
     private static final Object COMPLETED = new Object();
     private final BlockingDeque<Object> results = new LinkedBlockingDeque<>();
+    private final CompletableFuture<BatchCursor> batchCursorCompletableFuture = new CompletableFuture<>();
     private final Integer batchSize;
     private int countToBatchSize;
     private Subscription subscription;
@@ -50,7 +59,11 @@ class SyncMongoCursor<T> implements MongoCursor<T> {
 
     SyncMongoCursor(final Publisher<T> publisher, @Nullable final Integer batchSize) {
         this.batchSize = batchSize;
-        CountDownLatch latch = new CountDownLatch(1);
+        CountDownLatch subscriptionLatch = new CountDownLatch(1);
+        Hooks.onEachOperator(Operators.lift((sc, sub) ->
+                new BatchCursorInterceptSubscriber(sub, batchCursorCompletableFuture
+                )));
+
         //noinspection ReactiveStreamsSubscriberImplementation
         Flux.from(publisher).contextWrite(CONTEXT).subscribe(new Subscriber<T>() {
             @Override
@@ -61,7 +74,7 @@ public void onSubscribe(final Subscription s) {
                 } else {
                     subscription.request(batchSize);
                 }
-                latch.countDown();
+                subscriptionLatch.countDown();
             }
 
             @Override
@@ -80,12 +93,20 @@ public void onComplete() {
             }
         });
         try {
-            if (!latch.await(TIMEOUT, TimeUnit.SECONDS)) {
+            if (!subscriptionLatch.await(TIMEOUT, TimeUnit.SECONDS)) {
                 throw new MongoTimeoutException("Timeout waiting for subscription");
             }
+            batchCursorCompletableFuture.get(TIMEOUT, TimeUnit.SECONDS);
+            Hooks.resetOnEachOperator();
             sleep(getSleepAfterCursorOpen());
         } catch (InterruptedException e) {
             throw interruptAndCreateMongoInterruptedException("Interrupted waiting for asynchronous cursor establishment", e);
+        } catch (ExecutionException | TimeoutException e) {
+            Throwable cause = e.getCause();
+            if (cause instanceof RuntimeException) {
+                throw (RuntimeException) cause;
+            }
+            throw new RuntimeException(e);
         }
     }
 
@@ -181,4 +202,50 @@ private RuntimeException translateError(final Throwable throwable) {
         }
         return new RuntimeException(throwable);
     }
+
+
+    static class BatchCursorInterceptSubscriber implements CoreSubscriber {
+
+        private final CoreSubscriber sub;
+        private final CompletableFuture<BatchCursor> batchCursorCompletableFuture;
+
+
+        BatchCursorInterceptSubscriber(final CoreSubscriber sub,
+                                       final CompletableFuture<BatchCursor> batchCursorCompletableFuture) {
+            this.sub = sub;
+            this.batchCursorCompletableFuture = batchCursorCompletableFuture;
+        }
+
+        @Override
+        public Context currentContext() {
+            return sub.currentContext();
+        }
+
+        @Override
+        public void onSubscribe(final Subscription s) {
+            sub.onSubscribe(s);
+        }
+
+        @Override
+        public void onNext(final Object o) {
+            if (o instanceof BatchCursor) {
+                // Interception of a cursor means that it has been created at this point.
+                batchCursorCompletableFuture.complete((BatchCursor) o);
+            }
+            sub.onNext(o);
+        }
+
+        @Override
+        public void onError(final Throwable t) {
+            if (!batchCursorCompletableFuture.isDone()) { // Cursor has not been created yet but an error occurred.
+                batchCursorCompletableFuture.completeExceptionally(t);
+            }
+            sub.onError(t);
+        }
+
+        @Override
+        public void onComplete() {
+            sub.onComplete();
+        }
+    }
 }

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

@@ -36,10 +36,7 @@ public final class ChangeStreamsTest extends UnifiedReactiveStreamsTest {
 
     private static final List<String> ERROR_REQUIRED_FROM_CHANGE_STREAM_INITIALIZATION_TESTS =
             Arrays.asList(
-                    "Test with document comment - pre 4.4",
-                    "Change Stream should error when an invalid aggregation stage is passed in",
-                    "The watch helper must not throw a custom exception when executed against a single server topology, "
-                            + "but instead depend on a server error"
+                    "Test with document comment - pre 4.4"
             );
 
     private static final List<String> EVENT_SENSITIVE_TESTS =