CSHARP-3919: Drivers should retry operations if connection handshake fails

Author: BorisDog

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

@@ -177,46 +177,30 @@ public void ReplaceChannelSource(IChannelSourceHandle channelSource)
         private void Initialize(CancellationToken cancellationToken)
         {
             _channelSource = _binding.GetReadChannelSource(cancellationToken);
-            var serverDescription = _channelSource.ServerDescription;
 
             try
             {
                 _channel = _channelSource.GetChannel(cancellationToken);
             }
-            catch (MongoConnectionPoolPausedException)
+            catch (Exception ex) when (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(this, ex))
             {
-                if (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(this, serverDescription))
-                {
-                    ReplaceChannelSource(_binding.GetReadChannelSource(cancellationToken));
-                    ReplaceChannel(_channelSource.GetChannel(cancellationToken));
-                }
-                else
-                {
-                    throw;
-                }
+                ReplaceChannelSource(_binding.GetReadChannelSource(cancellationToken));
+                ReplaceChannel(_channelSource.GetChannel(cancellationToken));
             }
         }
 
         private async Task InitializeAsync(CancellationToken cancellationToken)
         {
             _channelSource = await _binding.GetReadChannelSourceAsync(cancellationToken).ConfigureAwait(false);
-            var serverDescription = _channelSource.ServerDescription;
 
             try
             {
                 _channel = await _channelSource.GetChannelAsync(cancellationToken).ConfigureAwait(false);
             }
-            catch (MongoConnectionPoolPausedException)
+            catch (Exception ex) when (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(this, ex))
             {
-                if (RetryableReadOperationExecutor.ShouldConnectionAcquireBeRetried(this, serverDescription))
-                {
-                    ReplaceChannelSource(await _binding.GetReadChannelSourceAsync(cancellationToken).ConfigureAwait(false));
-                    ReplaceChannel(await _channelSource.GetChannelAsync(cancellationToken).ConfigureAwait(false));
-                }
-                else
-                {
-                    throw;
-                }
+                ReplaceChannelSource(await _binding.GetReadChannelSourceAsync(cancellationToken).ConfigureAwait(false));
+                ReplaceChannel(await _channelSource.GetChannelAsync(cancellationToken).ConfigureAwait(false));
             }
         }
     }

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

@@ -17,7 +17,6 @@
 using System.Threading;
 using System.Threading.Tasks;
 using MongoDB.Driver.Core.Bindings;
-using MongoDB.Driver.Core.Servers;
 
 namespace MongoDB.Driver.Core.Operations
 {
@@ -115,9 +114,11 @@ public static async Task<TResult> ExecuteAsync<TResult>(IRetryableReadOperation<
             }
         }
 
-        public static bool ShouldConnectionAcquireBeRetried(RetryableReadContext context, ServerDescription serverDescription)
+        public static bool ShouldConnectionAcquireBeRetried(RetryableReadContext context, Exception ex)
         {
-            return context.RetryRequested && !context.Binding.Session.IsInTransaction;
+            // According the spec error during handshake should be handle according to RetryableReads logic
+            var innerException = ex is MongoAuthenticationException mongoAuthenticationException ? mongoAuthenticationException.InnerException : ex;
+            return context.RetryRequested && !context.Binding.Session.IsInTransaction && RetryabilityHelper.IsRetryableReadException(innerException);
         }
 
         // private static methods

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

@@ -201,17 +201,10 @@ private void Initialize(CancellationToken cancellationToken)
             {
                 _channel = _channelSource.GetChannel(cancellationToken);
             }
-            catch (MongoConnectionPoolPausedException)
+            catch (Exception ex) when (RetryableWriteOperationExecutor.ShouldConnectionAcquireBeRetried(this, serverDescription, ex))
             {
-                if (RetryableWriteOperationExecutor.ShouldConnectionAcquireBeRetried(this, serverDescription))
-                {
-                    ReplaceChannelSource(_binding.GetWriteChannelSource(cancellationToken));
-                    ReplaceChannel(_channelSource.GetChannel(cancellationToken));
-                }
-                else
-                {
-                    throw;
-                }
+                ReplaceChannelSource(_binding.GetWriteChannelSource(cancellationToken));
+                ReplaceChannel(_channelSource.GetChannel(cancellationToken));
             }
         }
 
@@ -224,17 +217,10 @@ private async Task InitializeAsync(CancellationToken cancellationToken)
             {
                 _channel = await _channelSource.GetChannelAsync(cancellationToken).ConfigureAwait(false);
             }
