CSHARP-3537: CSOT: retryable reads and writes

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

@@ -139,10 +139,9 @@ public WriteConcern WriteConcern
         public BulkWriteOperationResult Execute(OperationContext operationContext, IWriteBinding binding)
         {
             using (BeginOperation())
-            using (var context = RetryableWriteContext.Create(operationContext, binding, _retryRequested))
+            using (var context = RetryableWriteContext.Create(operationContext, binding, IsOperationRetryable()))
             {
                 EnsureHintIsSupportedIfAnyRequestHasHint();
-                context.DisableRetriesIfAnyWriteRequestIsNotRetryable(_requests);
                 var helper = new BatchHelper(_requests, _isOrdered, _writeConcern);
                 foreach (var batch in helper.GetBatches())
                 {
@@ -155,10 +154,9 @@ 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, _retryRequested).ConfigureAwait(false))
+            using (var context = await RetryableWriteContext.CreateAsync(operationContext, binding, IsOperationRetryable()).ConfigureAwait(false))
             {
                 EnsureHintIsSupportedIfAnyRequestHasHint();
-                context.DisableRetriesIfAnyWriteRequestIsNotRetryable(_requests);
                 var helper = new BatchHelper(_requests, _isOrdered, _writeConcern);
                 foreach (var batch in helper.GetBatches())
                 {
@@ -168,6 +166,9 @@ public async Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operat
             }
         }
 
+        private bool IsOperationRetryable()
+            => _retryRequested && _requests.All(r => r.IsRetryable());
+
         private IDisposable BeginOperation() =>
             // Execution starts with the first request
             EventContext.BeginOperation(null, _requests.FirstOrDefault()?.RequestType.ToString().ToLower());

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

@@ -129,9 +129,8 @@ public BulkWriteOperationResult Execute(OperationContext operationContext, Retry
         public BulkWriteOperationResult Execute(OperationContext operationContext, IWriteBinding binding)
         {
             using (BeginOperation())
-            using (var context = RetryableWriteContext.Create(operationContext, binding, _retryRequested))
+            using (var context = RetryableWriteContext.Create(operationContext, binding, IsOperationRetryable()))
             {
-                context.DisableRetriesIfAnyWriteRequestIsNotRetryable(_requests);
                 return Execute(operationContext, context);
             }
         }
@@ -146,9 +145,8 @@ 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, _retryRequested).ConfigureAwait(false))
+            using (var context = await RetryableWriteContext.CreateAsync(operationContext, binding, IsOperationRetryable()).ConfigureAwait(false))
             {
-                context.DisableRetriesIfAnyWriteRequestIsNotRetryable(_requests);
                 return await ExecuteAsync(operationContext, context).ConfigureAwait(false);
             }
         }
@@ -159,6 +157,9 @@ public async Task<BulkWriteOperationResult> ExecuteAsync(OperationContext operat
         protected abstract bool RequestHasHint(TWriteRequest request);
 
         // private methods
+        private bool IsOperationRetryable()
+            => _retryRequested && _requests.All(r => r.IsRetryable());
+
         private IDisposable BeginOperation() =>
             EventContext.BeginOperation(null, _requests.FirstOrDefault()?.RequestType.ToString().ToLower());
 

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

@@ -14,7 +14,6 @@
 */
 
 using MongoDB.Bson;
-using MongoDB.Driver.Core.Connections;
 using MongoDB.Driver.Core.Misc;
 
 namespace MongoDB.Driver.Core.Operations
@@ -36,6 +35,6 @@ public DeleteRequest(BsonDocument filter)
         public int Limit { get; init; }
 
         // public methods
-        public override bool IsRetryable(ConnectionDescription connectionDescription) => Limit != 0;
+        public override bool IsRetryable() => Limit != 0;
     }
 }

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

@@ -14,7 +14,6 @@
 */
 
 using MongoDB.Bson;
-using MongoDB.Driver.Core.Connections;
 using MongoDB.Driver.Core.Misc;
 
 namespace MongoDB.Driver.Core.Operations
@@ -32,6 +31,6 @@ public InsertRequest(BsonDocument document)
         public BsonDocument Document { get; }
 
         // public methods
-        public override bool IsRetryable(ConnectionDescription connectionDescription) => true;
+        public override bool IsRetryable() => true;
     }
 }

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

@@ -1,4 +1,4 @@
-/* Copyright 2018-present MongoDB Inc.
+/* Copyright 2010-present MongoDB Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -135,20 +135,17 @@ public static bool IsResumableChangeStreamException(Exception exception, int max
             {
                 return exception is MongoException mongoException ? mongoException.HasErrorLabel(ResumableChangeStreamErrorLabel) : false;
             }
-            else
+
+            if (exception is MongoCommandException commandException)
             {
-                var commandException = exception as MongoCommandException;
-                if (commandException != null)
+                var code = (ServerErrorCode)commandException.Code;
+                if (__resumableChangeStreamErrorCodes.Contains(code))
                 {
-                    var code = (ServerErrorCode)commandException.Code;
-                    if (__resumableChangeStreamErrorCodes.Contains(code))
-                    {
-                        return true;
-                    }
+                    return true;
                 }
-
-                return __resumableChangeStreamExceptions.Contains(exception.GetType());
             }
+
+            return __resumableChangeStreamExceptions.Contains(exception.GetType());
         }
 
         /// <summary>

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

@@ -1,4 +1,4 @@
-/* Copyright 2019-present MongoDB Inc.
+/* Copyright 2010-present MongoDB Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -14,9 +14,11 @@
 */
 
 using System;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 using MongoDB.Driver.Core.Bindings;
 using MongoDB.Driver.Core.Misc;
