From 2238c786c457e35486f7cac8ab556b2f15f88d5e Mon Sep 17 00:00:00 2001
From: Sam Harwell <Sam.Harwell@microsoft.com>
Date: Thu, 29 Oct 2020 10:44:59 -0700
Subject: [PATCH] Improve resilience to COM RCW cleanup issues in tests

---
 .../Threading/StaTaskScheduler.cs             | 39 +++++++------------
 1 file changed, 14 insertions(+), 25 deletions(-)

diff --git a/src/EditorFeatures/TestUtilities/Threading/StaTaskScheduler.cs b/src/EditorFeatures/TestUtilities/Threading/StaTaskScheduler.cs
index 4156e506ea8f4..471464f8eb8f6 100644
--- a/src/EditorFeatures/TestUtilities/Threading/StaTaskScheduler.cs
+++ b/src/EditorFeatures/TestUtilities/Threading/StaTaskScheduler.cs
@@ -5,6 +5,7 @@
 #nullable disable
 
 using System;
+using System.Runtime.InteropServices;
 using System.Threading;
 using System.Windows.Threading;
 
@@ -45,8 +46,7 @@ static StaTaskScheduler()
             // to DomainUnload is a reasonable place to do it.
             AppDomain.CurrentDomain.DomainUnload += (sender, e) =>
             {
-                GC.Collect();
-                GC.WaitForPendingFinalizers();
+                GC.GetTotalMemory(forceFullCollection: true);
             };
         }
 
@@ -55,28 +55,12 @@ public StaTaskScheduler()
         {
             using (var threadStartedEvent = new ManualResetEventSlim(initialState: false))
             {
-                DispatcherSynchronizationContext synchronizationContext = null;
+                Dispatcher staDispatcher = null;
                 StaThread = new Thread(() =>
                 {
-                    var oldContext = SynchronizationContext.Current;
-                    try
-                    {
-                        // All WPF Tests need a DispatcherSynchronizationContext and we dont want to block pending keyboard
-                        // or mouse input from the user. So use background priority which is a single level below user input.
-                        synchronizationContext = new DispatcherSynchronizationContext();
-
-                        // xUnit creates its own synchronization context and wraps any existing context so that messages are
-                        // still pumped as necessary. So we are safe setting it here, where we are not safe setting it in test.
-                        SynchronizationContext.SetSynchronizationContext(synchronizationContext);
-
-                        threadStartedEvent.Set();
-
-                        Dispatcher.Run();
-                    }
-                    finally
-                    {
-                        SynchronizationContext.SetSynchronizationContext(oldContext);
-                    }
+                    staDispatcher = Dispatcher.CurrentDispatcher;
+                    threadStartedEvent.Set();
+                    Dispatcher.Run();
                 });
                 StaThread.Name = $"{nameof(StaTaskScheduler)} thread";
                 StaThread.IsBackground = true;
@@ -84,7 +68,11 @@ public StaTaskScheduler()
                 StaThread.Start();
 
                 threadStartedEvent.Wait();
-                DispatcherSynchronizationContext = synchronizationContext;
+#pragma warning disable VSTHRD001 // Avoid legacy thread switching APIs
+                DispatcherSynchronizationContext = (DispatcherSynchronizationContext)staDispatcher.Invoke(() => SynchronizationContext.Current);
+#pragma warning restore VSTHRD001 // Avoid legacy thread switching APIs
+
+                AppDomain.CurrentDomain.DomainUnload += (_, _) => Dispose();
             }
 
             // Work around the WeakEventTable Shutdown race conditions
@@ -104,8 +92,9 @@ public void Dispose()
         {
             if (StaThread.IsAlive)
             {
-                DispatcherSynchronizationContext.Post(_ => Dispatcher.ExitAllFrames(), null);
-                StaThread.Join();
+                // The message pump is not active during AppDomain.Unload callbacks, so our only option is to abort the
+                // thread directly. 😢
+                StaThread.Abort();
             }
         }
     }