Add cancellation to HtmlCopy GetAllClassificationSpans

microsoft · Mar 5, 2021 · 1abfe8e · 1abfe8e
1 parent ee92b43
commit 1abfe8e
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/src/CopyAsHtml/CopyAsHtml/FormattedStringBuilder.cs b/src/CopyAsHtml/CopyAsHtml/FormattedStringBuilder.cs
@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
+using System.Threading;
 using CopyAsHtml;
 using Microsoft.VisualStudio.Language.Intellisense.Utilities;
 using Microsoft.VisualStudio.Text.Classification;
@@ -97,16 +98,24 @@ private void Append(string text)
             _stringBuilder.Append(text);
         }
 
-
         private IList<ClassificationSpan> GetClassificationSpansSync(SnapshotSpan parentSpan)
         {
             IAccurateClassifier accurateClassifier = _classifier as IAccurateClassifier;
 
             if (accurateClassifier != null && _waitIndicator != null)
             {
-                using (var waitContext = WaitHelper.Wait(_waitIndicator, "HTML Copy", "Formatting document for copying"))
+                try
+                {
+                    using (var waitContext = WaitHelper.Wait(_waitIndicator, "HTML Copy", "Formatting document for copying"))
+                    using (var timeoutSource = new CancellationTokenSource(TimeSpan.FromMilliseconds(500)))
+                    using (var cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(waitContext.CancellationToken, timeoutSource.Token))
+                    {
+                        return accurateClassifier.GetAllClassificationSpans(parentSpan, cancellationTokenSource.Token);
+                    }
+                }
+                catch
                 {
-                    return accurateClassifier.GetAllClassificationSpans(parentSpan, waitContext.CancellationToken);
+                    return Array.Empty<ClassificationSpan>();
                 }
             }