+using MongoDB.Driver.Core.Servers;
 
 namespace MongoDB.Driver.Core.Operations
 {
@@ -29,41 +31,33 @@ public static RetryableReadContext Create(OperationContext operationContext, IRe
             var context = new RetryableReadContext(binding, retryRequested);
             try
             {
-                context.Initialize(operationContext);
-
-                ChannelPinningHelper.PinChannellIfRequired(
-                    context.ChannelSource,
-                    context.Channel,
-                    context.Binding.Session);
-
-                return context;
+                context.AcquireOrReplaceChannel(operationContext, null);
             }
             catch
             {
                 context.Dispose();
                 throw;
             }
+
+            ChannelPinningHelper.PinChannellIfRequired(context.ChannelSource, context.Channel, context.Binding.Session);
+            return context;
         }
 
         public static async Task<RetryableReadContext> CreateAsync(OperationContext operationContext, IReadBinding binding, bool retryRequested)
         {
             var context = new RetryableReadContext(binding, retryRequested);
             try
             {
-                await context.InitializeAsync(operationContext).ConfigureAwait(false);
-
-                ChannelPinningHelper.PinChannellIfRequired(
-                    context.ChannelSource,
-                    context.Channel,
-                    context.Binding.Session);
-
-                return context;
+                await context.AcquireOrReplaceChannelAsync(operationContext, null).ConfigureAwait(false);
             }
             catch
             {
                 context.Dispose();
                 throw;
             }
+
+            ChannelPinningHelper.PinChannellIfRequired(context.ChannelSource, context.Channel, context.Binding.Session);
+            return context;
         }
         #endregion
 
@@ -96,50 +90,58 @@ public void Dispose()
             }
         }
 
-        public void ReplaceChannel(IChannelHandle channel)
+        public void AcquireOrReplaceChannel(OperationContext operationContext, IReadOnlyCollection<ServerDescription> deprioritizedServers)
+        {
+            var attempt = 1;
+            while (true)
+            {
+                operationContext.ThrowIfTimedOutOrCanceled();
+                ReplaceChannelSource(Binding.GetReadChannelSource(operationContext, deprioritizedServers));
+                try
+                {
+                    ReplaceChannel(ChannelSource.GetChannel(operationContext));
+                    return;
+                }
+                catch (Exception ex) when (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(operationContext, this, ex, attempt))
+                {
+                    attempt++;
+                }
+            }
+        }
+
+        public async Task AcquireOrReplaceChannelAsync(OperationContext operationContext, IReadOnlyCollection<ServerDescription> deprioritizedServers)
+        {
+            var attempt = 1;
+            while (true)
+            {
+                operationContext.ThrowIfTimedOutOrCanceled();
+                ReplaceChannelSource(await Binding.GetReadChannelSourceAsync(operationContext, deprioritizedServers).ConfigureAwait(false));
+                try
+                {
+                    ReplaceChannel(await ChannelSource.GetChannelAsync(operationContext).ConfigureAwait(false));
+                    return;
+                }
+                catch (Exception ex) when (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(operationContext, this, ex, attempt))
+                {
+                    attempt++;
+                }
+            }
+        }
+
+        private void ReplaceChannel(IChannelHandle channel)
         {
             Ensure.IsNotNull(channel, nameof(channel));
             _channel?.Dispose();
             _channel = channel;
         }
 
-        public void ReplaceChannelSource(IChannelSourceHandle channelSource)
+        private void ReplaceChannelSource(IChannelSourceHandle channelSource)
         {
             Ensure.IsNotNull(channelSource, nameof(channelSource));
             _channelSource?.Dispose();
             _channel?.Dispose();
             _channelSource = channelSource;
             _channel = null;
         }
-
-        private void Initialize(OperationContext operationContext)
-        {
-            _channelSource = _binding.GetReadChannelSource(operationContext);
-
-            try
-            {
-                _channel = _channelSource.GetChannel(operationContext);
-            }
-            catch (Exception ex) when (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(this, ex))
-            {
-                ReplaceChannelSource(_binding.GetReadChannelSource(operationContext));
-                ReplaceChannel(_channelSource.GetChannel(operationContext));
-            }
-        }
-
-        private async Task InitializeAsync(OperationContext operationContext)
-        {
-            _channelSource = await _binding.GetReadChannelSourceAsync(operationContext).ConfigureAwait(false);
-
-            try
-            {
-                _channel = await _channelSource.GetChannelAsync(operationContext).ConfigureAwait(false);
-            }
-            catch (Exception ex) when (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(this, ex))
-            {
-                ReplaceChannelSource(await _binding.GetReadChannelSourceAsync(operationContext).ConfigureAwait(false));
-                ReplaceChannel(await _channelSource.GetChannelAsync(operationContext).ConfigureAwait(false));
-            }
-        }
     }
 }