Use limiter instead of custom batcher

extensions/markdown-language-features/src/util/limiter.ts

@@ -0,0 +1,67 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+interface ILimitedTaskFactory<T> {
+	factory: ITask<Promise<T>>;
+	c: (value: T | Promise<T>) => void;
+	e: (error?: unknown) => void;
+}
+
+interface ITask<T> {
+	(): T;
+}
+
+/**
+ * A helper to queue N promises and run them all with a max degree of parallelism. The helper
+ * ensures that at any time no more than M promises are running at the same time.
+ *
+ * Taken from 'src/vs/base/common/async.ts'
+ */
+export class Limiter<T> {
+
+	private _size = 0;
+	private runningPromises: number;
+	private readonly maxDegreeOfParalellism: number;
+	private readonly outstandingPromises: ILimitedTaskFactory<T>[];
+
+	constructor(maxDegreeOfParalellism: number) {
+		this.maxDegreeOfParalellism = maxDegreeOfParalellism;
+		this.outstandingPromises = [];
+		this.runningPromises = 0;
+	}
+
+	get size(): number {
+		return this._size;
+	}
+
+	queue(factory: ITask<Promise<T>>): Promise<T> {
+		this._size++;
+
+		return new Promise<T>((c, e) => {
+			this.outstandingPromises.push({ factory, c, e });
+			this.consume();
+		});
+	}
+
+	private consume(): void {
+		while (this.outstandingPromises.length && this.runningPromises < this.maxDegreeOfParalellism) {
+			const iLimitedTask = this.outstandingPromises.shift()!;
+			this.runningPromises++;
+
+			const promise = iLimitedTask.factory();
+			promise.then(iLimitedTask.c, iLimitedTask.e);
+			promise.then(() => this.consumed(), () => this.consumed());
+		}
+	}
+
+	private consumed(): void {
+		this._size--;
+		this.runningPromises--;
+
+		if (this.outstandingPromises.length > 0) {
+			this.consume();
+		}
+	}
+}

extensions/markdown-language-features/src/workspaceContents.ts

@@ -4,9 +4,11 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
+import { coalesce } from './util/arrays';
 import { Disposable } from './util/dispose';
 import { isMarkdownFile } from './util/file';
 import { InMemoryDocument } from './util/inMemoryDocument';
+import { Limiter } from './util/limiter';
 
 /**
  * Minimal version of {@link vscode.TextLine}. Used for mocking out in testing.
@@ -68,15 +70,14 @@ export class VsCodeMdWorkspaceContents extends Disposable implements MdWorkspace
 	 */
 	async getAllMarkdownDocuments(): Promise<SkinnyTextDocument[]> {
 		const maxConcurrent = 20;
-		const docList: SkinnyTextDocument[] = [];
+
 		const resources = await vscode.workspace.findFiles('**/*.md', '**/node_modules/**');
 
-		for (let i = 0; i < resources.length; i += maxConcurrent) {
-			const resourceBatch = resources.slice(i, i + maxConcurrent);
-			const documentBatch = (await Promise.all(resourceBatch.map(x => this.getMarkdownDocument(x)))).filter((doc) => !!doc) as SkinnyTextDocument[];
-			docList.push(...documentBatch);
-		}
-		return docList;
+		const limiter = new Limiter<SkinnyTextDocument | undefined>(maxConcurrent);
+		const results = await Promise.all(resources.map(resource => {
+			return limiter.queue(() => this.getMarkdownDocument(resource));
+		}));
+		return coalesce(results);
 	}
 
 	public get onDidChangeMarkdownDocument() {