Azure · morgilad · Oct 7, 2024 · Oct 7, 2024 · noam-shapira · Oct 7, 2024
diff --git a/README.md b/README.md
@@ -85,6 +85,10 @@ Every PR should come with a test that checks it.
 
 ## Changelog
 
+### 12.0.11
+
+-   fix: avoid race condition on worker disposal
+
 ### 12.0.10
 
 -   fix: Parameter docstrings not shown

diff --git a/package/package.json b/package/package.json
@@ -1,6 +1,6 @@
 {
     "name": "@kusto/monaco-kusto",
-    "version": "12.0.10",
+    "version": "12.0.11",
     "description": "CSL, KQL plugin for the Monaco Editor",
     "author": {
         "name": "Microsoft"

diff --git a/package/src/kustoMode.ts b/package/src/kustoMode.ts
@@ -30,14 +30,11 @@ let workerPromise: Promise<AugmentedWorkerAccessor> = new Promise((resolve, reje
  * Called when Kusto language is first needed (a model has the language set)
  * @param defaults
  */
-export function setupMode(defaults: LanguageServiceDefaults, monacoInstance: typeof globalThis.monaco): IDisposable {
+export function setupMode(defaults: LanguageServiceDefaults, monacoInstance: typeof globalThis.monaco) {
     const onSchemaChange = new monaco.Emitter<Schema>();
-    // TODO: when should we dispose of these? seems like monaco-css and monaco-typescript don't dispose of these.
-    const disposables: monaco.IDisposable[] = [];
     const semanticTokensProviderRegistrar = semanticTokensProviderRegistrarCreator();
 
     const client = new WorkerManager(monacoInstance, defaults);
-    disposables.push(client);
 
     const workerAccessor: AugmentedWorkerAccessor = (first, ...more) => {
         const augmentedSetSchema = async (schema: Schema, worker: KustoWorker) => {
@@ -71,64 +68,45 @@ export function setupMode(defaults: LanguageServiceDefaults, monacoInstance: typ
         );
     };
 
-    disposables.push(
-        monacoInstance.languages.registerCompletionItemProvider(
-            LANGUAGE_ID,
-            new languageFeatures.CompletionAdapter(workerAccessor, defaults.languageSettings)
-        )
+    monacoInstance.languages.registerCompletionItemProvider(
+        LANGUAGE_ID,
+        new languageFeatures.CompletionAdapter(workerAccessor, defaults.languageSettings)
     );
 
-    disposables.push(monacoInstance.languages.setMonarchTokensProvider(LANGUAGE_ID, kustoLanguageDefinition));
+    monacoInstance.languages.setMonarchTokensProvider(LANGUAGE_ID, kustoLanguageDefinition);
 
-    disposables.push(
-        new languageFeatures.DiagnosticsAdapter(
-            monacoInstance,
-            LANGUAGE_ID,
-            workerAccessor,
-            defaults,
-            onSchemaChange.event
-        )
+    new languageFeatures.DiagnosticsAdapter(
+        monacoInstance,
+        LANGUAGE_ID,
+        workerAccessor,
+        defaults,
+        onSchemaChange.event
     );
 
-    disposables.push(
-        monacoInstance.languages.registerDocumentRangeFormattingEditProvider(
-            LANGUAGE_ID,
-            new languageFeatures.FormatAdapter(workerAccessor)
-        )
+    monacoInstance.languages.registerDocumentRangeFormattingEditProvider(
+        LANGUAGE_ID,
+        new languageFeatures.FormatAdapter(workerAccessor)
     );
 
-    disposables.push(
-        monacoInstance.languages.registerFoldingRangeProvider(
-            LANGUAGE_ID,
-            new languageFeatures.FoldingAdapter(workerAccessor)
-        )
+    monacoInstance.languages.registerFoldingRangeProvider(
+        LANGUAGE_ID,
+        new languageFeatures.FoldingAdapter(workerAccessor)
     );
 
-    disposables.push(
-        monacoInstance.languages.registerDefinitionProvider(
-            LANGUAGE_ID,
-            new languageFeatures.DefinitionAdapter(workerAccessor)
-        )
+    monacoInstance.languages.registerDefinitionProvider(
+        LANGUAGE_ID,
+        new languageFeatures.DefinitionAdapter(workerAccessor)
     );
 
-    disposables.push(
-        monacoInstance.languages.registerRenameProvider(LANGUAGE_ID, new languageFeatures.RenameAdapter(workerAccessor))
-    );
+    monacoInstance.languages.registerRenameProvider(LANGUAGE_ID, new languageFeatures.RenameAdapter(workerAccessor));
 
-    disposables.push(
-        monacoInstance.languages.registerReferenceProvider(
-            LANGUAGE_ID,
-            new languageFeatures.ReferenceAdapter(workerAccessor)
-        )
+    monacoInstance.languages.registerReferenceProvider(
+        LANGUAGE_ID,
+        new languageFeatures.ReferenceAdapter(workerAccessor)
     );
 
     if (defaults.languageSettings.enableHover) {
-        disposables.push(
-            monacoInstance.languages.registerHoverProvider(
-                LANGUAGE_ID,
-                new languageFeatures.HoverAdapter(workerAccessor)
-            )
-        );
+        monacoInstance.languages.registerHoverProvider(LANGUAGE_ID, new languageFeatures.HoverAdapter(workerAccessor));
     }
 
     monacoInstance.languages.registerDocumentFormattingEditProvider(
@@ -138,21 +116,7 @@ export function setupMode(defaults: LanguageServiceDefaults, monacoInstance: typ
     kustoWorker = workerAccessor;
     resolveWorker(workerAccessor);
 
-    disposables.push(monacoInstance.languages.setLanguageConfiguration(LANGUAGE_ID, kanguageConfiguration));
-
-    return asDisposable(disposables);
-}
-
-function asDisposable(disposables: IDisposable[]): IDisposable {
-    return {
-        dispose: () => {
-            return disposeAll(disposables);
-        },
-    };
-}
-
-function disposeAll(disposables: IDisposable[]) {
-    disposables.forEach((d) => d.dispose());
+    monacoInstance.languages.setLanguageConfiguration(LANGUAGE_ID, kanguageConfiguration);
 }
 
 export function getKustoWorker(): Promise<AugmentedWorkerAccessor> {

diff --git a/package/src/monaco.contribution.ts b/package/src/monaco.contribution.ts
@@ -102,11 +102,7 @@ export const kustoDefaults = new LanguageServiceDefaultsImpl(defaultLanguageSett
 
 let disposable: monaco.IDisposable;
 monaco.languages.onLanguage('kusto', async () => {
-    disposable = await withMode((mode) => mode.setupMode(kustoDefaults, monaco as typeof globalThis.monaco));
-});
-
-monaco.editor.onWillDisposeModel((model) => {
-    disposable.dispose();
+    await withMode((mode) => mode.setupMode(kustoDefaults, monaco as typeof globalThis.monaco));
 });
 
 monaco.languages.register({

diff --git a/package/src/workerManager.ts b/package/src/workerManager.ts
@@ -3,41 +3,50 @@ import type * as monaco from 'monaco-editor/esm/vs/editor/editor.api';
 import type { LanguageServiceDefaults } from './monaco.contribution';
 import type { IKustoWorkerImpl } from './kustoWorker';
 
+interface WorkerDetails {
+    _worker: monaco.editor.MonacoWebWorker<IKustoWorkerImpl>;
+    _client: IKustoWorkerImpl;
+    _lastUsedTime: number;
+}
 export class WorkerManager {
     private _storedState: {
         schema: any;
     };
 
     private _defaults: LanguageServiceDefaults;
     private _idleCheckInterval: number;
-    private _lastUsedTime: number;
     private _configChangeListener: monaco.IDisposable;
 
-    private _worker: monaco.editor.MonacoWebWorker<IKustoWorkerImpl>;
-    private _client: Promise<IKustoWorkerImpl>;
+    private _workerDetails: WorkerDetails;
+    private _workerDetailsPromise: Promise<WorkerDetails>;
 
     constructor(private _monacoInstance: typeof monaco, defaults: LanguageServiceDefaults) {
         this._defaults = defaults;
-        this._worker = null;
         this._idleCheckInterval = self.setInterval(() => this._checkIfIdle(), 30 * 1000);
-        this._lastUsedTime = 0;
         this._configChangeListener = this._defaults.onDidChange(() => this._saveStateAndStopWorker());
     }
 
-    private _stopWorker(): void {
-        if (this._worker) {
-            this._worker.dispose();
-            this._worker = null;
-        }
-        this._client = null;
+    private _stopWorker() {
+        const workerToStop = this._workerDetails;
+        this._workerDetailsPromise = null;
+        this._workerDetails = null;
+
+        // Ensure disposal occurs only after the last request completes.
+        // This is necessary because setting the languageSettings disposes of the worker,
+        // causing the setSchema call to remain unresolved, which prevents the semantic tokens provider from being registered.
+        setTimeout(async () => {
+            if (workerToStop._worker) {
+                workerToStop._worker.dispose();
+            }
+        }, 5000);
     }
 
     private _saveStateAndStopWorker(): void {
-        if (!this._worker) {
+        if (!this._workerDetails?._worker) {
             return;
         }
 
-        this._worker.getProxy().then((proxy) => {
+        this._workerDetails?._worker.getProxy().then((proxy) => {
             proxy.getSchema().then((schema) => {
                 this._storedState = { schema: schema };
                 this._stopWorker();
@@ -52,24 +61,22 @@ export class WorkerManager {
     }
 
     private _checkIfIdle(): void {
-        if (!this._worker) {
+        if (!this._workerDetails?._worker) {
             return;
         }
         const maxIdleTime = this._defaults.getWorkerMaxIdleTime();
-        let timePassedSinceLastUsed = Date.now() - this._lastUsedTime;
+        let timePassedSinceLastUsed = Date.now() - this._workerDetails?._lastUsedTime;
         if (maxIdleTime > 0 && timePassedSinceLastUsed > maxIdleTime) {
             this._saveStateAndStopWorker();
         }
     }
 
-    private _getClient(): Promise<IKustoWorkerImpl> {
-        this._lastUsedTime = Date.now();
-
+    private _getClient(): Promise<WorkerDetails> {
         // Since onDidProvideCompletionItems is not used in web worker, and since functions cannot be trivially serialized (throws exception unable to clone), We remove it here.
         const { onDidProvideCompletionItems, ...languageSettings } = this._defaults.languageSettings;
 
-        if (!this._client) {
-            this._worker = this._monacoInstance.editor.createWebWorker<IKustoWorkerImpl>({
+        if (!this._workerDetailsPromise) {
+            const worker = this._monacoInstance.editor.createWebWorker<IKustoWorkerImpl>({
                 // module that exports the create() method and returns a `KustoWorker` instance
                 moduleId: 'vs/language/kusto/kustoWorker',
 
@@ -82,26 +89,35 @@ export class WorkerManager {
                 },
             });
 
-            this._client = this._worker.getProxy().then((proxy) => {
+            const client = worker.getProxy().then((proxy) => {
                 // push state we held onto before killing the client.
                 if (this._storedState) {
                     return proxy.setSchema(this._storedState.schema).then(() => proxy);
                 } else {
                     return proxy;
                 }
             });
+
+            this._workerDetailsPromise = client.then((client) => {
+                this._workerDetails = {
+                    _worker: worker,
+                    _client: client,
+                    _lastUsedTime: Date.now(),
+                };
+                return this._workerDetails;
+            });
         }
-        return this._client;
+        return this._workerDetailsPromise;
     }
 
     getLanguageServiceWorker(...resources: monaco.Uri[]): Promise<IKustoWorkerImpl> {
         let _client: IKustoWorkerImpl;
         return this._getClient()
             .then((client) => {
-                _client = client;
+                _client = client._client;
             })
             .then((_) => {
-                return this._worker.withSyncedResources(resources);
+                return this._workerDetails?._worker?.withSyncedResources(resources);
             })
             .then((_) => _client);
     }