Switch back to a shared thread for WpfFact and WpfTheory

src/EditorFeatures/TestUtilities/Threading/StaTaskScheduler.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System;
+using System.Threading;
+using System.Windows.Threading;
+
+namespace Roslyn.Test.Utilities
+{
+    public sealed class StaTaskScheduler : IDisposable
+    {
+        /// <summary>Gets a <see cref="StaTaskScheduler"/> for the current <see cref="AppDomain"/>.</summary>
+        /// <remarks>We use a count of 1, because the editor ends up re-using <see cref="DispatcherObject"/>
+        /// instances between tests, so we need to always use the same thread for our Sta tests.</remarks>
+        public static StaTaskScheduler DefaultSta { get; } = new StaTaskScheduler();
+
+        /// <summary>The STA threads used by the scheduler.</summary>
+        public Thread StaThread { get; }
+
+        public bool IsRunningInScheduler => StaThread.ManagedThreadId == Thread.CurrentThread.ManagedThreadId;
+
+        /// <summary>Initializes a new instance of the <see cref="StaTaskScheduler"/> class.</summary>
+        public StaTaskScheduler()
+        {
+            using (var threadStartedEvent = new ManualResetEventSlim(initialState: false))
+            {
+                DispatcherSynchronizationContext synchronizationContext = 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);
+                    }
+                });
+                StaThread.Name = $"{nameof(StaTaskScheduler)} thread";
+                StaThread.IsBackground = true;
+                StaThread.SetApartmentState(ApartmentState.STA);
+                StaThread.Start();
+
+                threadStartedEvent.Wait();
+                DispatcherSynchronizationContext = synchronizationContext;
+            };
+        }
+
+        public DispatcherSynchronizationContext DispatcherSynchronizationContext
+        {
+            get;
+        }
+
+        /// <summary>
+        /// Cleans up the scheduler by indicating that no more tasks will be queued.
+        /// This method blocks until all threads successfully shutdown.
+        /// </summary>
+        public void Dispose()
+        {
+            if (StaThread.IsAlive)
+            {
+                DispatcherSynchronizationContext.Post(_ => Dispatcher.ExitAllFrames(), null);
+                StaThread.Join();
+            }
+        }
+    }
+}

src/EditorFeatures/TestUtilities/Threading/WpfTestRunner.cs

@@ -50,56 +50,39 @@ public sealed class WpfTestRunner : XunitTestRunner
         protected override Task<decimal> InvokeTestMethodAsync(ExceptionAggregator aggregator)
         {
             SharedData.ExecutingTest(TestMethod);
-
-            DispatcherSynchronizationContext synchronizationContext = null;
-            var staThreadStartedEvent = new ManualResetEventSlim(initialState: false);
-            var staThread = new Thread((ThreadStart)(() =>
-            {
-                // 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);
-
-                staThreadStartedEvent.Set();
-
-                Dispatcher.Run();
-            }));
-
-            staThread.Name = nameof(WpfTestRunner);
-            staThread.SetApartmentState(ApartmentState.STA);
-            staThread.Start();
-
-            staThreadStartedEvent.Wait();
-            Debug.Assert(synchronizationContext != null);
-
-            var taskScheduler = new SynchronizationContextTaskScheduler(synchronizationContext);
+            var sta = StaTaskScheduler.DefaultSta;
             var task = Task.Factory.StartNew(async () =>
             {
-                Debug.Assert(SynchronizationContext.Current is DispatcherSynchronizationContext);
+                Debug.Assert(sta.StaThread == Thread.CurrentThread);
 
                 using (await SharedData.TestSerializationGate.DisposableWaitAsync(CancellationToken.None))
                 {
-                    // Sync up FTAO to the context that we are creating here. 
-                    ForegroundThreadAffinitizedObject.CurrentForegroundThreadData = new ForegroundThreadData(
-                        Thread.CurrentThread,
-                        new SynchronizationContextTaskScheduler(new DispatcherSynchronizationContext(Dispatcher.CurrentDispatcher, DispatcherPriority.Background)),
-                        ForegroundThreadDataKind.StaUnitTest);
-
-                    // Reset our flag ensuring that part of this test actually needs WpfFact
-                    s_wpfFactRequirementReason = null;
-
-                    // Just call back into the normal xUnit dispatch process now that we are on an STA Thread with no synchronization context.
-                    var invoker = new WpfTestInvoker(SharedData, Test, MessageBus, TestClass, ConstructorArguments, TestMethod, TestMethodArguments, BeforeAfterAttributes, aggregator, CancellationTokenSource);
-                    return await invoker.RunAsync();
+                    try
+                    {
+                        Debug.Assert(SynchronizationContext.Current is DispatcherSynchronizationContext);
+
+                        // Sync up FTAO to the context that we are creating here. 
+                        ForegroundThreadAffinitizedObject.CurrentForegroundThreadData = new ForegroundThreadData(
+                            Thread.CurrentThread,
+                            new SynchronizationContextTaskScheduler(new DispatcherSynchronizationContext(Dispatcher.CurrentDispatcher, DispatcherPriority.Background)),
+                            ForegroundThreadDataKind.StaUnitTest);
+
+                        // Reset our flag ensuring that part of this test actually needs WpfFact
+                        s_wpfFactRequirementReason = null;
+
+                        // Just call back into the normal xUnit dispatch process now that we are on an STA Thread with no synchronization context.
+                        var invoker = new WpfTestInvoker(SharedData, Test, MessageBus, TestClass, ConstructorArguments, TestMethod, TestMethodArguments, BeforeAfterAttributes, aggregator, CancellationTokenSource);
+                        return invoker.RunAsync().JoinUsingDispatcher(CancellationTokenSource.Token);
+                    }
+                    finally
+                    {
+                        // Cleanup the synchronization context even if the test is failing exceptionally
+                        SynchronizationContext.SetSynchronizationContext(null);
+                    }
                 }
-            }, CancellationTokenSource.Token, TaskCreationOptions.None, taskScheduler).Unwrap();
-
-            task.SafeContinueWith(_ => Dispatcher.ExitAllFrames(), taskScheduler);
+            }, CancellationTokenSource.Token, TaskCreationOptions.None, new SynchronizationContextTaskScheduler(sta.DispatcherSynchronizationContext));
 
-            return task;
+            return task.Unwrap();
         }
 
         /// <summary>

src/EditorFeatures/TestUtilities/Threading/WpfTestSharedData.cs

@@ -136,7 +136,7 @@ void checkForBad()
             void queueCheckForBad()
             {
                 var task = new Task((Action)checkForBad);
-                task.Start(new SynchronizationContextTaskScheduler(new DispatcherSynchronizationContext(dispatcher, DispatcherPriority.Send)));
+                task.Start(new SynchronizationContextTaskScheduler(StaTaskScheduler.DefaultSta.DispatcherSynchronizationContext));
             }
 
             queueCheckForBad();