Properly wait for last batch of results to be reported

src/Features/LanguageServer/Protocol/CustomProtocol/FindUsagesLSPContext.cs

@@ -71,18 +71,14 @@ internal class FindUsagesLSPContext : FindUsagesContext
             _position = position;
             _metadataAsSourceFileService = metadataAsSourceFileService;
             _workQueue = new AsyncBatchingWorkQueue<VSReferenceItem>(
-                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<VSReferenceItem> referencesToReport, CancellationToken cancellationToken)
+        private Task ReportReferencesAsync(ImmutableArray<VSReferenceItem> 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;

src/Workspaces/Core/Portable/Shared/Utilities/AsyncBatchingWorkQueue.cs

@@ -128,6 +128,8 @@ public void AddWork(IEnumerable<TItem> items)
             }
         }
 
+        public async Task WaitUntilCurrentBatchCompletesAsync() => await _updateTask.ConfigureAwait(false);
+
         private void AddItemsToBatch(IEnumerable<TItem> items)
         {
             // no equality comparer.  We want to process all items.