Code review updates

Author: rozza

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/BatchCursorPublisher.java

@@ -106,7 +106,6 @@ public Publisher<T> first() {
                 .flatMap(batchCursor -> Mono.create(sink -> {
                     batchCursor.setBatchSize(1);
                     Mono.from(batchCursor.next())
-                            .contextWrite(sink.contextView())
                             .doOnTerminate(batchCursor::close)
                             .doOnError(sink::error)
                             .doOnSuccess(results -> {

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/Crypt.java

@@ -305,6 +305,7 @@ private void collInfo(final MongoCryptContext cryptContext,
             sink.error(new IllegalStateException("Missing database name"));
         } else {
             collectionInfoRetriever.filter(databaseName, cryptContext.getMongoOperation())
+                    .contextWrite(sink.contextView())
                     .doOnSuccess(result -> {
                         if (result != null) {
                             cryptContext.addMongoOperationResult(result);
@@ -326,6 +327,7 @@ private void mark(final MongoCryptContext cryptContext,
             sink.error(wrapInClientException(new IllegalStateException("Missing database name")));
         } else {
             commandMarker.mark(databaseName, cryptContext.getMongoOperation())
+                    .contextWrite(sink.contextView())
                     .doOnSuccess(result -> {
                         cryptContext.addMongoOperationResult(result);
                         cryptContext.completeMongoOperation();
@@ -340,6 +342,7 @@ private void fetchKeys(final MongoCryptContext cryptContext,
                            @Nullable final String databaseName,
                            final MonoSink<RawBsonDocument> sink) {
         keyRetriever.find(cryptContext.getMongoOperation())
+                .contextWrite(sink.contextView())
                 .doOnSuccess(results -> {
                     for (BsonDocument result : results) {
                         cryptContext.addMongoOperationResult(result);
@@ -357,11 +360,13 @@ private void decryptKeys(final MongoCryptContext cryptContext,
         MongoKeyDecryptor keyDecryptor = cryptContext.nextKeyDecryptor();
         if (keyDecryptor != null) {
             keyManagementService.decryptKey(keyDecryptor)
+                    .contextWrite(sink.contextView())
                     .doOnSuccess(r -> decryptKeys(cryptContext, databaseName, sink))
                     .doOnError(e -> sink.error(wrapInClientException(e)))
                     .subscribe();
         } else {
             Mono.fromRunnable(cryptContext::completeKeyDecryptors)
+                    .contextWrite(sink.contextView())
                     .doOnSuccess(r -> executeStateMachineWithSink(cryptContext, databaseName, sink))
                     .doOnError(e -> sink.error(wrapInClientException(e)))
                     .subscribe();

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/gridfs/GridFSUploadPublisherImpl.java

@@ -100,7 +100,7 @@ public BsonValue getId() {
     public void subscribe(final Subscriber<? super Void> s) {
         Mono.<Void>create(sink -> {
             AtomicBoolean terminated = new AtomicBoolean(false);
-            sink.onCancel(() -> createCancellationMono(terminated).subscribe());
+            sink.onCancel(() -> createCancellationMono(terminated).contextWrite(sink.contextView()).subscribe());
 
             Consumer<Throwable> errorHandler = e -> createCancellationMono(terminated)
                     .contextWrite(sink.contextView())
@@ -163,7 +163,6 @@ private Mono<Void> createCheckAndCreateIndexesMono() {
         AtomicBoolean collectionExists = new AtomicBoolean(false);
 
         return Mono.create(sink -> Mono.from(findPublisher.projection(PROJECTION).first())
-                .contextWrite(sink.contextView())
                 .subscribe(
                         d -> collectionExists.set(true),
                         sink::error,