Add support for maxItemsComputed

README.md

@@ -39,6 +39,7 @@ The following settings are supported:
 - `yaml.schemas`: Helps you associate schemas with files in a glob pattern
 - `yaml.schemaStore.enable`: When set to true the YAML language server will pull in all available schemas from [JSON Schema Store](https://www.schemastore.org/json/)
 - `yaml.customTags`: Array of custom tags that the parser will validate against. It has two ways to be used. Either an item in the array is a custom tag such as "!Ref" and it will automatically map !Ref to scalar or you can specify the type of the object !Ref should be e.g. "!Ref sequence". The type of object can be either scalar (for strings and booleans), sequence (for arrays), map (for objects).
+- `yaml.maxItemsComputed`: The maximum number of outline symbols and folding regions computed (limited for performance reasons).
 - `[yaml].editor.tabSize`: the number of spaces to use when autocompleting. Takes priority over editor.tabSize.
 - `editor.tabSize`: the number of spaces to use when autocompleting. Default is 2.
 - `http.proxy`: The URL of the proxy server that will be used when attempting to download a schema. If it is not set or it is undefined no proxy server will be used.

src/languageserver/handlers/languageHandlers.ts

@@ -23,12 +23,16 @@ import { isKubernetesAssociatedDocument } from '../../languageservice/parser/isK
 import { LanguageService } from '../../languageservice/yamlLanguageService';
 import { SettingsState } from '../../yamlSettings';
 import { ValidationHandler } from './validationHandlers';
+import { ResultLimitReachedNotification } from '../../requestTypes';
+import * as path from 'path';
 
 export class LanguageHandlers {
   private languageService: LanguageService;
   private yamlSettings: SettingsState;
   private validationHandler: ValidationHandler;
 
+  pendingLimitExceededWarnings: { [uri: string]: { features: { [name: string]: string }; timeout?: NodeJS.Timeout } };
+
   constructor(
     private readonly connection: Connection,
     languageService: LanguageService,
@@ -38,6 +42,7 @@ export class LanguageHandlers {
     this.languageService = languageService;
     this.yamlSettings = yamlSettings;
     this.validationHandler = validationHandler;
+    this.pendingLimitExceededWarnings = {};
   }
 
   public registerHandlers(): void {
@@ -52,6 +57,9 @@ export class LanguageHandlers {
     this.connection.onDocumentOnTypeFormatting((params) => this.formatOnTypeHandler(params));
     this.connection.onCodeLens((params) => this.codeLensHandler(params));
     this.connection.onCodeLensResolve((params) => this.codeLensResolveHandler(params));
+
+    this.yamlSettings.documents.onDidChangeContent((change) => this.cancelLimitExceededWarnings(change.document.uri));
+    this.yamlSettings.documents.onDidClose((event) => this.cancelLimitExceededWarnings(event.document.uri));
   }
 
   documentLinkHandler(params: DocumentLinkParams): Promise<DocumentLink[]> {
@@ -74,10 +82,18 @@ export class LanguageHandlers {
       return;
     }
 
+    const onResultLimitExceeded = this.onResultLimitExceeded(
+      document.uri,
+      this.yamlSettings.maxItemsComputed,
+      'document symbols'
+    );
+
+    const context = { resultLimit: this.yamlSettings.maxItemsComputed, onResultLimitExceeded };
+
     if (this.yamlSettings.hierarchicalDocumentSymbolSupport) {
-      return this.languageService.findDocumentSymbols2(document);
+      return this.languageService.findDocumentSymbols2(document, context);
     } else {
-      return this.languageService.findDocumentSymbols(document);
+      return this.languageService.findDocumentSymbols(document, context);
     }
   }
 
@@ -169,7 +185,17 @@ export class LanguageHandlers {
       return;
     }
 
-    return this.languageService.getFoldingRanges(textDocument, this.yamlSettings.capabilities.textDocument.foldingRange);
+    const capabilities = this.yamlSettings.capabilities.textDocument.foldingRange;
+    const rangeLimit = this.yamlSettings.maxItemsComputed || capabilities.rangeLimit;
+    const onRangeLimitExceeded = this.onResultLimitExceeded(textDocument.uri, rangeLimit, 'folding ranges');
+
+    const context = {
+      rangeLimit,
+      onRangeLimitExceeded,
+      lineFoldingOnly: capabilities.lineFoldingOnly,
+    };
+
+    return this.languageService.getFoldingRanges(textDocument, context);
   }
 
   codeActionHandler(params: CodeActionParams): CodeAction[] | undefined {
@@ -192,4 +218,40 @@ export class LanguageHandlers {
   codeLensResolveHandler(param: CodeLens): Thenable<CodeLens> | CodeLens {
     return this.languageService.resolveCodeLens(param);
   }
+
+  // Adapted from:
+  // https://github.com/microsoft/vscode/blob/94c9ea46838a9a619aeafb7e8afd1170c967bb55/extensions/json-language-features/server/src/jsonServer.ts#L172
+  private cancelLimitExceededWarnings(uri: string): void {
+    const warning = this.pendingLimitExceededWarnings[uri];
+    if (warning && warning.timeout) {
+      clearTimeout(warning.timeout);
+      delete this.pendingLimitExceededWarnings[uri];
+    }
+  }
+
+  private onResultLimitExceeded(uri: string, resultLimit: number, name: string) {
+    return () => {
+      let warning = this.pendingLimitExceededWarnings[uri];
+      if (warning) {
+        if (!warning.timeout) {
+          // already shown
+          return;
+        }
+        warning.features[name] = name;
+        warning.timeout.refresh();
+      } else {
+        warning = { features: { [name]: name } };
+        warning.timeout = setTimeout(() => {
+          this.connection.sendNotification(
+            ResultLimitReachedNotification.type,
+            `${path.basename(uri)}: For performance reasons, ${Object.keys(warning.features).join(
+              ' and '
+            )} have been limited to ${resultLimit} items.`
+          );
+          warning.timeout = undefined;
+        }, 2000);
+        this.pendingLimitExceededWarnings[uri] = warning;
+      }
+    };
+  }
 }

src/languageserver/handlers/settingsHandlers.ts

@@ -47,6 +47,8 @@ export class SettingsHandler {
       }
       this.yamlSettings.customTags = settings.yaml.customTags ? settings.yaml.customTags : [];
 
+      this.yamlSettings.maxItemsComputed = Math.trunc(Math.max(0, Number(settings.yaml.maxItemsComputed))) || 5000;
+
       if (settings.yaml.schemaStore) {
         this.yamlSettings.schemaStoreEnabled = settings.yaml.schemaStore.enable;
       }

src/languageservice/services/documentSymbols.ts

@@ -8,6 +8,7 @@
 import { SymbolInformation, DocumentSymbol } from 'vscode-languageserver-types';
 import { YAMLSchemaService } from './yamlSchemaService';
 import { JSONDocumentSymbols } from 'vscode-json-languageservice/lib/umd/services/jsonDocumentSymbols';
+import { DocumentSymbolsContext } from 'vscode-json-languageservice/lib/umd/jsonLanguageTypes';
 import { TextDocument } from 'vscode-languageserver-textdocument';
 import { yamlDocumentsCache } from '../parser/yaml-documents';
 
@@ -29,7 +30,10 @@ export class YAMLDocumentSymbols {
     };
   }
 
-  public findDocumentSymbols(document: TextDocument): SymbolInformation[] {
+  public findDocumentSymbols(
+    document: TextDocument,
+    context: DocumentSymbolsContext = { resultLimit: Number.MAX_VALUE }
+  ): SymbolInformation[] {
     const doc = yamlDocumentsCache.getYamlDocument(document);
     if (!doc || doc['documents'].length === 0) {
       return null;
@@ -38,14 +42,17 @@ export class YAMLDocumentSymbols {
     let results = [];
     for (const yamlDoc of doc['documents']) {
       if (yamlDoc.root) {
-        results = results.concat(this.jsonDocumentSymbols.findDocumentSymbols(document, yamlDoc));
+        results = results.concat(this.jsonDocumentSymbols.findDocumentSymbols(document, yamlDoc, context));
       }
     }
 
     return results;
   }
 
-  public findHierarchicalDocumentSymbols(document: TextDocument): DocumentSymbol[] {
+  public findHierarchicalDocumentSymbols(
+    document: TextDocument,
+    context: DocumentSymbolsContext = { resultLimit: Number.MAX_VALUE }
+  ): DocumentSymbol[] {
     const doc = yamlDocumentsCache.getYamlDocument(document);
     if (!doc || doc['documents'].length === 0) {
       return null;
@@ -54,7 +61,7 @@ export class YAMLDocumentSymbols {
     let results = [];
     for (const yamlDoc of doc['documents']) {
       if (yamlDoc.root) {
-        results = results.concat(this.jsonDocumentSymbols.findDocumentSymbols2(document, yamlDoc));
+        results = results.concat(this.jsonDocumentSymbols.findDocumentSymbols2(document, yamlDoc, context));
       }
     }
 

src/languageservice/services/yamlFolding.ts

@@ -34,8 +34,11 @@ export function getFoldingRanges(document: TextDocument, context: FoldingRangesC
   if (typeof rangeLimit !== 'number' || result.length <= rangeLimit) {
     return result;
   }
+  if (context && context.onRangeLimitExceeded) {
+    context.onRangeLimitExceeded(document.uri);
+  }
 
-  return result.slice(0, context.rangeLimit - 1);
+  return result.slice(0, context.rangeLimit);
 }
 
 function creteNormalizedFolding(document: TextDocument, node: ASTNode): FoldingRange {

src/languageservice/yamlLanguageService.ts

@@ -28,7 +28,7 @@ import { YAMLCompletion } from './services/yamlCompletion';
 import { YAMLHover } from './services/yamlHover';
 import { YAMLValidation } from './services/yamlValidation';
 import { YAMLFormatter } from './services/yamlFormatter';
-import { JSONDocument, DefinitionLink, TextDocument } from 'vscode-json-languageservice';
+import { JSONDocument, DefinitionLink, TextDocument, DocumentSymbolsContext } from 'vscode-json-languageservice';
 import { findLinks } from './services/yamlLinks';
 import {
   FoldingRange,
@@ -117,8 +117,8 @@ export interface LanguageService {
   doComplete(document: TextDocument, position: Position, isKubernetes: boolean): Promise<CompletionList>;
   doValidation(document: TextDocument, isKubernetes: boolean): Promise<Diagnostic[]>;
   doHover(document: TextDocument, position: Position): Promise<Hover | null>;
-  findDocumentSymbols(document: TextDocument): SymbolInformation[];
-  findDocumentSymbols2(document: TextDocument): DocumentSymbol[];
+  findDocumentSymbols(document: TextDocument, context: DocumentSymbolsContext): SymbolInformation[];
+  findDocumentSymbols2(document: TextDocument, context: DocumentSymbolsContext): DocumentSymbol[];
   findDefinition(document: TextDocument, position: Position, doc: JSONDocument): Promise<DefinitionLink[]>;
   findLinks(document: TextDocument): Promise<DocumentLink[]>;
   resetSchema(uri: string): boolean;

src/languageservice/yamlTypes.ts

@@ -8,6 +8,10 @@ export interface FoldingRangesContext {
    * The maximal number of ranges returned.
    */
   rangeLimit?: number;
+  /**
+   * Called when the result was cropped.
+   */
+  onRangeLimitExceeded?: (uri: string) => void;
   /**
    * If set, the client signals that it only supports folding complete lines. If set, client will
    * ignore specified `startCharacter` and `endCharacter` properties in a FoldingRange.

src/requestTypes.ts

@@ -21,6 +21,10 @@ export namespace VSCodeContentRequestRegistration {
   export const type: NotificationType<{}> = new NotificationType('yaml/registerContentRequest');
 }
 
+export namespace ResultLimitReachedNotification {
+  export const type: NotificationType<string> = new NotificationType('yaml/resultLimitReached');
+}
+
 export namespace VSCodeContentRequest {
   export const type: RequestType<{}, {}, {}> = new RequestType('vscode/content');
 }

src/yamlSettings.ts

@@ -17,6 +17,7 @@ export interface Settings {
     schemaStore: {
       enable: boolean;
     };
+    maxItemsComputed: number;
   };
   http: {
     proxy: string;
@@ -54,6 +55,7 @@ export class SettingsState {
   customTags = [];
   schemaStoreEnabled = true;
   indentation: string | undefined = undefined;
+  maxItemsComputed = 5000;
 
   // File validation helpers
   pendingValidationRequests: { [uri: string]: NodeJS.Timer } = {};

test/documentSymbols.test.ts

@@ -19,13 +19,29 @@ describe('Document Symbols Tests', () => {
   let languageHandler: LanguageHandlers;
   let yamlSettings: SettingsState;
 
+  const limitContent = `
+    a: [1, 2, 3]
+    b: [4, 5, 6]
+  `;
+
   before(() => {
     const languageSettingsSetup = new ServiceSetup();
     const { languageHandler: langHandler, yamlSettings: settings } = setupLanguageService(languageSettingsSetup.languageSettings);
     languageHandler = langHandler;
     yamlSettings = settings;
   });
 
+  afterEach(() => {
+    yamlSettings.maxItemsComputed = 5000;
+  });
+
+  function assertLimitWarning(): void {
+    const warnings = languageHandler.pendingLimitExceededWarnings;
+    assert.deepEqual(Object.keys(warnings), [TEST_URI]);
+    assert.deepEqual(warnings[TEST_URI].features, { 'document symbols': 'document symbols' });
+    assert(warnings[TEST_URI].timeout);
+  }
+
   describe('Document Symbols Tests (Non Hierarchical)', function () {
     function parseNonHierarchicalSetup(content: string): SymbolInformation[] | DocumentSymbol[] {
       const testTextDocument = setupTextDocument(content);
@@ -121,6 +137,16 @@ describe('Document Symbols Tests', () => {
       assert.deepEqual(symbols[0], createExpectedSymbolInformation('analytics', 17, '', TEST_URI, 1, 0, 1, 15));
       assert.deepEqual(symbols[1], createExpectedSymbolInformation('json', 15, '', TEST_URI, 4, 0, 4, 10));
     });
+
+    it('Document symbols with a limit', () => {
+      yamlSettings.maxItemsComputed = 1;
+
+      const symbols = parseNonHierarchicalSetup(limitContent);
+      assert.equal(symbols.length, 1);
+      assert.deepEqual(symbols[0], createExpectedSymbolInformation('a', SymbolKind.Array, '', TEST_URI, 1, 4, 1, 16));
+
+      assertLimitWarning();
+    });
   });
 
   describe('Document Symbols Tests (Hierarchical)', function () {
@@ -275,5 +301,21 @@ describe('Document Symbols Tests', () => {
         createExpectedDocumentSymbol('conditions', SymbolKind.Module, 6, 12, 10, 28, 6, 12, 6, 22, [root2])
       );
     });
+
+    it('Document symbols with a limit', () => {
+      yamlSettings.maxItemsComputed = 3;
+
+      const symbols = parseHierarchicalSetup(limitContent) as DocumentSymbol[];
+      assert.equal(symbols.length, 2);
+      assert.equal(symbols[0].children.length, 1);
+      assert.equal(symbols[1].children.length, 0);
+
+      const el = createExpectedDocumentSymbolNoDetail('0', SymbolKind.Number, 1, 8, 1, 9, 1, 8, 1, 9, []);
+      const root = createExpectedDocumentSymbol('a', SymbolKind.Array, 1, 4, 1, 16, 1, 4, 1, 5, [el]);
+
+      assert.deepEqual(symbols[0], root);
+
+      assertLimitWarning();
+    });
   });
 });

test/yamlFolding.test.ts

@@ -7,7 +7,7 @@ import { expect } from 'chai';
 import { FoldingRange } from 'vscode-languageserver';
 import { getFoldingRanges } from '../src/languageservice/services/yamlFolding';
 import { FoldingRangesContext } from '../src/languageservice/yamlTypes';
-import { setupTextDocument } from './utils/testHelper';
+import { setupTextDocument, TEST_URI } from './utils/testHelper';
 
 const context: FoldingRangesContext = { rangeLimit: 10_0000 };
 
@@ -100,4 +100,31 @@ SecondDict:
     const ranges = getFoldingRanges(doc, context);
     expect(ranges).to.deep.include.members([FoldingRange.create(1, 4, 2, 15)]);
   });
+
+  it('should respect range limits', () => {
+    const yaml = `
+      a:
+        - 1
+      b:
+        - 2
+    `;
+
+    const warnings = [];
+
+    const doc = setupTextDocument(yaml);
+
+    const unlimitedRanges = getFoldingRanges(doc, {
+      rangeLimit: 10,
+      onRangeLimitExceeded: (uri) => warnings.push(uri),
+    });
+    expect(unlimitedRanges.length).to.equal(2);
+    expect(warnings).to.be.empty;
+
+    const limitedRanges = getFoldingRanges(doc, {
+      rangeLimit: 1,
+      onRangeLimitExceeded: (uri) => warnings.push(uri),
+    });
+    expect(limitedRanges.length).to.equal(1);
+    expect(warnings).to.deep.equal([TEST_URI]);
+  });
 });