diff --git a/src/Features/LanguageServer/Protocol/CustomProtocol/FindUsagesLSPContext.cs b/src/Features/LanguageServer/Protocol/CustomProtocol/FindUsagesLSPContext.cs index bf6105e5ec5cb1c5dce293b19611f95749f2b874..bc9506aaeda4f849532bf1e80d5f5bd821d8f251 100644 --- a/src/Features/LanguageServer/Protocol/CustomProtocol/FindUsagesLSPContext.cs +++ b/src/Features/LanguageServer/Protocol/CustomProtocol/FindUsagesLSPContext.cs @@ -71,18 +71,14 @@ internal class FindUsagesLSPContext : FindUsagesContext _position = position; _metadataAsSourceFileService = metadataAsSourceFileService; _workQueue = new AsyncBatchingWorkQueue( - TimeSpan.FromSeconds(0.5), ReportIfNotEmptyAsync, cancellationToken); + TimeSpan.FromSeconds(0.5), ReportReferencesAsync, cancellationToken); CancellationToken = cancellationToken; } - public override Task OnCompletedAsync() - { - // Upon completion, we wait an additional 0.5s (the time in between batches) to ensure - // that all results have been reported. - Thread.Sleep(500); - return Task.CompletedTask; - } + // After all definitions/references have been found, wait here until all results have been reported. + public override async Task OnCompletedAsync() + => await _workQueue.WaitUntilCurrentBatchCompletesAsync().ConfigureAwait(false); public override async Task OnDefinitionFoundAsync(DefinitionItem definition) { @@ -271,13 +267,8 @@ public override async Task OnReferenceFoundAsync(SourceReferenceItem reference) } } - private Task ReportIfNotEmptyAsync(ImmutableArray referencesToReport, CancellationToken cancellationToken) + private Task ReportReferencesAsync(ImmutableArray referencesToReport, CancellationToken cancellationToken) { - if (referencesToReport.IsEmpty) - { - return Task.CompletedTask; - } - // We can report outside of the lock here since _progress is thread-safe. _progress.Report(referencesToReport.ToArray()); return Task.CompletedTask; diff --git a/src/Workspaces/Core/Portable/Shared/Utilities/AsyncBatchingWorkQueue.cs b/src/Workspaces/Core/Portable/Shared/Utilities/AsyncBatchingWorkQueue.cs index e8d54d5945b922f354e7484949a6a20cefc4ef1d..bc06aeedea0efae825337b11bb88949addf366e4 100644 --- a/src/Workspaces/Core/Portable/Shared/Utilities/AsyncBatchingWorkQueue.cs +++ b/src/Workspaces/Core/Portable/Shared/Utilities/AsyncBatchingWorkQueue.cs @@ -128,6 +128,8 @@ public void AddWork(IEnumerable items) } } + public async Task WaitUntilCurrentBatchCompletesAsync() => await _updateTask.ConfigureAwait(false); + private void AddItemsToBatch(IEnumerable items) { // no equality comparer. We want to process all items.