diff --git a/driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java b/driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java index c4bd72a4775..f10cd09c68b 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/ChangeStreamBatchCursor.java +++ b/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) changeStreamOperation.execute(new SourceAwareReadBinding(source, binding))).getWrapped(); return null; }); - wrapped = ((ChangeStreamBatchCursor) 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(); + } + } } diff --git a/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java b/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java index 6d013df59ba..f8bbacf53ed 100644 --- a/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java +++ b/driver-core/src/main/com/mongodb/internal/operation/SyncOperationHelper.java @@ -97,6 +97,10 @@ interface CommandWriteTransformer { 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 withReadConnectionSource(final ReadBinding binding, final CallableWithSource callable) { ConnectionSource source = binding.getReadConnectionSource(); try { diff --git a/driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorTest.java b/driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorTest.java index 48c3a50e79a..7835af045dc 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/operation/ChangeStreamBatchCursorTest.java +++ b/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 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 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; + }); } }