Reuse ConnectionSource to avoid extra server selection.

driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java

@@ -19,10 +19,14 @@
 import com.mongodb.MongoChangeStreamException;
 import com.mongodb.MongoException;
 import com.mongodb.MongoOperationTimeoutException;
+import com.mongodb.ReadPreference;
 import com.mongodb.ServerAddress;
 import com.mongodb.ServerCursor;
+import com.mongodb.assertions.Assertions;
 import com.mongodb.internal.TimeoutContext;
+import com.mongodb.internal.binding.ConnectionSource;
 import com.mongodb.internal.binding.ReadBinding;
+import com.mongodb.internal.connection.OperationContext;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
 import org.bson.BsonTimestamp;
@@ -244,9 +248,9 @@ private void resumeChangeStream() {
 
         withReadConnectionSource(binding, source -> {
             changeStreamOperation.setChangeStreamOptionsForResume(resumeToken, source.getServerDescription().getMaxWireVersion());
+            wrapped = ((ChangeStreamBatchCursor<T>) changeStreamOperation.execute(new SourceAwareReadBinding(source, binding))).getWrapped();
             return null;
         });
-        wrapped = ((ChangeStreamBatchCursor<T>) changeStreamOperation.execute(binding)).getWrapped();
         binding.release(); // release the new change stream batch cursor's reference to the binding
     }
 
@@ -257,4 +261,49 @@ private boolean hasPreviousNextTimedOut() {
     private static boolean isTimeoutException(final Throwable exception) {
         return exception instanceof MongoOperationTimeoutException;
     }
+
+    private static class SourceAwareReadBinding implements ReadBinding {
+        private final ConnectionSource source;
+        private final ReadBinding binding;
+
+        SourceAwareReadBinding(final ConnectionSource source, final ReadBinding binding) {
+            this.source = source;
+            this.binding = binding;
+        }
+
+        @Override
+        public ReadPreference getReadPreference() {
+            return binding.getReadPreference();
+        }
+
+        @Override
+        public ConnectionSource getReadConnectionSource() {
+            return source;
+        }
+
+        @Override
+        public ConnectionSource getReadConnectionSource(final int minWireVersion, final ReadPreference fallbackReadPreference) {
+            throw Assertions.fail();
+        }
+
+        @Override
+        public int getCount() {
+            return binding.getCount();
+        }
+
+        @Override
+        public ReadBinding retain() {
+            return binding.retain();
+        }
+
+        @Override
+        public int release() {
+            return binding.release();
+        }
+
+        @Override
+        public OperationContext getOperationContext() {
+            return binding.getOperationContext();
+        }
+    }
 }

driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java

