Init and cleanup timeout

src/Adapter/MSTest.TestAdapter/Execution/TestAssemblyInfo.cs

@@ -59,6 +59,8 @@ internal set
     /// </summary>
     internal int? AssemblyInitializeMethodTimeoutMilliseconds { get; set; }
 
+    internal int? AssemblyCleanupMethodTimeoutMilliseconds { get; set; }
+
     /// <summary>
     /// Gets <c>AssemblyCleanup</c> method for the assembly.
     /// </summary>
@@ -210,30 +212,48 @@ public void RunAssemblyInitialize(TestContext testContext)
 
         lock (_assemblyInfoExecuteSyncObject)
         {
+            Exception? assemblyCleanupException = null;
             try
             {
-                AssemblyCleanupMethod.InvokeAsSynchronousTask(null);
-
-                return null;
+                assemblyCleanupException = MethodRunner.RunWithTimeout(
+                     () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),
+                     AssemblyCleanupMethodTimeoutMilliseconds,
+                     AssemblyCleanupMethod,
+                     Resource.AssemblyCleanupWasCancelled,
+                     Resource.AssemblyCleanupTimedOut);
             }
             catch (Exception ex)
             {
-                var realException = ex.InnerException ?? ex;
+                assemblyCleanupException = ex;
+            }
 
-                // special case AssertFailedException to trim off part of the stack trace
-                string errorMessage = realException is AssertFailedException or AssertInconclusiveException
-                    ? realException.Message
-                    : realException.GetFormattedExceptionMessage();
+            // If assemblyCleanup was successful, then don't do anything
+            if (assemblyCleanupException is null)
+            {
+                return null;
+            }
 
-                DebugEx.Assert(AssemblyCleanupMethod.DeclaringType?.Name is not null, "AssemblyCleanupMethod.DeclaringType.Name is null");
-                return string.Format(
-                    CultureInfo.CurrentCulture,
-                    Resource.UTA_AssemblyCleanupMethodWasUnsuccesful,
-                    AssemblyCleanupMethod.DeclaringType.Name,
-                    AssemblyCleanupMethod.Name,
-                    errorMessage,
-                    realException.GetStackTraceInformation()?.ErrorStackTrace);
+            // Cache and return an already created TestFailedException.
+            if (assemblyCleanupException is TestFailedException)
+            {
+                throw assemblyCleanupException;
             }
+
+            var realException = assemblyCleanupException.InnerException ?? assemblyCleanupException;
+
+            // special case AssertFailedException to trim off part of the stack trace
+            string errorMessage = realException is AssertFailedException or AssertInconclusiveException
+                ? realException.Message
+                : realException.GetFormattedExceptionMessage();
+
+            DebugEx.Assert(AssemblyCleanupMethod.DeclaringType?.Name is not null, "AssemblyCleanupMethod.DeclaringType.Name is null");
+            return string.Format(
+                CultureInfo.CurrentCulture,
+                Resource.UTA_AssemblyCleanupMethodWasUnsuccesful,
+                AssemblyCleanupMethod.DeclaringType.Name,
+                AssemblyCleanupMethod.Name,
+                errorMessage,
+                realException.GetStackTraceInformation()?.ErrorStackTrace);
         }
     }
 
@@ -255,7 +275,24 @@ internal void ExecuteAssemblyCleanup()
         {
             try
             {
-                AssemblyCleanupMethod.InvokeAsSynchronousTask(null);
+                Exception? assemblyCleanupException = MethodRunner.RunWithTimeout(
+                     () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),
+                     AssemblyCleanupMethodTimeoutMilliseconds,
+                     AssemblyCleanupMethod,
+                     Resource.AssemblyCleanupWasCancelled,
+                     Resource.AssemblyCleanupTimedOut);
+
+                // If assemblyCleanup was successful, then don't do anything
+                if (assemblyCleanupException is null)
+                {
+                    return;
+                }
+
+                // Cache and return an already created TestFailedException.
+                if (assemblyCleanupException is TestFailedException)
+                {
+                    throw assemblyCleanupException;
+                }
             }
             catch (Exception ex)
             {

src/Adapter/MSTest.TestAdapter/Execution/TypeCache.cs

@@ -449,6 +449,22 @@ private TestAssemblyInfo GetAssemblyInfo(Type type)
                 else if (IsAssemblyOrClassCleanupMethod<AssemblyCleanupAttribute>(methodInfo))
                 {
                     assemblyInfo.AssemblyCleanupMethod = methodInfo;
+                    if (_reflectionHelper.IsAttributeDefined<TimeoutAttribute>(methodInfo, false))
+                    {
+                        if (!methodInfo.HasCorrectTimeout())
+                        {
+                            var message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_ErrorInvalidTimeout, methodInfo.DeclaringType!.FullName, methodInfo.Name);
+                            throw new TypeInspectionException(message);
+                        }
+
+                        var timeoutAttribute = _reflectionHelper.GetAttribute<TimeoutAttribute>(methodInfo);
+                        DebugEx.Assert(timeoutAttribute != null, "TimeoutAttribute cannot be null");
+                        assemblyInfo.AssemblyCleanupMethodTimeoutMilliseconds = timeoutAttribute.Timeout;
+                    }
+                    else if (MSTestSettings.CurrentSettings.AssemblyInitializeTimeout > 0)
+                    {
+                        assemblyInfo.AssemblyCleanupMethodTimeoutMilliseconds = MSTestSettings.CurrentSettings.AssemblyCleanupTimeout;
+                    }
                 }
             }
         }

