CSHARP-5798: Implement test to track if any UnobservedTaskExceptions were raised while test run (#1836)

Author: sanych-sun

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

@@ -165,10 +165,17 @@ private void ConfigureConnectedSocket(Socket socket)
 
         private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancellationToken)
         {
-            var isSocketDisposed = false;
+            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(DisposeSocket);
+            using var cancellationSubscription = combinedCancellationTokenSource.Token.Register(state =>
+            {
+                var operationState = (OperationCallbackState<Socket>)state;
+                if (operationState.TryChangeStatusFromInProgress(OperationCallbackState<Socket>.OperationStatus.Interrupted))
+                {
+                    DisposeSocket(operationState.Subject);
+                }
+            }, callbackState);
 
             try
             {
@@ -185,13 +192,14 @@ private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancell
 #else
                 socket.Connect(endPoint);
 #endif
+                if (!callbackState.TryChangeStatusFromInProgress(OperationCallbackState<Socket>.OperationStatus.Done))
+                {
+                    throw new ObjectDisposedException(nameof(Socket));
+                }
             }
             catch
             {
-                if (!isSocketDisposed)
-                {
-                    DisposeSocket();
-                }
+                DisposeSocket(socket);
 
                 cancellationToken.ThrowIfCancellationRequested();
                 if (timeoutCancellationTokenSource.IsCancellationRequested)
@@ -202,9 +210,8 @@ private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancell
                 throw;
             }
 
-            void DisposeSocket()
+            static void DisposeSocket(Socket socket)
             {
-                isSocketDisposed = true;
                 try
                 {
                     socket.Dispose();

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

@@ -0,0 +1,35 @@
+/* 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,
-        }
     }
 }

tests/MongoDB.TestHelpers/XunitExtensions/TimeoutEnforcing/UnobservedExceptionTestDiscoverer.cs

@@ -42,13 +42,16 @@ public UnobservedExceptionTestDiscoverer(IMessageSink diagnosticsMessageSink)
 
     public IEnumerable<IXunitTestCase> Discover(ITestFrameworkDiscoveryOptions discoveryOptions, ITestMethod testMethod, IAttributeInfo factAttribute)
     {
-        return [new XunitTestCase(_diagnosticsMessageSink, TestMethodDisplay.Method, TestMethodDisplayOptions.All, testMethod)
+        var testCase = new XunitTestCase(_diagnosticsMessageSink, TestMethodDisplay.Method, TestMethodDisplayOptions.All, testMethod);
+        if (!testCase.Traits.TryGetValue("Category", out var categories))
         {
-            Traits =
-            {
-                { "Category", ["UnobservedExceptionTracking"] }
-            }
-        }];
+            categories = new List<string>();
+            testCase.Traits.Add("Category", categories);
+        }
+
+        categories.Add("UnobservedExceptionTracking");
+
+        return [testCase];
     }
 
     void UnobservedTaskExceptionEventHandler(object sender, UnobservedTaskExceptionEventArgs unobservedException) =>