Add TaskSchedulers to the dispatcher abstraction

src/Microsoft.CodeAnalysis.Razor.Workspaces/ForegroundDispatcher.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host;
 
 namespace Microsoft.CodeAnalysis.Razor
@@ -11,6 +12,10 @@ namespace Microsoft.CodeAnalysis.Razor
     {
         public abstract bool IsForegroundThread { get; }
 
+        public abstract TaskScheduler ForegroundScheduler { get; }
+
+        public abstract TaskScheduler BackgroundScheduler { get; }
+
         public virtual void AssertForegroundThread([CallerMemberName] string caller = null)
         {
             if (!IsForegroundThread)

src/Microsoft.VisualStudio.LanguageServices.Razor/VisualStudioForegroundDispatcher.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System.Composition;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.VisualStudio.Shell;
@@ -12,6 +13,10 @@ namespace Microsoft.VisualStudio.LanguageServices.Razor
     [ExportWorkspaceService(typeof(ForegroundDispatcher), ServiceLayer.Host)]
     internal class VisualStudioForegroundDispatcher : ForegroundDispatcher
     {
+        public override TaskScheduler BackgroundScheduler { get; } = TaskScheduler.Default;
+
+        public override TaskScheduler ForegroundScheduler { get; } = TaskScheduler.FromCurrentSynchronizationContext();
+
         public override bool IsForegroundThread => ThreadHelper.CheckAccess();
     }
 }

test/Microsoft.VisualStudio.LanguageServices.Razor.Test/ForegroundDispatcherTestBase.cs

@@ -1,20 +0,0 @@
-// Copyright (c) .NET Foundation. All rights reserved.
-// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
-
-using System.Threading;
-using Microsoft.CodeAnalysis.Razor;
-
-namespace Microsoft.VisualStudio.LanguageServices.Razor
-{
-    public abstract class ForegroundDispatcherTestBase
-    {
-        internal ForegroundDispatcher Dispatcher { get; } = new SingleThreadedForegroundDispatcher();
-
-        private class SingleThreadedForegroundDispatcher : ForegroundDispatcher
-        {
-            private Thread Thread { get; } = Thread.CurrentThread;
-
-            public override bool IsForegroundThread => Thread.CurrentThread == Thread;
-        }
-    }
-}

test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Microsoft.VisualStudio.LanguageServices.Razor.Test.csproj

@@ -13,6 +13,10 @@
     <None Include="xunit.runner.json" CopyToOutputDirectory="PreserveNewest" />
   </ItemGroup>
 
+  <ItemGroup>
+    <Reference Include="WindowsBase" />
+  </ItemGroup>
+
   <ItemGroup>
     <ProjectReference Include="..\..\src\Microsoft.VisualStudio.LanguageServices.Razor\Microsoft.VisualStudio.LanguageServices.Razor.csproj" />
     <ProjectReference Include="..\..\src\Microsoft.AspNetCore.Razor.Runtime\Microsoft.AspNetCore.Razor.Runtime.csproj" />

test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Xunit/ForegroundDispatcherTestBase.cs

@@ -0,0 +1,52 @@
+// Copyright (c) .NET Foundation. 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.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Razor;
+
+namespace Xunit
+{
+    public abstract class ForegroundDispatcherTestBase
+    {
+        internal ForegroundDispatcher Dispatcher { get; } = new SingleThreadedForegroundDispatcher();
+
+        private class SingleThreadedForegroundDispatcher : ForegroundDispatcher
+        {
+            public SingleThreadedForegroundDispatcher()
+            {
+                ForegroundScheduler = SynchronizationContext.Current == null ? new ThrowingTaskScheduler() : TaskScheduler.FromCurrentSynchronizationContext();
+                BackgroundScheduler = TaskScheduler.Default;
+            }
+
+            public override TaskScheduler ForegroundScheduler { get; }
+
+            public override TaskScheduler BackgroundScheduler { get; }
+
+            private Thread Thread { get; } = Thread.CurrentThread;
+
+            public override bool IsForegroundThread => Thread.CurrentThread == Thread;
+        }
+
+        private class ThrowingTaskScheduler : TaskScheduler
+        {
+            protected override IEnumerable<Task> GetScheduledTasks()
+            {
+                return Enumerable.Empty<Task>();
+            }
+
+            protected override void QueueTask(Task task)
+            {
+                throw new InvalidOperationException("Use [ForegroundFactAttribute]");
+            }
+
+            protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
+            {
+                throw new InvalidOperationException("Use [ForegroundFactAttribute]");
+            }
+        }
+    }
+}

test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Xunit/ForegroundFactAttribute.cs

@@ -0,0 +1,16 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using Xunit;
+using Xunit.Sdk;
+
+namespace Microsoft.VisualStudio.LanguageServices.Razor.Test
+{
+    // Similar to WpfFactAttribute https://github.com/xunit/samples.xunit/blob/969d9f7e887836f01a6c525324bf3db55658c28f/STAExamples/WpfFactAttribute.cs
+    [AttributeUsage(AttributeTargets.Method, AllowMultiple = false)]
+    [XunitTestCaseDiscoverer(nameof(ForegroundFactAttribute), nameof(Microsoft.VisualStudio.LanguageServices.Razor))]
+    internal class ForegroundFactAttribute : FactAttribute
+    {
+    }
+}

test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Xunit/ForegroundFactDiscoverer.cs

@@ -0,0 +1,117 @@
+// Copyright (c) .NET Foundation. 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.Collections.Generic;
+using System.ComponentModel;
+using System.Threading;
+using System.Threading.Tasks;
+using System.Windows.Threading;
+using Xunit.Abstractions;
+using Xunit.Sdk;
+
+namespace Xunit
+{
+    internal class ForegroundFactTestCase : LongLivedMarshalByRefObject, IXunitTestCase
+    {
+        private IXunitTestCase _inner;
+
+        [EditorBrowsable(EditorBrowsableState.Never)]
+        [Obsolete("Called by the de-serializer", error: true)]
+        public ForegroundFactTestCase()
+        {
+        }
+
+        public ForegroundFactTestCase(IXunitTestCase testCase)
+        {
+            _inner = testCase;
+        }
+
+        public string DisplayName => _inner.DisplayName;
+
+        public IMethodInfo Method => _inner.Method;
+
+        public string SkipReason => _inner.SkipReason;
+
+        public ISourceInformation SourceInformation
+        {
+            get => _inner.SourceInformation;
+            set => _inner.SourceInformation = value;
+        }
+
+        public ITestMethod TestMethod => _inner.TestMethod;
+
+        public object[] TestMethodArguments => _inner.TestMethodArguments;
+
+        public Dictionary<string, List<string>> Traits => _inner.Traits;
+
+        public string UniqueID => _inner.UniqueID;
+
+        public void Deserialize(IXunitSerializationInfo info)
+        {
+            _inner = info.GetValue<IXunitTestCase>("InnerTestCase");
+        }
+
+        public void Serialize(IXunitSerializationInfo info)
+        {
+            info.AddValue("InnerTestCase", _inner);
+        }
+
+        public Task<RunSummary> RunAsync(
+            IMessageSink diagnosticMessageSink,
+            IMessageBus messageBus,
+            object[] constructorArguments,
+            ExceptionAggregator aggregator,
+            CancellationTokenSource cancellationTokenSource)
+        {
+            var tcs = new TaskCompletionSource<RunSummary>();
+            var thread = new Thread(() =>
+            {
+                try
+                {
+                    SynchronizationContext.SetSynchronizationContext(new DispatcherSynchronizationContext());
+                    
+                    var worker = _inner.RunAsync(diagnosticMessageSink, messageBus, constructorArguments, aggregator, cancellationTokenSource);
+
+                    Exception caught = null;
+                    var frame = new DispatcherFrame();
+                    Task.Run(async () =>
+                    {
+                        try
+                        {
+                            await worker;
+                        }
+                        catch (Exception ex)
+                        {
+                            caught = ex;
+                        }
+                        finally
+                        {
+                            frame.Continue = false;
+                        }
+                    });
+
+                    Dispatcher.PushFrame(frame);
+                    
+                    if (caught == null)
+                    {
+                        tcs.SetException(caught);
+                    }
+                    else
+                    {
+                        tcs.SetResult(worker.Result);
+                    }
+                }
+                catch (Exception e)
+                {
+                    tcs.SetException(e);
+                }
+            });
+
+            thread.SetApartmentState(ApartmentState.STA);
+            thread.Start();
+
+            return tcs.Task;
+        }
+    }
+}