Fix MdDocumentInfoCache computing values twice (microsoft#152799)

* Fix MdDocumentInfoCache computing values twice

Fixes a race where values could be computed twice before being cached

* Remove only

Author: mjbvz

extensions/markdown-language-features/src/test/documentInfoCache.test.ts

@@ -0,0 +1,33 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import 'mocha';
+import { InMemoryDocument } from '../util/inMemoryDocument';
+import { MdDocumentInfoCache } from '../util/workspaceCache';
+import { InMemoryWorkspaceMarkdownDocuments } from './inMemoryWorkspace';
+import { workspacePath } from './util';
+
+suite('DocumentInfoCache', () => {
+	test('Repeated calls should only compute value once', async () => {
+		const doc = workspacePath('doc.md');
+		const workspace = new InMemoryWorkspaceMarkdownDocuments([
+			new InMemoryDocument(doc, '')
+		]);
+
+		let i = 0;
+		const cache = new MdDocumentInfoCache<number>(workspace, async () => {
+			return ++i;
+		});
+
+		const a = cache.get(doc);
+		const b = cache.get(doc);
+
+		assert.strictEqual(await a, 1);
+		assert.strictEqual(i, 1);
+		assert.strictEqual(await b, 1);
+		assert.strictEqual(i, 1);
+	});
+});

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

@@ -43,48 +43,70 @@ class LazyResourceMap<T> {
 export class MdDocumentInfoCache<T> extends Disposable {
 
 	private readonly _cache = new LazyResourceMap<T>();
+	private readonly _loadingDocuments = new ResourceMap<Promise<SkinnyTextDocument | undefined>>();
 
 	public constructor(
 		private readonly workspaceContents: MdWorkspaceContents,
 		private readonly getValue: (document: SkinnyTextDocument) => Promise<T>,
 	) {
 		super();
 
-		this._register(this.workspaceContents.onDidChangeMarkdownDocument(doc => this.onDidChangeDocument(doc)));
-		this._register(this.workspaceContents.onDidCreateMarkdownDocument(doc => this.onDidChangeDocument(doc)));
+		this._register(this.workspaceContents.onDidChangeMarkdownDocument(doc => this.invalidate(doc)));
 		this._register(this.workspaceContents.onDidDeleteMarkdownDocument(this.onDidDeleteDocument, this));
 	}
 
 	public async get(resource: vscode.Uri): Promise<T | undefined> {
-		const existing = this._cache.get(resource);
+		let existing = this._cache.get(resource);
 		if (existing) {
 			return existing;
 		}
 
-		const doc = await this.workspaceContents.getOrLoadMarkdownDocument(resource);
-		return doc && this.onDidChangeDocument(doc, true)?.value;
+		const doc = await this.loadDocument(resource);
+		if (!doc) {
+			return undefined;
+		}
+
+		// Check if we have invalidated
+		existing = this._cache.get(resource);
+		if (existing) {
+			return existing;
+		}
+
+		return this.resetEntry(doc)?.value;
 	}
 
 	public async getForDocument(document: SkinnyTextDocument): Promise<T> {
 		const existing = this._cache.get(document.uri);
 		if (existing) {
 			return existing;
 		}
+		return this.resetEntry(document).value;
+	}
 
-		return this.onDidChangeDocument(document, true)!.value;
+	private loadDocument(resource: vscode.Uri): Promise<SkinnyTextDocument | undefined> {
+		const existing = this._loadingDocuments.get(resource);
+		if (existing) {
+			return existing;
+		}
+
+		const p = this.workspaceContents.getOrLoadMarkdownDocument(resource);
+		this._loadingDocuments.set(resource, p);
+		p.finally(() => {
+			this._loadingDocuments.delete(resource);
+		});
+		return p;
 	}
 
-	public async entries(): Promise<Array<[vscode.Uri, T]>> {
-		return this._cache.entries();
+	private resetEntry(document: SkinnyTextDocument): Lazy<Promise<T>> {
+		const value = lazy(() => this.getValue(document));
+		this._cache.set(document.uri, value);
+		return value;
 	}
 
-	private onDidChangeDocument(document: SkinnyTextDocument, forceAdd = false): Lazy<Promise<T>> | undefined {
-		if (forceAdd || this._cache.has(document.uri)) {
-			const value = lazy(() => this.getValue(document));
-			this._cache.set(document.uri, value);
-			return value;
+	private invalidate(document: SkinnyTextDocument): void {
+		if (this._cache.has(document.uri)) {
+			this.resetEntry(document);
 		}
-		return undefined;
 	}
 
 	private onDidDeleteDocument(resource: vscode.Uri) {