microsoft · CyrusNajmabadi · Jul 6, 2015 · Jun 18, 2015 · Jun 18, 2015 · Jun 18, 2015
diff --git a/src/services/services.ts b/src/services/services.ts
@@ -1618,7 +1618,6 @@ namespace ts {
     export class OperationCanceledException { }
 
     export class CancellationTokenObject {
-
         public static None: CancellationTokenObject = new CancellationTokenObject(null)
 
         constructor(private cancellationToken: CancellationToken) {
@@ -6064,6 +6063,26 @@ namespace ts {
             return convertClassifications(getEncodedSemanticClassifications(fileName, span));
         }
 
+        function checkForClassificationCancellation(kind: SyntaxKind) {
+            // We don't want to actually call back into our host on every node to find out if we've
+            // been canceled.  That would be an enormous amount of chattyness, along with the all
+            // the overhead of marshalling the data to/from the host.  So instead we pick a few
+            // reasonable node kinds to bother checking on.  These node kinds represent high level
+            // constructs that we would expect to see commonly, but just at a far less frequent 
+            // interval.
+            //
+            // For example, in checker.ts (around 750k) we only have around 600 of these constructs.
+            // That means we're calling back into the host around every 1.2k of the file we process.
+            // Lib.d.ts has similar numbers.
+            switch (kind) {
+                case SyntaxKind.ModuleDeclaration:
+                case SyntaxKind.ClassDeclaration:
+                case SyntaxKind.InterfaceDeclaration:
+                case SyntaxKind.FunctionDeclaration:
+                    cancellationToken.throwIfCancellationRequested();
+            }
+        }
+
         function getEncodedSemanticClassifications(fileName: string, span: TextSpan): Classifications {
             synchronizeHostData();
 
@@ -6131,7 +6150,10 @@ namespace ts {
             function processNode(node: Node) {
                 // Only walk into nodes that intersect the requested span.
                 if (node && textSpanIntersectsWith(span, node.getFullStart(), node.getFullWidth())) {
-                    if (node.kind === SyntaxKind.Identifier && !nodeIsMissing(node)) {
+                    let kind = node.kind;
+                    checkForClassificationCancellation(kind);
+
+                    if (kind === SyntaxKind.Identifier && !nodeIsMissing(node)) {
                         let identifier = <Identifier>node;
 
                         // Only bother calling into the typechecker if this is an identifier that
@@ -6498,6 +6520,8 @@ namespace ts {
 
                 // Ignore nodes that don't intersect the original span to classify.
                 if (decodedTextSpanIntersectsWith(spanStart, spanLength, element.pos, element.getFullWidth())) {
+                    checkForClassificationCancellation(element.kind);
+
                     let children = element.getChildren(sourceFile);
                     for (let i = 0, n = children.length; i < n; i++) {
                         let child = children[i];

diff --git a/src/services/shims.ts b/src/services/shims.ts
@@ -327,7 +327,8 @@ namespace ts {
         }
 
         public getCancellationToken(): CancellationToken {
-            return this.shimHost.getCancellationToken();
+            var hostCancellationToken = this.shimHost.getCancellationToken();
+            return new ThrottledCancellationToken(hostCancellationToken);
         }
 
         public getCurrentDirectory(): string {
@@ -346,6 +347,29 @@ namespace ts {
         }
     }
 
+    /** A cancellation that throttles calls to the host */
+    class ThrottledCancellationToken implements CancellationToken {
+        // Store when we last tried to cancel.  Checking cancellation can be expensive (as we have
+        // to marshall over to the host layer).  So we only bother actually checking once enough
+        // time has passed.
+        private lastCancellationCheckTime = 0;
+
+        constructor(private hostCancellationToken: CancellationToken) {
+        }
+
+        public isCancellationRequested(): boolean {
+            var time = Date.now();
+            var duration = Math.abs(time - this.lastCancellationCheckTime);
+            if (duration > 10) {
+                // Check no more than once every 10 ms.
+                this.lastCancellationCheckTime = time;
+                return this.hostCancellationToken.isCancellationRequested();
+            }
+
+            return false;
+        }
+    }
+
     export class CoreServicesShimHostAdapter implements ParseConfigHost {
 
         constructor(private shimHost: CoreServicesShimHost) {

diff --git a/tests/cases/fourslash/semanticClassificationsCancellation1.ts b/tests/cases/fourslash/semanticClassificationsCancellation1.ts
@@ -0,0 +1,15 @@
+/// <reference path="fourslash.ts" />
+
+////module M {
+////}
+////module N {
+////}
+
+var c = classification;
+cancellation.setCancelled(1);
+verifyOperationIsCancelled(() => verify.semanticClassificationsAre());
+cancellation.resetCancelled();
+
+verify.semanticClassificationsAre(
+    c.moduleName("M"),
+    c.moduleName("N"));
diff --git a/tests/cases/fourslash/shims/cancellationWhenfindingAllRefsOnDefinition.ts b/tests/cases/fourslash/shims/cancellationWhenfindingAllRefsOnDefinition.ts
diff --git a/tests/cases/fourslash/syntacticClassificationsCancellation1.ts b/tests/cases/fourslash/syntacticClassificationsCancellation1.ts
@@ -0,0 +1,21 @@
+/// <reference path="fourslash.ts" />
+
+////module M {
+////}
+////module N {
+////}
+
+var c = classification;
+cancellation.setCancelled(1);
+verifyOperationIsCancelled(() => verify.syntacticClassificationsAre());
+cancellation.resetCancelled();
+
+verify.syntacticClassificationsAre(
+    c.keyword("module"),
+    c.moduleName("M"),
+    c.punctuation("{"),
+    c.punctuation("}"),
+    c.keyword("module"),
+    c.moduleName("N"),
+    c.punctuation("{"),
+    c.punctuation("}"));