src/Adapter/MSTest.TestAdapter/Helpers/MethodRunner.cs

@@ -12,6 +12,72 @@ namespace Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Helpers;
 
 internal static class MethodRunner
 {
+    internal static TestFailedException? RunWithTimeout(
+        Action action, int? timeout, MethodInfo methodInfo,
+        string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
+    {
+        Exception? realException = null;
+        Task? executionTask = null;
+        try
+        {
+            executionTask = Task.Run(
+                () =>
+                {
+                    try
+                    {
+                        action();
+                    }
+                    catch (Exception ex)
+                    {
+                        realException = ex;
+                        throw;
+                    }
+                });
+        }
+        catch (TaskCanceledException)
+        {
+            return new(
+                UnitTestOutcome.Timeout,
+                string.Format(CultureInfo.InvariantCulture, methodCancelledMessageFormat, methodInfo.DeclaringType!.FullName, methodInfo.Name));
+        }
+
+        try
+        {
+            if (timeout.HasValue && executionTask.Wait(timeout.Value))
+            {
+                // If we reach this line then the task has completed successfully.
+                return null;
+            }
+
+            // Timed out. For cancellation, either OCE or AggregateException with TCE will be thrown.
+            return new(
+                UnitTestOutcome.Timeout,
+                string.Format(
+                    CultureInfo.InvariantCulture,
+                    methodTimedOutMessageFormat,
+                    methodInfo.DeclaringType!.FullName,
+                    methodInfo.Name));
+        }
+        catch (Exception ex) when
+            (ex is OperationCanceledException
+            || (ex is AggregateException aggregateEx && aggregateEx.InnerExceptions.OfType<TaskCanceledException>().Any()))
+        {
+            return new(
+                UnitTestOutcome.Timeout,
+                string.Format(CultureInfo.InvariantCulture, methodCancelledMessageFormat, methodInfo.DeclaringType!.FullName, methodInfo.Name));
+        }
+        catch (Exception)
+        {
+            // We throw the real exception to have the original stack trace to elaborate up the chain.
+            if (realException is not null)
+            {
+                throw realException;
+            }
+
+            throw;
+        }
+    }
+
     internal static TestFailedException? RunWithTimeoutAndCancellation(
         Action action, CancellationTokenSource cancellationTokenSource, int? timeout, MethodInfo methodInfo,
         string methodCancelledMessageFormat, string methodTimedOutMessageFormat)

src/Adapter/MSTest.TestAdapter/MSTestSettings.cs

@@ -59,6 +59,7 @@ public MSTestSettings()
         TestTimeout = 0;
         AssemblyInitializeTimeout = 0;
         ClassInitializeTimeout = 0;
+        AssemblyCleanupTimeout = 0;
         TreatClassAndAssemblyCleanupWarningsAsErrors = false;
     }
 
@@ -151,6 +152,11 @@ public static RunConfigurationSettings RunConfigurationSettings
     /// </summary>
     internal int AssemblyInitializeTimeout { get; private set; }
 
+    /// <summary>
+    ///  Gets specified global AssemblyCleanup timeout.
+    /// </summary>
+    internal int AssemblyCleanupTimeout { get; private set; }
+
     /// <summary>
     ///  Gets specified global ClassInitializeTimeout timeout.
     /// </summary>
@@ -186,6 +192,7 @@ public static void PopulateSettings(MSTestSettings settings)
         CurrentSettings.TestTimeout = settings.TestTimeout;
         CurrentSettings.TreatClassAndAssemblyCleanupWarningsAsErrors = settings.TreatClassAndAssemblyCleanupWarningsAsErrors;
         CurrentSettings.AssemblyInitializeTimeout = settings.AssemblyInitializeTimeout;
+        CurrentSettings.AssemblyCleanupTimeout = settings.AssemblyCleanupTimeout;
         CurrentSettings.ClassInitializeTimeout = settings.ClassInitializeTimeout;
     }
 
@@ -436,6 +443,16 @@ private static MSTestSettings ToSettings(XmlReader reader)
                             break;
                         }
 
+                    case "ASSEMBLYCLEANUPTIMEOUT":
+                        {
+                            if (int.TryParse(reader.ReadInnerXml(), out int assemblyCleanupTimeout) && assemblyCleanupTimeout > 0)
+                            {
+                                settings.AssemblyCleanupTimeout = assemblyCleanupTimeout;
+                            }
+
+                            break;
+                        }
+
                     case "ASSEMBLYINITIALIZETIMEOUT":
                         {
                             if (int.TryParse(reader.ReadInnerXml(), out int assemblyInitializeTimeout) && assemblyInitializeTimeout > 0)

src/Adapter/MSTest.TestAdapter/Resources/Resource.resx

@@ -375,4 +375,10 @@ Error: {1}</value>
   <data name="UTA_ErrorTestClassIsGenericNonAbstract" xml:space="preserve">
     <value>TestClass attribute defined on generic non-abstract class {0}</value>
   </data>
+  <data name="AssemblyCleanupTimedOut" xml:space="preserve">
+    <value>Assembly cleanup method '{0}.{1}' timed out</value>
+  </data>
+  <data name="AssemblyCleanupWasCancelled" xml:space="preserve">
+    <value>Assembly cleanup method '{0}.{1}' was cancelled</value>
+  </data>
 </root>