@@ -97,6 +97,10 @@ interface CommandWriteTransformer<T, R> {
 
     private static final BsonDocumentCodec BSON_DOCUMENT_CODEC = new BsonDocumentCodec();
 
+    /**
+     * Gets a {@link ConnectionSource} from the {@link ReadBinding#getReadConnectionSource()} and executes
+     * the provided {@link  CallableWithSource} with it.
+     */
     static <T> T withReadConnectionSource(final ReadBinding binding, final CallableWithSource<T> callable) {
         ConnectionSource source = binding.getReadConnectionSource();
         try {

driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorTest.java

@@ -153,7 +153,7 @@ void shouldResumeOnlyOnceOnSubsequentCallsAfterTimeoutError() {
         verifyNoMoreInteractions(commandBatchCursor);
         verify(changeStreamOperation).setChangeStreamOptionsForResume(resumeToken, maxWireVersion);
         verify(changeStreamOperation, times(1)).getDecoder();
-        verify(changeStreamOperation, times(1)).execute(readBinding);
+        verify(changeStreamOperation, times(1)).execute(any(ReadBinding.class));
         verifyNoMoreInteractions(changeStreamOperation);
         verify(newCommandBatchCursor, times(1)).next();
         verify(newCommandBatchCursor, atLeastOnce()).getPostBatchResumeToken();
@@ -180,7 +180,7 @@ void shouldResumeOnlyOnceOnSubsequentCallsAfterTimeoutError() {
     void shouldPropagateAnyErrorsOccurredInAggregateOperation() {
         when(commandBatchCursor.next()).thenThrow(new MongoOperationTimeoutException("timeout"));
         MongoNotPrimaryException resumableError = new MongoNotPrimaryException(new BsonDocument(), new ServerAddress());
-        when(changeStreamOperation.execute(readBinding)).thenThrow(resumableError);
+        when(changeStreamOperation.execute(any(ReadBinding.class))).thenThrow(resumableError);
 
         ChangeStreamBatchCursor<Document> cursor = createChangeStreamCursor();
         //when
@@ -208,11 +208,12 @@ void shouldResumeAfterTimeoutInAggregateOnNextCall() {
         clearInvocations(commandBatchCursor, newCommandBatchCursor, timeoutContext, changeStreamOperation, readBinding);
 
         //second next operation times out on resume attempt when creating change stream
-        when(changeStreamOperation.execute(readBinding)).thenThrow(new MongoOperationTimeoutException("timeout during resumption"));
+        when(changeStreamOperation.execute(any(ReadBinding.class))).thenThrow(
+                new MongoOperationTimeoutException("timeout during resumption"));
         assertThrows(MongoOperationTimeoutException.class, cursor::next);
-        clearInvocations(commandBatchCursor, newCommandBatchCursor, timeoutContext, changeStreamOperation);
+        clearInvocations(commandBatchCursor, newCommandBatchCursor, timeoutContext, changeStreamOperation, readBinding);
 
-        doReturn(newChangeStreamCursor).when(changeStreamOperation).execute(readBinding);
+        doReturn(newChangeStreamCursor).when(changeStreamOperation).execute(any(ReadBinding.class));
 
         //when third operation succeeds to resume and call next
         List<Document> next = cursor.next();
@@ -242,7 +243,8 @@ void shouldCloseChangeStreamWhenResumeOperationFailsDueToNonTimeoutError() {
         clearInvocations(commandBatchCursor, newCommandBatchCursor, timeoutContext, changeStreamOperation, readBinding);
 
         //when second next operation errors on resume attempt when creating change stream
-        when(changeStreamOperation.execute(readBinding)).thenThrow(new MongoNotPrimaryException(new BsonDocument(), new ServerAddress()));
+        when(changeStreamOperation.execute(any(ReadBinding.class))).thenThrow(
+                new MongoNotPrimaryException(new BsonDocument(), new ServerAddress()));
         assertThrows(MongoNotPrimaryException.class, cursor::next);
 
         //then
@@ -280,7 +282,8 @@ private void verifyNoResumeAttemptCalled() {
     private void verifyResumeAttemptCalled() {
         verify(commandBatchCursor, times(1)).close();
         verify(changeStreamOperation).setChangeStreamOptionsForResume(resumeToken, maxWireVersion);
-        verify(changeStreamOperation, times(1)).execute(readBinding);
+        verify(changeStreamOperation, times(1)).execute(any(ReadBinding.class));
+        verify(readBinding, times(1)).getReadConnectionSource();
         verifyNoMoreInteractions(commandBatchCursor);
     }
 
@@ -326,7 +329,11 @@ void setUp() {
         changeStreamOperation = mock(ChangeStreamOperation.class);
         when(changeStreamOperation.getDecoder()).thenReturn(new DocumentCodec());
         doNothing().when(changeStreamOperation).setChangeStreamOptionsForResume(resumeToken, maxWireVersion);
-        when(changeStreamOperation.execute(readBinding)).thenReturn(newChangeStreamCursor);
+        when(changeStreamOperation.execute(any(ReadBinding.class))).thenAnswer(invocation -> {
+            ReadBinding binding = invocation.getArgument(0);
+            binding.getReadConnectionSource();
+            return newChangeStreamCursor;
+        });
     }
 
 }