-            catch (MongoConnectionPoolPausedException)
+            catch (Exception ex) when (RetryableWriteOperationExecutor.ShouldConnectionAcquireBeRetried(this, serverDescription, ex))
             {
-                if (RetryableWriteOperationExecutor.ShouldConnectionAcquireBeRetried(this, serverDescription))
-                {
-                    ReplaceChannelSource(await _binding.GetWriteChannelSourceAsync(cancellationToken).ConfigureAwait(false));
-                    ReplaceChannel(await _channelSource.GetChannelAsync(cancellationToken).ConfigureAwait(false));
-                }
-                else
-                {
-                    throw;
-                }
+                ReplaceChannelSource(await _binding.GetWriteChannelSourceAsync(cancellationToken).ConfigureAwait(false));
+                ReplaceChannel(await _channelSource.GetChannelAsync(cancellationToken).ConfigureAwait(false));
             }
         }
     }

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

@@ -127,12 +127,16 @@ public static async Task<TResult> ExecuteAsync<TResult>(IRetryableWriteOperation
             }
         }
 
-        public static bool ShouldConnectionAcquireBeRetried(RetryableWriteContext context, ServerDescription serverDescription)
+        public static bool ShouldConnectionAcquireBeRetried(RetryableWriteContext context, ServerDescription serverDescription, Exception exception)
         {
+            var innerException = exception is MongoAuthenticationException mongoAuthenticationException ? mongoAuthenticationException.InnerException : exception;
+
+            // According the spec error during handshake should be handle according to RetryableReads logic
             return context.RetryRequested &&
                 AreRetryableWritesSupported(serverDescription) &&
                 context.Binding.Session.Id != null &&
-                !context.Binding.Session.IsInTransaction;
+                !context.Binding.Session.IsInTransaction &&
+                RetryabilityHelper.IsRetryableReadException(innerException);
         }
 
         // private static methods

tests/MongoDB.Driver.Tests/Specifications/retryable-reads/RetryableReadsTestRunner.cs

@@ -481,7 +481,7 @@ public class TestCaseFactory : JsonDrivenTestCaseFactory
             #endregion
             // protected properties
             // the path is "retryable-reads" but the namespace is "retryable_reads"
-            protected override string PathPrefix => "MongoDB.Driver.Tests.Specifications.retryable_reads.tests.";
+            protected override string PathPrefix => "MongoDB.Driver.Tests.Specifications.retryable_reads.tests.legacy";
 
             // protected methods
             protected override IEnumerable<JsonDrivenTestCase> CreateTestCases(BsonDocument document)

tests/MongoDB.Driver.Tests/Specifications/retryable-reads/RetryableReadsUnifiedTestRunner.cs

@@ -0,0 +1,68 @@
+/* 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.
+* You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System.Collections.Generic;
+using MongoDB.Bson;
+using MongoDB.Bson.TestHelpers.JsonDrivenTests;
+using MongoDB.Driver.Core.TestHelpers.Logging;
+using MongoDB.Driver.Tests.UnifiedTestOperations;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace MongoDB.Driver.Tests.Specifications.retryable_reads
+{
+    [Trait("Category", "SupportLoadBalancing")]
+    [Trait("Category", "Serverless")]
+    public sealed class RetryableReadsUnifiedTestRunner : LoggableTestClass
+    {
+        // public constructors
+        public RetryableReadsUnifiedTestRunner(ITestOutputHelper testOutputHelper)
+            : base(testOutputHelper)
+        {
+        }
+
+        // public methods
+        [SkippableTheory]
+        [ClassData(typeof(TestCaseFactory))]
+        public void Run(JsonDrivenTestCase testCase)
+        {
+            using (var runner = new UnifiedTestRunner())
+            {
+                runner.Run(testCase);
+            }
+        }
+
+        // nested types
+        public class TestCaseFactory : JsonDrivenTestCaseFactory
+        {
+            // protected properties
+            protected override string PathPrefix => "MongoDB.Driver.Tests.Specifications.retryable_reads.tests.unified";
+
+            // protected methods
+            protected override IEnumerable<JsonDrivenTestCase> CreateTestCases(BsonDocument document)
+            {
+                foreach (var testCase in base.CreateTestCases(document))
+                {
+                    foreach (var async in new[] { false, true })
+                    {
+                        var name = $"{testCase.Name}:async={async}";
+                        var test = testCase.Test.DeepClone().AsBsonDocument.Add("async", async);
+                        yield return new JsonDrivenTestCase(name, testCase.Shared, test);
+                    }
+                }
+            }
+        }
+    }
+}