CSHARP-5798: Implement test to track if any UnobservedTaskExceptions were raised while test run

Author: sanych-sun

src/MongoDB.Driver/Core/Connections/TcpStreamFactory.cs

@@ -167,17 +167,16 @@ private void ConfigureConnectedSocket(Socket socket)
 
         private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancellationToken)
         {
-            var callbackState = new ConnectOperationState(socket);
+            var callbackState = new OperationCallbackState<Socket>(socket);
             using var timeoutCancellationTokenSource = new CancellationTokenSource(_settings.ConnectTimeout);
             using var combinedCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCancellationTokenSource.Token);
             using var cancellationSubscription = combinedCancellationTokenSource.Token.Register(state =>
             {
-                var operationState = (ConnectOperationState)state;
-                if (operationState.IsSucceeded)
+                var operationState = (OperationCallbackState<Socket>)state;
+                if (operationState.TryChangeStatusFromInProgress(OperationCallbackState<Socket>.OperationStatus.Interrupted))
                 {
-                    return;
+                    DisposeSocket(operationState.Subject);
                 }
-                DisposeSocket(operationState.Socket);
             }, callbackState);
 
             try
@@ -195,8 +194,10 @@ private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancell
 #else
                 socket.Connect(endPoint);
 #endif
-                EnsureConnected(socket);
-                callbackState.IsSucceeded = true;
+                if (!callbackState.TryChangeStatusFromInProgress(OperationCallbackState<Socket>.OperationStatus.Done))
+                {
+                    throw new ObjectDisposedException(nameof(Socket));
+                }
             }
             catch
             {
@@ -222,23 +223,6 @@ static void DisposeSocket(Socket socket)
                     // Ignore any exceptions.
                 }
             }
-
-            static void EnsureConnected(Socket socket)
-            {
-                bool originalBlockingState = socket.Blocking;
-
-                try
-                {
-                    socket.Blocking = false;
-                    // Try to use the socket to ensure it's connected. On MacOS with net6.0 sometimes Connect is completed successfully even after the socket disposal.
-                    socket.Send(__ensureConnectedBuffer, 0, 0);
-                }
-                finally
-                {
-                    // Restore original blocking state
-                    socket.Blocking = originalBlockingState;
-                }
-            }
         }
 
         private async Task ConnectAsync(Socket socket, EndPoint endPoint, CancellationToken cancellationToken)
@@ -401,10 +385,5 @@ public int Compare(EndPoint x, EndPoint y)
                 return 0;
             }
         }
-
-        private sealed record ConnectOperationState(Socket Socket)
-        {
-            public bool IsSucceeded { get; set; }
-        }
     }
 }

src/MongoDB.Driver/Core/Misc/OperationCallbackState.cs

@@ -0,0 +1,37 @@
+/* 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.Threading;
+
+namespace MongoDB.Driver.Core.Misc;
+
+internal sealed class OperationCallbackState<T>(T subject)
+{
+    private int _status = (int)OperationStatus.InProgress;
+
+    public OperationStatus Status => (OperationStatus)_status;
+    public T Subject => subject;
+    public bool TryChangeStatusFromInProgress(OperationStatus newState) =>
+        Interlocked.CompareExchange(ref _status, (int)newState, (int)OperationStatus.InProgress) == (int)OperationStatus.InProgress;
+
+    public enum OperationStatus
+    {
+        InProgress = 0,
+        Done,
+        Interrupted,
+    }
+}
+
+

src/MongoDB.Driver/Core/Misc/StreamExtensionMethods.cs

@@ -287,33 +287,33 @@ private static void ExecuteOperationWithTimeout<TState>(Stream stream, TState st
                 throw new TimeoutException();
             }
 
-            StreamDisposeCallbackState callbackState = null;
+            OperationCallbackState<Stream> callbackState = null;
             Timer timer = null;
             CancellationTokenRegistration cancellationSubscription = default;
             if (timeoutMs > 0)
             {
-                callbackState = new StreamDisposeCallbackState(stream);
+                callbackState = new OperationCallbackState<Stream>(stream);
                 timer = new Timer(DisposeStreamCallback, callbackState, timeoutMs, Timeout.Infinite);
             }
 
             if (cancellationToken.CanBeCanceled)
             {
-                callbackState ??= new StreamDisposeCallbackState(stream);
+                callbackState ??= new OperationCallbackState<Stream>(stream);
                 cancellationSubscription = cancellationToken.Register(DisposeStreamCallback, callbackState);
             }
 
             try
             {
                 operation(stream, state);
-                if (callbackState?.TryChangeStateFromInProgress(OperationState.Done) == false)
+                if (callbackState?.TryChangeStatusFromInProgress(OperationCallbackState<Stream>.OperationStatus.Done) == false)
                 {
                     // If the state can't be changed - then the stream was/will be disposed, throw here
                     throw new IOException();
                 }
             }
             catch (Exception ex)
             {
-                if (callbackState?.OperationState == OperationState.Interrupted)
+                if (callbackState?.Status == OperationCallbackState<Stream>.OperationStatus.Interrupted)
                 {
                     cancellationToken.ThrowIfCancellationRequested();
                     throw new TimeoutException();
@@ -334,39 +334,22 @@ private static void ExecuteOperationWithTimeout<TState>(Stream stream, TState st
 
             static void DisposeStreamCallback(object state)
             {
-                var disposeCallbackState = (StreamDisposeCallbackState)state;
-                if (!disposeCallbackState.TryChangeStateFromInProgress(OperationState.Interrupted))
+                var disposeCallbackState = (OperationCallbackState<Stream>)state;
+                if (!disposeCallbackState.TryChangeStatusFromInProgress(OperationCallbackState<Stream>.OperationStatus.Interrupted))
                 {
                     // If the state can't be changed - then I/O had already succeeded
                     return;
                 }
 
                 try
                 {
-                    disposeCallbackState.Stream.Dispose();
+                    disposeCallbackState.Subject.Dispose();
                 }
                 catch (Exception)
                 {
                     // callbacks should not fail, suppress any exceptions here
                 }
             }
         }
-
-        private record StreamDisposeCallbackState(Stream Stream)
-        {
-            private int _operationState = (int)OperationState.InProgress;
-
-            public OperationState OperationState => (OperationState)_operationState;
-
-            public bool TryChangeStateFromInProgress(OperationState newState) =>
-                Interlocked.CompareExchange(ref _operationState, (int)newState, (int)OperationState.InProgress) == (int)OperationState.InProgress;
-        }
-
-        private enum OperationState
-        {
-            InProgress = 0,
-            Done,
-            Interrupted,
-        }
     }
 }