Unify the behavior of AsyncBatchingWorkQueue's Dispose() and clean up uses

src/Dependencies/Threading/AsyncBatchingWorkQueue`0.cs

@@ -18,7 +18,7 @@ internal sealed class AsyncBatchingWorkQueue(
     TimeSpan delay,
     Func<CancellationToken, ValueTask> processBatchAsync,
     IAsynchronousOperationListener asyncListener,
-    CancellationToken cancellationToken) : AsyncBatchingWorkQueue<VoidResult>(delay, Convert(processBatchAsync), EqualityComparer<VoidResult>.Default, asyncListener, cancellationToken)
+    CancellationToken cancellationToken = default) : AsyncBatchingWorkQueue<VoidResult>(delay, Convert(processBatchAsync), EqualityComparer<VoidResult>.Default, asyncListener, cancellationToken)
 {
     private static Func<ImmutableSegmentedList<VoidResult>, CancellationToken, ValueTask> Convert(Func<CancellationToken, ValueTask> processBatchAsync)
         => (items, ct) => processBatchAsync(ct);

src/Dependencies/Threading/AsyncBatchingWorkQueue`1.cs

@@ -19,13 +19,13 @@ internal class AsyncBatchingWorkQueue<TItem>(
     Func<ImmutableSegmentedList<TItem>, CancellationToken, ValueTask> processBatchAsync,
     IEqualityComparer<TItem>? equalityComparer,
     IAsynchronousOperationListener asyncListener,
-    CancellationToken cancellationToken) : AsyncBatchingWorkQueue<TItem, VoidResult>(delay, Convert(processBatchAsync), equalityComparer, asyncListener, cancellationToken)
+    CancellationToken cancellationToken = default) : AsyncBatchingWorkQueue<TItem, VoidResult>(delay, Convert(processBatchAsync), equalityComparer, asyncListener, cancellationToken)
 {
     public AsyncBatchingWorkQueue(
         TimeSpan delay,
         Func<ImmutableSegmentedList<TItem>, CancellationToken, ValueTask> processBatchAsync,
         IAsynchronousOperationListener asyncListener,
-        CancellationToken cancellationToken)
+        CancellationToken cancellationToken = default)
         : this(delay,
                processBatchAsync,
                equalityComparer: null,

src/Dependencies/Threading/AsyncBatchingWorkQueue`2.cs

@@ -21,9 +21,9 @@ namespace Microsoft.CodeAnalysis.Threading;
 /// processing happen serially, only starting up after a previous round has completed.
 /// <para>
 /// Failure to complete a particular batch (either due to cancellation or some faulting error) will not prevent
-/// further batches from executing. The only thing that will permenantly stop this queue from processing items is if
-/// the <see cref="CancellationToken"/> passed to the constructor switches to <see
-/// cref="CancellationToken.IsCancellationRequested"/>.
+/// further batches from executing. The only thing that will permanently stop this queue from processing items is
+/// calling <see cref="Dispose()"/>, or cancelling the <see cref="CancellationToken"/> passed to the constructor,
+/// which is equivalent.
 /// </para>
 /// </summary>
 internal class AsyncBatchingWorkQueue<TItem, TResult> : IDisposable
@@ -46,9 +46,10 @@ internal class AsyncBatchingWorkQueue<TItem, TResult> : IDisposable
 
     /// <summary>
     /// Cancellation token controlling the entire queue.  Once this is triggered, we don't want to do any more work
-    /// at all.
+    /// at all. This is cancelled by a call to <see cref="Dispose()"/>; the IsCancellationRequested flag of this token
+    /// can be used as the "is disposed" flag for this object.
     /// </summary>
-    private readonly CancellationToken _entireQueueCancellationToken;
+    private readonly CancellationTokenSource _entireQueueCancellationTokenSource;
 
     /// <summary>
     /// Cancellation series we use so we can cancel individual batches of work if requested.  The client of the
@@ -59,6 +60,11 @@ internal class AsyncBatchingWorkQueue<TItem, TResult> : IDisposable
     /// </summary>
     private readonly CancellationSeries _cancellationSeries;
 
+    /// <summary>
+    /// If our constructor was given a CancellationToken, the registration against that token to call Dispose().
+    /// </summary>
+    private readonly CancellationTokenRegistration _externalCancellationTokenRegistration;
+
     #region protected by lock
 
     /// <summary>
@@ -107,23 +113,44 @@ public AsyncBatchingWorkQueue(
         Func<ImmutableSegmentedList<TItem>, CancellationToken, ValueTask<TResult>> processBatchAsync,
         IEqualityComparer<TItem>? equalityComparer,
         IAsynchronousOperationListener asyncListener,
-        CancellationToken cancellationToken)
+        CancellationToken cancellationToken = default)
     {
         _delay = delay;
         _processBatchAsync = processBatchAsync;
         _equalityComparer = equalityComparer;
         _asyncListener = asyncListener;
-        _entireQueueCancellationToken = cancellationToken;
+        _entireQueueCancellationTokenSource = new CancellationTokenSource();
 
         _uniqueItems = new SegmentedHashSet<TItem>(equalityComparer);
 
         // Combine with the queue cancellation token so that any batch is controlled by that token as well.
-        _cancellationSeries = new CancellationSeries(_entireQueueCancellationToken);
+        _cancellationSeries = new CancellationSeries(_entireQueueCancellationTokenSource.Token);
         CancelExistingWork();
+
+        // As a convenience, if we were given a cancellation token, this should be equivalent to calling Dispose().
+        // We don't link _entireQueueCancellationTokenSource to this, since we want to ensure the Dispose() also cleans up any
+        // queued items that were in our lists.
+        _externalCancellationTokenRegistration = cancellationToken.Register(static @this => ((AsyncBatchingWorkQueue<TItem, TResult>)@this!).Dispose(), this);
     }
 
     public void Dispose()
     {
+        lock (_gate)
+        {
+            // If we've previously disposed, we don't need to do anything further
+            if (_entireQueueCancellationTokenSource.IsCancellationRequested)
+                return;
+
+            // Cancel all work in the queue; this .Cancel() should stop the work, but we'll call CancelExistingWork() too to ensure
+            // we've cleared out all items that haven't ran.
+            CancelExistingWork();
+            _entireQueueCancellationTokenSource.Cancel();
+        }
+
+        // This must be done outside of the lock: disposing a registration blocks if the registered callback is currently running.
+        // If we did this inside the lock, the callback might be blocked waiting for a call to Dispose() to release the lock, but the
+        // caller of Dispose() would be blocked on that registration. If we could drop netstandard support, we could just call Unregister() instead.
+        _externalCancellationTokenRegistration.Dispose();
         _cancellationSeries.Dispose();
     }
 
@@ -135,6 +162,10 @@ public void CancelExistingWork()
     {
         lock (_gate)
         {
+            // If we've previously disposed, we don't need to do anything further
+            if (_entireQueueCancellationTokenSource.IsCancellationRequested)
+                return;
+
             // Cancel out the current executing batch, and create a new token for the next batch.
             _nextBatchCancellationToken = _cancellationSeries.CreateNext();
 
@@ -151,12 +182,12 @@ public void AddWork(TItem item, bool cancelExistingWork = false)
 
     public void AddWork(ReadOnlySpan<TItem> items, bool cancelExistingWork = false)
     {
-        // Don't do any more work if we've been asked to shutdown.
-        if (_entireQueueCancellationToken.IsCancellationRequested)
-            return;
-
         lock (_gate)
         {
+            // Don't do any more work if we've been asked to shutdown.
+            if (_entireQueueCancellationTokenSource.IsCancellationRequested)
+                return;
+
             // if we were asked to cancel the prior set of items, do so now.
             if (cancelExistingWork)
                 CancelExistingWork();
@@ -203,18 +234,18 @@ void AddItemsToBatch(ReadOnlySpan<TItem> items)
             await lastTask.NoThrowAwaitableInternal(captureContext: false);
 
             // If we were asked to shutdown, immediately transition to the canceled state without doing any more work.
-            if (_entireQueueCancellationToken.IsCancellationRequested)
+            if (_entireQueueCancellationTokenSource.IsCancellationRequested)
                 return (ranToCompletion: false, default(TResult?));
 
             // Ensure that we always yield the current thread this is necessary for correctness as we are called
             // inside a lock that _taskInFlight to true.  We must ensure that the work to process the next batch
             // must be on another thread that runs afterwards, can only grab the thread once we release it and will
             // then reset that bool back to false
             await Task.Yield().ConfigureAwait(false);
-            await _asyncListener.Delay(_delay, _entireQueueCancellationToken).NoThrowAwaitableInternal(false);
+            await _asyncListener.Delay(_delay, _entireQueueCancellationTokenSource.Token).NoThrowAwaitableInternal(false);
 
             // If we were asked to shutdown, immediately transition to the canceled state without doing any more work.
-            if (_entireQueueCancellationToken.IsCancellationRequested)
+            if (_entireQueueCancellationTokenSource.IsCancellationRequested)
                 return (ranToCompletion: false, default(TResult?));
 
             return (ranToCompletion: true, await ProcessNextBatchAsync().ConfigureAwait(false));
@@ -237,16 +268,16 @@ void AddItemsToBatch(ReadOnlySpan<TItem> items)
         var (ranToCompletion, result) = await updateTask.ConfigureAwait(false);
         if (!ranToCompletion)
         {
-            Debug.Assert(_entireQueueCancellationToken.IsCancellationRequested);
-            _entireQueueCancellationToken.ThrowIfCancellationRequested();
+            Debug.Assert(_entireQueueCancellationTokenSource.IsCancellationRequested);
+            _entireQueueCancellationTokenSource.Token.ThrowIfCancellationRequested();
         }
 
         return result;
     }
 
     private async ValueTask<TResult?> ProcessNextBatchAsync()
     {
-        _entireQueueCancellationToken.ThrowIfCancellationRequested();
+        _entireQueueCancellationTokenSource.Token.ThrowIfCancellationRequested();
         try
         {
             var (nextBatch, batchCancellationToken) = GetNextBatchAndResetQueue();
@@ -259,7 +290,7 @@ void AddItemsToBatch(ReadOnlySpan<TItem> items)
             await batchResultTask.NoThrowAwaitableInternal(false);
             if (batchResultTask.IsCompletedSuccessfully)
                 return batchResultTask.Result;
-            else if (batchResultTask.IsCanceled && !_entireQueueCancellationToken.IsCancellationRequested)
+            else if (batchResultTask.IsCanceled && !_entireQueueCancellationTokenSource.IsCancellationRequested)
             {
                 // Don't bubble up cancellation to the queue for the nested batch cancellation.  Just because we decided
                 // to cancel this batch isn't something that should stop processing further batches.

src/EditorFeatures/Core/Classification/Syntactic/SyntacticClassificationTaggerProvider.TagComputer.cs

@@ -45,8 +45,6 @@ private sealed record CachedServices(
         private readonly ITextBuffer _subjectBuffer;
         private readonly ITaggerEventSource _taggerEventSource;
 
-        private readonly CancellationTokenSource _disposalCancellationSource = new();
-
         /// <summary>
         /// Work queue we use to batch up notifications about changes that will cause
         /// us to classify.  This ensures that if we hear a flurry of changes, we don't
@@ -97,8 +95,7 @@ public TagComputer(
                 DelayTimeSpan.NearImmediate,
                 ProcessChangesAsync,
                 equalityComparer: null,
-                taggerProvider._listener,
-                _disposalCancellationSource.Token);
+                taggerProvider._listener);
 
             _lineCache = new ClassifiedLineCache(taggerProvider.ThreadingContext);
 
@@ -170,7 +167,7 @@ internal void DecrementReferenceCount()
             if (_taggerReferenceCount == 0)
             {
                 // stop any bg work we're doing.
-                _disposalCancellationSource.Cancel();
+                _workQueue.Dispose();
 
                 _taggerEventSource.Changed -= OnEventSourceChanged;
                 _taggerEventSource.Disconnect();

src/EditorFeatures/Core/CodeDefinitionWindow/DefinitionContextTracker.cs

@@ -37,7 +37,7 @@ namespace Microsoft.CodeAnalysis.CodeDefinitionWindow;
 [Export(typeof(DefinitionContextTracker))]
 [ContentType(ContentTypeNames.RoslynContentType)]
 [TextViewRole(PredefinedTextViewRoles.Interactive)]
-internal sealed class DefinitionContextTracker : ITextViewConnectionListener
+internal sealed class DefinitionContextTracker : ITextViewConnectionListener, IDisposable
 {
     private readonly HashSet<ITextView> _subscribedViews = [];
     private readonly IMetadataAsSourceFileService _metadataAsSourceFileService;
@@ -66,10 +66,11 @@ public DefinitionContextTracker(
         _workQueue = new AsyncBatchingWorkQueue<SnapshotPoint>(
             DelayTimeSpan.Short,
             ProcessWorkAsync,
-            _asyncListener,
-            _threadingContext.DisposalToken);
+            _asyncListener);
     }
 
+    public void Dispose() => _workQueue.Dispose();
+
     void ITextViewConnectionListener.SubjectBuffersConnected(ITextView textView, ConnectionReason reason, IReadOnlyCollection<ITextBuffer> subjectBuffers)
     {
         Contract.ThrowIfFalse(_threadingContext.JoinableTaskContext.IsOnMainThread);

src/EditorFeatures/Core/Copilot/CopilotWpfTextCreationListener.cs

@@ -10,7 +10,6 @@
 using Microsoft.CodeAnalysis.Collections;
 using Microsoft.CodeAnalysis.Editor;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
-using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.CodeAnalysis.Remote;
@@ -27,10 +26,9 @@ namespace Microsoft.CodeAnalysis.Copilot;
 [Export(typeof(IWpfTextViewCreationListener))]
 [ContentType(ContentTypeNames.RoslynContentType)]
 [TextViewRole(PredefinedTextViewRoles.Document)]
-internal sealed class CopilotWpfTextViewCreationListener : IWpfTextViewCreationListener
+internal sealed class CopilotWpfTextViewCreationListener : IWpfTextViewCreationListener, IDisposable
 {
     private readonly IGlobalOptionService _globalOptions;
-    private readonly IThreadingContext _threadingContext;
     private readonly Lazy<SuggestionServiceBase> _suggestionServiceBase;
     private readonly IAsynchronousOperationListener _listener;
 
@@ -42,22 +40,21 @@ internal sealed class CopilotWpfTextViewCreationListener : IWpfTextViewCreationL
     [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
     public CopilotWpfTextViewCreationListener(
         IGlobalOptionService globalOptions,
-        IThreadingContext threadingContext,
         Lazy<SuggestionServiceBase> suggestionServiceBase,
         IAsynchronousOperationListenerProvider listenerProvider)
     {
         _globalOptions = globalOptions;
-        _threadingContext = threadingContext;
         _suggestionServiceBase = suggestionServiceBase;
         _listener = listenerProvider.GetListener(FeatureAttribute.CopilotChangeAnalysis);
 
         _completionWorkQueue = new AsyncBatchingWorkQueue<(bool accepted, ProposalBase proposal)>(
             DelayTimeSpan.Idle,
             ProcessCompletionEventsAsync,
-            _listener,
-            _threadingContext.DisposalToken);
+            _listener);
     }
 
+    public void Dispose() => _completionWorkQueue.Dispose();
+
     public void TextViewCreated(IWpfTextView textView)
     {
         // On the first roslyn text view created, kick off work to hydrate the suggestion service and register to events

src/EditorFeatures/Core/Remote/SolutionChecksumUpdater.cs

@@ -19,7 +19,7 @@ namespace Microsoft.CodeAnalysis.Remote;
 /// This class runs against the in-process workspace, and when it sees changes proactively pushes them to
 /// the out-of-process workspace through the <see cref="IRemoteAssetSynchronizationService"/>.
 /// </summary>
-internal sealed class SolutionChecksumUpdater
+internal sealed class SolutionChecksumUpdater : IDisposable
 {
     private readonly Workspace _workspace;
 
@@ -66,14 +66,12 @@ public SolutionChecksumUpdater(
         _synchronizeWorkspaceQueue = new AsyncBatchingWorkQueue(
             DelayTimeSpan.Short,
             SynchronizePrimaryWorkspaceAsync,
-            listener,
-            shutdownToken);
+            listener);
 
         _synchronizeActiveDocumentQueue = new AsyncBatchingWorkQueue(
             TimeSpan.Zero,
             SynchronizeActiveDocumentAsync,
-            listener,
-            shutdownToken);
+            listener);
 
         // start listening workspace change event
         _workspaceChangedDisposer = _workspace.RegisterWorkspaceChangedHandler(this.OnWorkspaceChanged);
@@ -85,12 +83,13 @@ public SolutionChecksumUpdater(
         _synchronizeWorkspaceQueue.AddWork();
     }
 
-    public void Shutdown()
+    public void Dispose()
     {
-        // Try to stop any work that is in progress.
+        // Stop any work that is in progress, and prevent any further work from being queued up.
         lock (_gate)
         {
-            _synchronizeWorkspaceQueue.CancelExistingWork();
+            _synchronizeWorkspaceQueue.Dispose();
+            _synchronizeActiveDocumentQueue.Dispose();
         }
 
         _documentTrackingService.ActiveDocumentChanged -= OnActiveDocumentChanged;