Preserve diagnostic pull state after quick document reopen

client-node-tests/src/integration.test.ts

@@ -89,6 +89,10 @@ function isFullDocumentDiagnosticReport(value: vsdiag.DocumentDiagnosticReport):
 	assert.ok(value.kind === vsdiag.DocumentDiagnosticReportKind.full);
 }
 
+function waitForNextTurn(): Promise<void> {
+	return new Promise((resolve) => setTimeout(resolve, 0));
+}
+
 function processExists(pid: number): boolean {
 	try {
 		process.kill(pid, 0);
@@ -247,12 +251,13 @@ suite('Client integration', () => {
 			run: { module: serverModule, transport: lsclient.TransportKind.ipc },
 			debug: { module: serverModule, transport: lsclient.TransportKind.ipc, options: { execArgv: ['--nolazy', '--inspect=6014'] } }
 		};
-		const documentSelector: lsclient.DocumentSelector = [{ scheme: 'lsptests', language: 'bat' }];
+		const documentSelector: lsclient.DocumentSelector = [{ scheme: 'lsptests', language: 'bat' }, { scheme: 'lsptests', language: 'plaintext' }];
 
 		middleware = {};
 		const clientOptions: lsclient.LanguageClientOptions = {
 			documentSelector, synchronize: {}, initializationOptions: {}, middleware,
 			workspaceFolder: { index: 0, name: 'test_folder', uri: vscode.Uri.parse(`${fsProvider.scheme}:///`) },
+			diagnosticPullOptions: { onChange: true }
 		};
 
 		client = new lsclient.LanguageClient('test svr', 'Test Language Server', serverOptions, clientOptions);
@@ -1444,6 +1449,73 @@ suite('Client integration', () => {
 		assert.strictEqual(reporterCalled, true);
 	});
 
+	test('Document diagnostic pull after quick reopen', async () => {
+		await vscode.window.showTextDocument(document);
+		let initialPullFinished: (() => void) | undefined;
+		const initialPull = new Promise<void>((resolve) => {
+			initialPullFinished = resolve;
+		});
+		(middleware as DiagnosticProviderMiddleware).provideDiagnostics = async (document, previousResultId, token, next) => {
+			const result = await next(document, previousResultId, token);
+			initialPullFinished?.();
+			initialPullFinished = undefined;
+			return result;
+		};
+		document = await vscode.languages.setTextDocumentLanguage(document, 'plaintext');
+		await initialPull;
+		await waitForNextTurn();
+
+		let pullCount = 0;
+		let releasePull!: () => void;
+		const holdPull = new Promise<void>((resolve) => {
+			releasePull = resolve;
+		});
+		let firstPullStarted!: () => void;
+		const firstPull = new Promise<void>((resolve) => {
+			firstPullStarted = resolve;
+		});
+		let firstPullFinished!: () => void;
+		const firstPullDone = new Promise<void>((resolve) => {
+			firstPullFinished = resolve;
+		});
+		let expectedEditPull: number | undefined;
+		let editPullFinished!: () => void;
+		const editPull = new Promise<void>((resolve) => {
+			editPullFinished = resolve;
+		});
+		(middleware as DiagnosticProviderMiddleware).provideDiagnostics = async (document, previousResultId, token, next) => {
+			const currentPull = ++pullCount;
+			if (currentPull === 1) {
+				firstPullStarted();
+				await holdPull;
+			}
+			const result = await next(document, previousResultId, token);
+			if (currentPull === 1) {
+				firstPullFinished();
+			}
+			if (currentPull === expectedEditPull) {
+				expectedEditPull = undefined;
+				editPullFinished();
+			}
+			return result;
+		};
+
+		const changeLanguage = vscode.languages.setTextDocumentLanguage(document, 'bat');
+		await firstPull;
+		document = await changeLanguage;
+		releasePull();
+		await firstPullDone;
+		await waitForNextTurn();
+
+		expectedEditPull = pullCount + 1;
+		const edit = new vscode.WorkspaceEdit();
+		edit.insert(document.uri, new vscode.Position(0, 0), ' ');
+		await vscode.workspace.applyEdit(edit);
+		await editPull;
+		(middleware as DiagnosticProviderMiddleware).provideDiagnostics = undefined;
+		await revertAllDirty();
+	}).timeout(5000);
+
 	test('Type Hierarchy', async () => {
 		const provider = client.getFeature(lsclient.TypeHierarchyPrepareRequest.method).getProvider(document);
 		isDefined(provider);

client/src/common/diagnostic.ts

@@ -307,6 +307,7 @@ class DiagnosticRequestor implements Disposable {
 	public readonly provider: vsdiag.DiagnosticProvider;
 	private readonly diagnostics: DiagnosticCollection;
 	private readonly openRequests: Map<string, RequestState>;
+	private readonly pendingDocumentForgets: Map<string, symbol>;
 	private readonly documentStates: DocumentPullStateTracker;
 
 	private workspaceErrorCounter: number;
@@ -324,6 +325,7 @@ class DiagnosticRequestor implements Disposable {
 
 		this.diagnostics = this.createDiagnosticCollection();
 		this.openRequests = new Map();
+		this.pendingDocumentForgets = new Map();
 		this.documentStates = new DocumentPullStateTracker();
 		this.workspaceErrorCounter = 0;
 	}
@@ -361,6 +363,10 @@ class DiagnosticRequestor implements Disposable {
 		this.documentStates.unTrack(kind, document);
 	}
 
+	public cancelPendingForget(document: TextDocument | Uri): void {
+		this.pendingDocumentForgets.delete(DocumentOrUri.asKey(document));
+	}
+
 	public pull(document: TextDocument | Uri, cb?: () => void): void {
 		if (this.isDisposed) {
 			return;
@@ -454,8 +460,13 @@ class DiagnosticRequestor implements Disposable {
 			if (request !== undefined) {
 				this.openRequests.set(key, { state: RequestStateKind.reschedule, document: document });
 			} else {
+				const pendingForget = Symbol();
+				this.pendingDocumentForgets.set(key, pendingForget);
 				this.pull(document, () => {
-					this.forget(PullState.document, document);
+					if (this.pendingDocumentForgets.get(key) === pendingForget) {
+						this.pendingDocumentForgets.delete(key);
+						this.forget(PullState.document, document);
+					}
 				});
 			}
 
@@ -916,11 +927,12 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 		const openFeature = client.getFeature(DidOpenTextDocumentNotification.method);
 		disposables.push(openFeature.onNotificationSent((event) => {
 			const textDocument = event.textDocument;
-			// We already know about this document. This can happen via a tab open.
-			if (this.diagnosticRequestor.knowsSameVersion(PullState.document, textDocument)) {
-				return;
-			}
 			if (matches(textDocument)) {
+				this.diagnosticRequestor.cancelPendingForget(textDocument);
+				// We already know about this document. This can happen via a tab open.
+				if (this.diagnosticRequestor.knowsSameVersion(PullState.document, textDocument)) {
+					return;
+				}
 				this.diagnosticRequestor.pull(textDocument, () => { addToBackgroundIfNeeded(textDocument); });
 			}
 		}));
@@ -929,6 +941,7 @@ class DiagnosticFeatureProviderImpl implements DiagnosticProviderShape {
 			// Send a pull for all opened cells in the notebook.
 			for (const cell of event.getCells()) {
 				if (matchesCell(cell)) {
+					this.diagnosticRequestor.cancelPendingForget(cell.document);
 					this.diagnosticRequestor.pull(cell.document, () => { addToBackgroundIfNeeded(cell.document); });
 				}
 			}