Merge ProjectSnapshotManagerDispatcher and ProjectSnapshotManagerDisp…

…atcherBase

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/LSPProjectSnapshotManagerDispatcher.cs

@@ -2,12 +2,11 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using Microsoft.CodeAnalysis.Razor;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer;
 
 internal class LSPProjectSnapshotManagerDispatcher(IErrorReporter errorReporter)
-    : ProjectSnapshotManagerDispatcherBase(ThreadName, errorReporter)
+    : ProjectSnapshotManagerDispatcher(ThreadName, errorReporter)
 {
     private const string ThreadName = "Razor." + nameof(LSPProjectSnapshotManagerDispatcher);
 }

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

@@ -2,17 +2,37 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace Microsoft.CodeAnalysis.Razor;
 
-internal abstract class ProjectSnapshotManagerDispatcher
+internal abstract class ProjectSnapshotManagerDispatcher : IDisposable
 {
-    public abstract bool IsDispatcherThread { get; }
+    private readonly ProjectSnapshotManagerTaskScheduler _dispatcherScheduler;
 
-    public abstract TaskScheduler DispatcherScheduler { get; }
+    public bool IsDispatcherThread
+        => Thread.CurrentThread.ManagedThreadId == _dispatcherScheduler.ThreadId;
+    public TaskScheduler DispatcherScheduler
+        => _dispatcherScheduler;
+
+    protected ProjectSnapshotManagerDispatcher(string threadName, IErrorReporter errorReporter)
+    {
+        if (threadName is null)
+        {
+            throw new ArgumentNullException(nameof(threadName));
+        }
+
+        _dispatcherScheduler = new ProjectSnapshotManagerTaskScheduler(threadName, errorReporter);
+    }
+
+    public void Dispose()
+    {
+        _dispatcherScheduler.Dispose();
+    }
 
     public Task RunOnDispatcherThreadAsync(Action action, CancellationToken cancellationToken)
         => Task.Factory.StartNew(action, cancellationToken, TaskCreationOptions.None, DispatcherScheduler);
@@ -31,4 +51,103 @@ public virtual void AssertDispatcherThread([CallerMemberName] string? caller = n
             throw new InvalidOperationException(caller + " must be called on the project snapshot manager's thread.");
         }
     }
+
+    private class ProjectSnapshotManagerTaskScheduler : TaskScheduler, IDisposable
+    {
+        private readonly Thread _thread;
+        private readonly BlockingCollection<Task> _tasks = new();
+        private readonly IErrorReporter _errorReporter;
+        private bool _disposed;
+        private readonly object _disposalLock = new();
+
+        public ProjectSnapshotManagerTaskScheduler(string threadName, IErrorReporter errorReporter)
+        {
+            _thread = new Thread(ThreadStart)
+            {
+                Name = threadName,
+                IsBackground = true,
+            };
+
+            _thread.Start();
+            _errorReporter = errorReporter;
+        }
+
+        public int ThreadId => _thread.ManagedThreadId;
+
+        public override int MaximumConcurrencyLevel => 1;
+
+        protected override void QueueTask(Task task)
+        {
+            lock (_disposalLock)
+            {
+                if (_disposed)
+                {
+                    return;
+                }
+            }
+
+            _tasks.Add(task);
+        }
+
+        protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
+        {
+            // If the task was previously queued it means that we're ensuring it's running on our single threaded scheduler.
+            // Otherwise, we can't enforce that behavior and therefore need it to be re-queued before execution.
+            if (taskWasPreviouslyQueued)
+            {
+                return TryExecuteTask(task);
+            }
+
+            return false;
+        }
+
+        protected override IEnumerable<Task> GetScheduledTasks() => _tasks.ToArray();
+
+        private void ThreadStart()
+        {
+            while (!_disposed)
+            {
+                try
+                {
+                    var task = _tasks.Take();
+                    TryExecuteTask(task);
+                }
+                catch (ThreadAbortException ex)
+                {
+                    // Fires when things shut down or in tests. Log exception and bail out.
+                    _errorReporter.ReportError(ex);
+                    return;
+                }
+                catch (Exception ex)
+                {
+                    lock (_disposalLock)
+                    {
+                        if (_disposed)
+                        {
+                            // Graceful teardown
+                            return;
+                        }
+                    }
+
+                    // Unexpected exception. Log and throw.
+                    _errorReporter.ReportError(ex);
+                    throw;
+                }
+            }
+        }
+
+        public void Dispose()
+        {
+            lock (_disposalLock)
+            {
+                if (_disposed)
+                {
+                    return;
+                }
+
+                _disposed = true;
+                _tasks.CompleteAdding();
+            }
+        }
+    }
 }

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

@@ -3,14 +3,13 @@
 
 using System.ComponentModel.Composition;
 using Microsoft.CodeAnalysis.Razor;
-using Microsoft.CodeAnalysis.Razor.Workspaces;
 
 namespace Microsoft.VisualStudio.LanguageServices.Razor;
 
 [Export(typeof(ProjectSnapshotManagerDispatcher))]
 [method: ImportingConstructor]
 internal class VisualStudioProjectSnapshotManagerDispatcher(IErrorReporter errorReporter)
-    : ProjectSnapshotManagerDispatcherBase(ThreadName, errorReporter)
+    : ProjectSnapshotManagerDispatcher(ThreadName, errorReporter)
 {
     private const string ThreadName = "Razor." + nameof(VisualStudioProjectSnapshotManagerDispatcher);
 }