Wait for schemas to be loaded

Author: satyakigh

src/datastore/LMDB.ts

@@ -12,7 +12,6 @@ import { stats } from './lmdb/Stats';
 import { encryptionStrategy } from './lmdb/Utils';
 
 export class LMDBStoreFactory implements DataStoreFactory {
-    private isClosed: boolean = false;
     private readonly log = LoggerFactory.getLogger('LMDB.Global');
     @Telemetry({ scope: 'LMDB.Global' }) private readonly telemetry!: ScopedTelemetry;
 
@@ -87,18 +86,14 @@ export class LMDBStoreFactory implements DataStoreFactory {
     }
 
     async close(): Promise<void> {
-        if (this.isClosed) {
-            return;
-        }
+        clearInterval(this.metricsInterval);
+        clearTimeout(this.timeout);
 
         // Clear the stores map but don't close individual stores
         // LMDB will close them when we close the environment
-        clearInterval(this.metricsInterval);
-        clearTimeout(this.timeout);
         this.stores.clear();
         await this.env.flushed;
         await this.env.close();
-        this.isClosed = true;
     }
 
     private cleanupOldVersions(): void {

src/server/CfnServer.ts

@@ -61,7 +61,6 @@ import { ServerComponents } from './ServerComponents';
 
 const log = LoggerFactory.getLogger('CfnServer');
 export class CfnServer {
-    private isClosed: boolean = false;
     private readonly components: ServerComponents;
 
     constructor(
@@ -290,10 +289,6 @@ export class CfnServer {
     }
 
     async close(): Promise<void> {
-        if (this.isClosed) {
-            return;
-        }
         await closeSafely(this.providers, this.external, this.core);
-        this.isClosed = true;
     }
 }

tst/unit/utils/wheel-download.test.ts tst/unit/tools/wheel-download.test.tstst/unit/utils/wheel-download.test.ts renamed to tst/unit/tools/wheel-download.test.ts

@@ -53,6 +53,7 @@ import { MultiDataStoreFactoryProvider } from '../../src/datastore/DataStore';
 import { FeatureFlagProvider } from '../../src/featureFlag/FeatureFlagProvider';
 import { LspCapabilities } from '../../src/protocol/LspCapabilities';
 import { LspConnection } from '../../src/protocol/LspConnection';
+import { SamStoreKey } from '../../src/schema/SamSchemas';
 import { SchemaRetriever } from '../../src/schema/SchemaRetriever';
 import { SchemaStore } from '../../src/schema/SchemaStore';
 import { CfnExternal } from '../../src/server/CfnExternal';
@@ -61,6 +62,7 @@ import { CfnLspProviders } from '../../src/server/CfnLspProviders';
 import { CfnServer } from '../../src/server/CfnServer';
 import { AwsMetadata, ExtendedInitializeParams } from '../../src/server/InitParams';
 import { RelationshipSchemaService } from '../../src/services/RelationshipSchemaService';
+import { DefaultSettings } from '../../src/settings/Settings';
 import { LoggerFactory } from '../../src/telemetry/LoggerFactory';
 import { Closeable } from '../../src/utils/Closeable';
 import { ExtensionName } from '../../src/utils/ExtensionConfig';
@@ -108,8 +110,7 @@ export class TestExtension implements Closeable {
     providers!: CfnLspProviders;
     server!: CfnServer;
 
-    private lspClientReady = false;
-    private lspServerReady = false;
+    private isReady = false;
 
     constructor() {
         this.serverConnection = new LspConnection(
@@ -154,16 +155,9 @@ export class TestExtension implements Closeable {
                     this.server = new CfnServer(lsp, this.core, this.external, this.providers);
                     return LspCapabilities;
                 },
-                onInitialized: (params) => {
-                    this.server.initialized(params);
-                    this.lspServerReady = true;
-                },
-                onShutdown: () => {
-                    return this.server.close();
-                },
-                onExit: () => {
-                    return this.server.close();
-                },
+                onInitialized: (params) => this.server.initialized(params),
+                onShutdown: () => this.server.close(),
+                onExit: () => this.server.close(),
             },
         );
 
@@ -185,40 +179,39 @@ export class TestExtension implements Closeable {
     }
 
     async ready() {
-        if (!this.lspClientReady) {
+        if (!this.isReady) {
             await this.clientConnection.sendRequest(InitializeRequest.type, this.initializeParams);
             await this.clientConnection.sendNotification(InitializedNotification.type, {});
-            this.lspClientReady = true;
-        }
 
-        await WaitFor.waitFor(() => {
-            if (!this.lspServerReady) {
-                throw new Error('Server is not ready yet');
-            }
-        }, 5000);
+            await WaitFor.waitFor(() => {
+                const store = this.external.schemaStore;
+                const pbSchemas = store?.publicSchemas?.get(DefaultSettings.profile.region);
+                const samSchemas = store?.samSchemas?.get(SamStoreKey);
+
+                if (pbSchemas === undefined || samSchemas === undefined) {
+                    throw new Error('Schemas not loaded yet');
+                }
+            }, 2_500);
 
-        await flushAllPromises();
+            await flushAllPromises();
+            this.isReady = true;
+        }
     }
 
     async send(method: string, params: any) {
         await this.ready();
-        const value = await this.clientConnection.sendRequest(method, params);
-        await wait(100);
-        return value;
+        return await this.clientConnection.sendRequest(method, params);
     }
 
     async notify(method: string, params: any) {
         await this.ready();
-        const value = await this.clientConnection.sendNotification(method, params);
-        await wait(100);
-        return value;
+        return await this.clientConnection.sendNotification(method, params);
     }
 
     async close() {
         await this.clientConnection.sendRequest(ShutdownRequest.type);
         await this.clientConnection.sendNotification(ExitNotification.type);
         this.clientConnection.dispose();
-        await flushAllPromises();
     }
 
     // ====================================================================
@@ -227,18 +220,22 @@ export class TestExtension implements Closeable {
 
     async openDocument(params: DidOpenTextDocumentParams) {
         await this.notify(DidOpenTextDocumentNotification.method, params);
+        await wait(10);
     }
 
     async changeDocument(params: DidChangeTextDocumentParams) {
         await this.notify(DidChangeTextDocumentNotification.method, params);
+        await wait(10);
     }
 
     async closeDocument(params: DidCloseTextDocumentParams) {
         await this.notify(DidCloseTextDocumentNotification.method, params);
+        await wait(10);
     }
 
-    saveDocument(params: DidSaveTextDocumentParams) {
-        return this.notify(DidSaveTextDocumentNotification.method, params);
+    async saveDocument(params: DidSaveTextDocumentParams) {
+        await this.notify(DidSaveTextDocumentNotification.method, params);
+        await wait(10);
     }
 
     completion(params: CompletionParams) {

tst/utils/TestExtension.ts

@@ -1,7 +1,7 @@
-export function flushAllPromises() {
-    return new Promise((resolve) => {
-        setTimeout(resolve, 0);
-    });
+import { setImmediate } from 'node:timers/promises';
+
+export async function flushAllPromises() {
+    await setImmediate();
 }
 
 export class WaitFor {
@@ -33,7 +33,11 @@ export class WaitFor {
         throw lastError!;
     }
 
-    static async waitFor(code: () => void | Promise<void>, timeoutMs: number = 100): Promise<void> {
-        await new WaitFor(timeoutMs).wait(code);
+    static async waitFor(
+        code: () => void | Promise<void>,
+        timeoutMs: number = 100,
+        intervalMs: number = 10,
+    ): Promise<void> {
+        await new WaitFor(timeoutMs, intervalMs).wait(code);
     }
 }