CSHARP-5887: Simplify retryable read and writes

src/MongoDB.Driver/Core/Operations/AggregateOperation.cs

@@ -29,7 +29,7 @@
 
 namespace MongoDB.Driver.Core.Operations
 {
-    internal sealed class AggregateOperation<TResult> : IReadOperation<IAsyncCursor<TResult>>, IExecutableInRetryableReadContext<IAsyncCursor<TResult>>
+    internal sealed class AggregateOperation<TResult> : IReadOperation<IAsyncCursor<TResult>>, IExecutableInRetryableReadContext<IAsyncCursor<TResult>>, ICommandCreator
     {
         // fields
         private bool? _allowDiskUse;
@@ -279,7 +279,7 @@ public IAsyncCursor<TResult> Execute(OperationContext operationContext, IReadBin
             Ensure.IsNotNull(binding, nameof(binding));
 
             using (BeginOperation())
-            using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
+            using (var context = new RetryableReadContext(binding, _retryRequested))
             {
                 return Execute(operationContext, context);
             }
@@ -308,7 +308,7 @@ public async Task<IAsyncCursor<TResult>> ExecuteAsync(OperationContext operation
             Ensure.IsNotNull(binding, nameof(binding));
 
             using (BeginOperation())
-            using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, _retryRequested).ConfigureAwait(false))
+            using (var context = new RetryableReadContext(binding, _retryRequested))
             {
                 return await ExecuteAsync(operationContext, context).ConfigureAwait(false);
             }
@@ -331,7 +331,7 @@ public async Task<IAsyncCursor<TResult>> ExecuteAsync(OperationContext operation
             }
         }
 
-        internal BsonDocument CreateCommand(OperationContext operationContext, ICoreSession session, ConnectionDescription connectionDescription)
+        public BsonDocument CreateCommand(OperationContext operationContext, ICoreSession session, ConnectionDescription connectionDescription)
         {
             var readConcern = ReadConcernHelper.GetReadConcernForCommand(session, connectionDescription, _readConcern);
             var command = new BsonDocument
@@ -362,9 +362,13 @@ internal BsonDocument CreateCommand(OperationContext operationContext, ICoreSess
         private ReadCommandOperation<AggregateResult> CreateOperation(OperationContext operationContext, RetryableReadContext context)
         {
             var databaseNamespace = _collectionNamespace == null ? _databaseNamespace : _collectionNamespace.DatabaseNamespace;
-            var command = CreateCommand(operationContext, context.Binding.Session, context.Channel.ConnectionDescription);
             var serializer = new AggregateResultDeserializer(_resultSerializer);
-            return new ReadCommandOperation<AggregateResult>(databaseNamespace, command, serializer, MessageEncoderSettings, OperationName)
+            return new ReadCommandOperation<AggregateResult>(
+                databaseNamespace,
+                this,
+                serializer,
+                MessageEncoderSettings,
+                OperationName)
             {
                 RetryRequested = _retryRequested // might be overridden by retryable read context
             };

src/MongoDB.Driver/Core/Operations/BulkMixedWriteOperation.cs

@@ -145,7 +145,7 @@ public WriteConcern WriteConcern
         public BulkWriteOperationResult Execute(OperationContext operationContext, IWriteBinding binding)
         {
             using (BeginOperation())
-            using (var context = RetryableWriteContext.Create(operationContext, binding, IsOperationRetryable()))
+            using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
             {
                 EnsureHintIsSupportedIfAnyRequestHasHint();
                 var helper = new BatchHelper(_requests, _isOrdered, _writeConcern);
@@ -160,7 +160,7 @@ public BulkWriteOperationResult Execute(OperationContext operationContext, IWrit
         public async Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operationContext, IWriteBinding binding)
         {
             using (BeginOperation())
-            using (var context = await RetryableWriteContext.CreateAsync(operationContext, binding, IsOperationRetryable()).ConfigureAwait(false))
+            using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
             {
                 EnsureHintIsSupportedIfAnyRequestHasHint();
                 var helper = new BatchHelper(_requests, _isOrdered, _writeConcern);

src/MongoDB.Driver/Core/Operations/BulkUnmixedWriteOperationBase.cs

@@ -131,7 +131,7 @@ public BulkWriteOperationResult Execute(OperationContext operationContext, Retry
         public BulkWriteOperationResult Execute(OperationContext operationContext, IWriteBinding binding)
         {
             using (BeginOperation())
-            using (var context = RetryableWriteContext.Create(operationContext, binding, IsOperationRetryable()))
+            using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
             {
                 return Execute(operationContext, context);
             }
@@ -147,7 +147,7 @@ public Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operationCon
         public async Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operationContext, IWriteBinding binding)
         {
             using (BeginOperation())
-            using (var context = await RetryableWriteContext.CreateAsync(operationContext, binding, IsOperationRetryable()).ConfigureAwait(false))
+            using (var context = new RetryableWriteContext(binding, IsOperationRetryable()))
             {
                 return await ExecuteAsync(operationContext, context).ConfigureAwait(false);
             }
@@ -286,7 +286,7 @@ public BulkWriteOperationResult CreateFinalResultOrThrow(IChannelHandle channel)
             {
                 var combiner = new BulkWriteBatchResultCombiner(_batchResults, _writeConcern.IsAcknowledged);
                 var remainingRequests = _requests.GetUnprocessedItems();
-                return combiner.CreateResultOrThrowIfHasErrors(channel.ConnectionDescription.ConnectionId, remainingRequests);
+                return combiner.CreateResultOrThrowIfHasErrors(channel?.ConnectionDescription.ConnectionId, remainingRequests);
             }
 
             // private methods

src/MongoDB.Driver/Core/Operations/ChangeStreamOperation.cs

@@ -266,7 +266,7 @@ public IChangeStreamCursor<TResult> Execute(OperationContext operationContext, I
             IAsyncCursor<RawBsonDocument> cursor;
             ICursorBatchInfo cursorBatchInfo;
             BsonTimestamp initialOperationTime;
-            using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
+            using (var context = new RetryableReadContext(binding, _retryRequested))
             {
                 cursor = ExecuteAggregateOperation(operationContext, context);
                 cursorBatchInfo = (ICursorBatchInfo)cursor;
@@ -301,7 +301,7 @@ public async Task<IChangeStreamCursor<TResult>> ExecuteAsync(OperationContext op
             IAsyncCursor<RawBsonDocument> cursor;
             ICursorBatchInfo cursorBatchInfo;
             BsonTimestamp initialOperationTime;
-            using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, _retryRequested).ConfigureAwait(false))
+            using (var context = new RetryableReadContext(binding, _retryRequested))
             {
                 cursor = await ExecuteAggregateOperationAsync(operationContext, context).ConfigureAwait(false);
                 cursorBatchInfo = (ICursorBatchInfo)cursor;
@@ -326,7 +326,7 @@ public async Task<IChangeStreamCursor<TResult>> ExecuteAsync(OperationContext op
         /// <inheritdoc />
         public IAsyncCursor<RawBsonDocument> Resume(OperationContext operationContext, IReadBinding binding)
         {
-            using (var context = RetryableReadContext.Create(operationContext, binding, retryRequested: false))
+            using (var context = new RetryableReadContext(binding, retryRequested: false))
             {
                 return ExecuteAggregateOperation(operationContext, context);
             }
@@ -335,7 +335,7 @@ public IAsyncCursor<RawBsonDocument> Resume(OperationContext operationContext, I
         /// <inheritdoc />
         public async Task<IAsyncCursor<RawBsonDocument>> ResumeAsync(OperationContext operationContext, IReadBinding binding)
         {
-            using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, retryRequested: false).ConfigureAwait(false))
+            using (var context = new RetryableReadContext(binding, retryRequested: false))
             {
                 return await ExecuteAggregateOperationAsync(operationContext, context).ConfigureAwait(false);
             }

src/MongoDB.Driver/Core/Operations/ClientBulkWriteOperation.cs

@@ -96,7 +96,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
             var bulkWriteResults = new BulkWriteRawResult();
             while (true)
             {
-                using var context = RetryableWriteContext.Create(operationContext, binding, GetEffectiveRetryRequested());
+                using var context = new RetryableWriteContext(binding, GetEffectiveRetryRequested());
                 BsonDocument serverResponse = null;
                 try
                 {
@@ -112,7 +112,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
                     bulkWriteResults.TopLevelException = commandException;
                     serverResponse = commandException.Result;
                 }
-                catch (Exception exception)
+                catch (Exception exception) when (context.Channel is not null)
                 {
                     bulkWriteResults.TopLevelException = exception;
                 }
@@ -154,7 +154,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
             var bulkWriteResults = new BulkWriteRawResult();
             while (true)
             {
-                using var context = RetryableWriteContext.Create(operationContext, binding, GetEffectiveRetryRequested());
+                using var context = new RetryableWriteContext(binding, GetEffectiveRetryRequested());
                 BsonDocument serverResponse = null;
                 try
                 {
@@ -170,7 +170,7 @@ protected override IEnumerable<BatchableCommandMessageSection> CreateCommandPayl
                     bulkWriteResults.TopLevelException = commandException;
                     serverResponse = commandException.Result;
                 }
-                catch (Exception exception)
+                catch (Exception exception) when (context.Channel is not null)
                 {
                     bulkWriteResults.TopLevelException = exception;
                 }

src/MongoDB.Driver/Core/Operations/CommandOperationBase.cs

@@ -27,21 +27,17 @@ namespace MongoDB.Driver.Core.Operations
     internal abstract class CommandOperationBase<TCommandResult>
     {
         private BsonDocument _additionalOptions;
-        private BsonDocument _command;
         private IElementNameValidator _commandValidator = NoOpElementNameValidator.Instance;
         private string _comment;
         private DatabaseNamespace _databaseNamespace;
         private MessageEncoderSettings _messageEncoderSettings;
         private IBsonSerializer<TCommandResult> _resultSerializer;
 
-        protected CommandOperationBase(
-            DatabaseNamespace databaseNamespace,
-            BsonDocument command,
+        protected CommandOperationBase(DatabaseNamespace databaseNamespace,
             IBsonSerializer<TCommandResult> resultSerializer,
             MessageEncoderSettings messageEncoderSettings)
         {
             _databaseNamespace = Ensure.IsNotNull(databaseNamespace, nameof(databaseNamespace));
-            _command = Ensure.IsNotNull(command, nameof(command));
             _resultSerializer = Ensure.IsNotNull(resultSerializer, nameof(resultSerializer));
             _messageEncoderSettings = messageEncoderSettings;
         }
@@ -52,11 +48,6 @@ public BsonDocument AdditionalOptions
             set { _additionalOptions = value; }
         }
 
-        public BsonDocument Command
-        {
-            get { return _command; }
-        }
-
         public IElementNameValidator CommandValidator
         {
             get { return _commandValidator; }
@@ -84,7 +75,7 @@ public IBsonSerializer<TCommandResult> ResultSerializer
             get { return _resultSerializer; }
         }
 
-        protected TCommandResult ExecuteProtocol(OperationContext operationContext, IChannelHandle channel, ICoreSessionHandle session, ReadPreference readPreference)
+        protected TCommandResult ExecuteProtocol(OperationContext operationContext, IChannelHandle channel, ICoreSessionHandle session, ReadPreference readPreference, BsonDocument command)
         {
             var additionalOptions = GetEffectiveAdditionalOptions();
 
@@ -93,7 +84,7 @@ protected TCommandResult ExecuteProtocol(OperationContext operationContext, ICha
                 session,
                 readPreference,
                 _databaseNamespace,
-                _command,
+                command,
                 null, // commandPayloads
                 _commandValidator,
                 additionalOptions,
@@ -107,15 +98,16 @@ protected TCommandResult ExecuteProtocol(
             OperationContext operationContext,
             IChannelSource channelSource,
             ICoreSessionHandle session,
-            ReadPreference readPreference)
+            ReadPreference readPreference,
+            BsonDocument command)
         {
             using (var channel = channelSource.GetChannel(operationContext))
             {
-                return ExecuteProtocol(operationContext, channel, session, readPreference);
+                return ExecuteProtocol(operationContext, channel, session, readPreference, command);
             }
         }
 
-        protected Task<TCommandResult> ExecuteProtocolAsync(OperationContext operationContext, IChannelHandle channel, ICoreSessionHandle session, ReadPreference readPreference)
+        protected Task<TCommandResult> ExecuteProtocolAsync(OperationContext operationContext, IChannelHandle channel, ICoreSessionHandle session, ReadPreference readPreference, BsonDocument command)
         {
             var additionalOptions = GetEffectiveAdditionalOptions();
 
@@ -124,7 +116,7 @@ protected Task<TCommandResult> ExecuteProtocolAsync(OperationContext operationCo
                 session,
                 readPreference,
                 _databaseNamespace,
-                _command,
+                command,
                 null, // TODO: support commandPayloads
                 _commandValidator,
                 additionalOptions,
@@ -138,11 +130,12 @@ protected async Task<TCommandResult> ExecuteProtocolAsync(
             OperationContext operationContext,
             IChannelSource channelSource,
             ICoreSessionHandle session,
-            ReadPreference readPreference)
+            ReadPreference readPreference,
+            BsonDocument command)
         {
             using (var channel = await channelSource.GetChannelAsync(operationContext).ConfigureAwait(false))
             {
-                return await ExecuteProtocolAsync(operationContext, channel, session, readPreference).ConfigureAwait(false);
+                return await ExecuteProtocolAsync(operationContext, channel, session, readPreference, command).ConfigureAwait(false);
             }
         }
 

src/MongoDB.Driver/Core/Operations/CountOperation.cs

@@ -25,7 +25,7 @@
 
 namespace MongoDB.Driver.Core.Operations
 {
-    internal sealed class CountOperation : IReadOperation<long>, IExecutableInRetryableReadContext<long>
+    internal sealed class CountOperation : IReadOperation<long>, IExecutableInRetryableReadContext<long>, ICommandCreator
     {
         private Collation _collation;
         private readonly CollectionNamespace _collectionNamespace;
@@ -133,7 +133,7 @@ public long Execute(OperationContext operationContext, IReadBinding binding)
             Ensure.IsNotNull(binding, nameof(binding));
 
             using (BeginOperation())
-            using (var context = RetryableReadContext.Create(operationContext, binding, _retryRequested))
+            using (var context = new RetryableReadContext(binding, _retryRequested))
             {
                 return Execute(operationContext, context);
             }
@@ -151,7 +151,7 @@ public async Task<long> ExecuteAsync(OperationContext operationContext, IReadBin
             Ensure.IsNotNull(binding, nameof(binding));
 
             using (BeginOperation())
-            using (var context = await RetryableReadContext.CreateAsync(operationContext, binding, _retryRequested).ConfigureAwait(false))
+            using (var context = new RetryableReadContext(binding, _retryRequested))
             {
                 return await ExecuteAsync(operationContext, context).ConfigureAwait(false);
             }
@@ -168,8 +168,12 @@ public async Task<long> ExecuteAsync(OperationContext operationContext, Retryabl
 
         private ReadCommandOperation<BsonDocument> CreateOperation(OperationContext operationContext, RetryableReadContext context)
         {
-            var command = CreateCommand(operationContext, context.Binding.Session, context.Channel.ConnectionDescription);
-            return new ReadCommandOperation<BsonDocument>(_collectionNamespace.DatabaseNamespace, command, BsonDocumentSerializer.Instance, _messageEncoderSettings, OperationName)
+            return new ReadCommandOperation<BsonDocument>(
+                _collectionNamespace.DatabaseNamespace,
+                this,
+                BsonDocumentSerializer.Instance,
+                _messageEncoderSettings,
+                OperationName)
             {
                 RetryRequested = _retryRequested // might be overridden by retryable read context
             };