fix: race condition in diagnostics

Author: mhuisi

vscode-lean4/src/diagnostics.ts

@@ -1,6 +1,7 @@
-import { DiagnosticCollection, Disposable, EventEmitter } from 'vscode'
-import { LeanDiagnostic, LeanPublishDiagnosticsParams, p2cConverter } from './utils/converters'
+import { CancellationToken, DiagnosticCollection, Disposable, EventEmitter } from 'vscode'
 import { DocumentUri } from 'vscode-languageserver-protocol'
+import { CoalescingSyncQueue } from './utils/coalescingSyncQueue'
+import { LeanDiagnostic, LeanPublishDiagnosticsParams, p2cConverter } from './utils/converters'
 
 export type DiagnosticChangeKind = 'replace' | 'append'
 
@@ -26,15 +27,39 @@ export class DiagnosticChangeEvent {
     }
 }
 
+type SyncQueueEntry = {
+    accumulatedParams: LeanPublishDiagnosticsParams
+    pendingKind: DiagnosticChangeKind
+    pendingBatch: LeanDiagnostic[]
+}
+
+function combineEntries(existing: SyncQueueEntry, incoming: SyncQueueEntry): SyncQueueEntry {
+    if (incoming.pendingKind === 'replace') {
+        return incoming
+    }
+    incoming.pendingKind satisfies 'append'
+    return {
+        accumulatedParams: incoming.accumulatedParams,
+        pendingKind: existing.pendingKind,
+        pendingBatch: [...existing.pendingBatch, ...incoming.pendingBatch],
+    }
+}
+
 export class LeanClientDiagnosticCollection implements Disposable {
     readonly vsCodeCollection: DiagnosticCollection
     private diags: Map<DocumentUri, LeanPublishDiagnosticsParams> = new Map()
 
     private diagnosticsChangedEmitter = new EventEmitter<DiagnosticChangeEvent>()
     onDidChangeDiagnostics = this.diagnosticsChangedEmitter.event
 
+    private syncQueue: CoalescingSyncQueue<SyncQueueEntry>
+
     constructor(vsCodeCollection: DiagnosticCollection) {
         this.vsCodeCollection = vsCodeCollection
+        this.syncQueue = new CoalescingSyncQueue(
+            (uri: string, entry: SyncQueueEntry, token: CancellationToken) => this.syncToCollection(uri, entry, token),
+            (existing, incoming) => combineEntries(existing, incoming)
+        )
     }
 
     private static determineChangeKind(
@@ -63,20 +88,33 @@ export class LeanClientDiagnosticCollection implements Disposable {
         }
 
         const accumulatedParams = { ...params, diagnostics: accumulated }
-
         this.diags.set(accumulatedParams.uri, accumulatedParams)
-        void this.syncToCollection(accumulatedParams)
-        this.diagnosticsChangedEmitter.fire(new DiagnosticChangeEvent(kind, params, accumulated))
+
+        const entry: SyncQueueEntry = {
+            accumulatedParams,
+            pendingKind: kind,
+            pendingBatch: [...params.diagnostics],
+        }
+
+        this.syncQueue.enqueue(accumulatedParams.uri, entry)
     }
 
-    private async syncToCollection(p: LeanPublishDiagnosticsParams): Promise<void> {
-        const nonSilent = p.diagnostics.filter(d => !d.isSilent)
-        const uri = p2cConverter.asUri(p.uri)
-        const vsCodeDiags = await p2cConverter.asDiagnostics(nonSilent)
-        this.vsCodeCollection.set(uri, vsCodeDiags)
+    private async syncToCollection(_uri: string, entry: SyncQueueEntry, token: CancellationToken): Promise<void> {
+        const nonSilentDiagnostics = entry.accumulatedParams.diagnostics.filter(d => !d.isSilent)
+        const vsUri = p2cConverter.asUri(entry.accumulatedParams.uri)
+        const vsDiags = await p2cConverter.asDiagnostics(nonSilentDiagnostics, token)
+        if (token.isCancellationRequested) {
+            return
+        }
+        this.vsCodeCollection.set(vsUri, vsDiags)
+        const collapsedParams = { ...entry.accumulatedParams, diagnostics: entry.pendingBatch }
+        this.diagnosticsChangedEmitter.fire(
+            new DiagnosticChangeEvent(entry.pendingKind, collapsedParams, entry.accumulatedParams.diagnostics),
+        )
     }
 
     dispose(): void {
+        this.syncQueue.dispose()
         this.diagnosticsChangedEmitter.dispose()
     }
 }

vscode-lean4/src/utils/coalescingSyncQueue.ts

@@ -0,0 +1,79 @@
+import { CancellationToken, CancellationTokenSource, Disposable } from 'vscode'
+
+/**
+ * A queue that coalesces pending work by string key, processes one entry at a time,
+ * and supports cancellation of in-flight work when the same key is re-enqueued.
+ *
+ * When a key is enqueued while its worker is already running, the in-flight
+ * CancellationToken is cancelled. The worker should check the token and bail early.
+ * If the worker completes without cancellation, the entry is removed from the queue.
+ * If cancelled or errored, the entry stays so the next enqueue can update it.
+ */
+export class CoalescingSyncQueue<T> implements Disposable {
+    private queue: Map<string, T> = new Map()
+    private state: { kind: 'idle' } | { kind: 'busy'; key: string; tokenSource: CancellationTokenSource } = {
+        kind: 'idle',
+    }
+    private disposed = false
+
+    constructor(
+        private readonly worker: (key: string, value: T, token: CancellationToken) => Promise<void>,
+        private readonly merge?: (existing: T, incoming: T) => T,
+    ) {}
+
+    enqueue(key: string, value: T): void {
+        if (this.state.kind === 'busy' && this.state.key === key) {
+            this.state.tokenSource.cancel()
+        }
+
+        const existing = this.queue.get(key)
+        if (this.merge !== undefined && existing !== undefined) {
+            this.queue.set(key, this.merge(existing, value))
+        } else {
+            this.queue.set(key, value)
+        }
+
+        if (this.state.kind === 'idle') {
+            this.trigger()
+        }
+    }
+
+    private trigger(): void {
+        queueMicrotask(() => void this.work())
+    }
+
+    private async work(): Promise<void> {
+        if (this.state.kind === 'busy' || this.disposed) {
+            return
+        }
+
+        const next = this.queue.entries().next()
+        if (next.done) {
+            return
+        }
+        const [key, value] = next.value
+
+        const tokenSource = new CancellationTokenSource()
+        this.state = { kind: 'busy', key, tokenSource }
+        try {
+            await this.worker(key, value, tokenSource.token)
+            if (!tokenSource.token.isCancellationRequested && !this.disposed) {
+                this.queue.delete(key)
+            }
+        } finally {
+            tokenSource.dispose()
+            this.state = { kind: 'idle' }
+            if (this.queue.size > 0 && !this.disposed) {
+                this.trigger()
+            }
+        }
+    }
+
+    dispose(): void {
+        this.disposed = true
+        if (this.state.kind === 'busy') {
+            this.state.tokenSource.cancel()
+        }
+        this.queue.clear()
+    }
+}