perf(startup): setup cfn template registry asynchronously #3370

Problem:
The cfn registry setup can take a long time if there
are a lot of yaml files. This can significantly slow
down extension startup.

Solution:
- Initialize the registry asynchronously.
- If the user triggers functionality that
  requires the template registry and it is not yet
  ready, show a progress message.

Signed-off-by: Nikolas Komonen <[email protected]>

Author: nkomonen-amazon

src/shared/cloudformation/activation.ts

@@ -7,7 +7,7 @@ import * as vscode from 'vscode'
 import { getLogger } from '../logger'
 import { localize } from '../utilities/vsCodeUtils'
 
-import { CloudFormationTemplateRegistry } from '../fs/templateRegistry'
+import { AsyncCloudFormationTemplateRegistry, CloudFormationTemplateRegistry } from '../fs/templateRegistry'
 import { getIdeProperties } from '../extensionUtilities'
 import { NoopWatcher } from '../fs/watchedFiles'
 import { createStarterTemplateFile } from './cloudformation'
@@ -35,11 +35,8 @@ export const devfileExcludePattern = /.*devfile\.(yaml|yml)/i
 export async function activate(extensionContext: vscode.ExtensionContext): Promise<void> {
     try {
         const registry = new CloudFormationTemplateRegistry()
-        globals.templateRegistry = registry
-        await registry.addExcludedPattern(devfileExcludePattern)
-        await registry.addExcludedPattern(templateFileExcludePattern)
-        await registry.addWatchPattern(templateFileGlobPattern)
-        await registry.watchUntitledFiles()
+        extensionContext.subscriptions.push(registry)
+        setTemplateRegistryInGlobals(registry)
     } catch (e) {
         vscode.window.showErrorMessage(
             localize(
@@ -48,15 +45,49 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi
                 getIdeProperties().codelenses
             )
         )
-        getLogger().error('Failed to activate template registry', e)
+        getLogger().error('Failed to activate template registry: %s', e)
         // This prevents us from breaking for any reason later if it fails to load. Since
         // Noop watcher is always empty, we will get back empty arrays with no issues.
         globals.templateRegistry = new NoopWatcher() as unknown as CloudFormationTemplateRegistry
     }
     // If setting it up worked, add it to subscriptions so it is cleaned up at exit
     extensionContext.subscriptions.push(
-        globals.templateRegistry,
         Commands.register('aws.cloudFormation.newTemplate', () => createStarterTemplateFile(false)),
         Commands.register('aws.sam.newTemplate', () => createStarterTemplateFile(true))
     )
 }
+
+/**
+ * Sets the `templateRegistry` property in the `globals` variable,
+ * where the value of the property depends on whether the registry
+ * is fully set up.
+ *
+ * This function exists to resolve the registry setup taking a long time
+ * and slowing down the extension starting up.
+ */
+function setTemplateRegistryInGlobals(registry: CloudFormationTemplateRegistry) {
+    const registrySetupFunc = async (registry: CloudFormationTemplateRegistry) => {
+        await registry.addExcludedPattern(devfileExcludePattern)
+        await registry.addExcludedPattern(templateFileExcludePattern)
+        await registry.addWatchPattern(templateFileGlobPattern)
+        await registry.watchUntitledFiles()
+    }
+
+    const asyncRegistry = new AsyncCloudFormationTemplateRegistry(registry, registrySetupFunc)
+
+    Object.defineProperty(globals, 'templateRegistry', {
+        set(newInstance: CloudFormationTemplateRegistry) {
+            this.cfnInstance = newInstance
+        },
+        get() {
+            // This condition handles testing scenarios where we may have
+            // already set a mock object before activation.
+            // Though in prod nothing should be calling this 'set' function.
+            if (this.cfnInstance) {
+                return this.cfnInstance
+            }
+
+            return asyncRegistry.getInstance()
+        },
+    })
+}

src/shared/fs/templateRegistry.ts

@@ -11,10 +11,12 @@ import * as path from 'path'
 import { isInDirectory } from '../filesystemUtilities'
 import { dotNetRuntimes, goRuntimes, javaRuntimes } from '../../lambda/models/samLambdaRuntime'
 import { getLambdaDetails } from '../../lambda/utils'
-import { WatchedFiles, WatchedItem } from './watchedFiles'
+import { NoopWatcher, WatchedFiles, WatchedItem } from './watchedFiles'
 import { getLogger } from '../logger'
 import globals from '../extensionGlobals'
 import { isUntitledScheme, normalizeVSCodeUri } from '../utilities/vsCodeUtils'
+import { sleep } from '../utilities/timeoutUtils'
+import { localize } from '../utilities/vsCodeUtils'
 
 export class CloudFormationTemplateRegistry extends WatchedFiles<CloudFormation.Template> {
     protected name: string = 'CloudFormationTemplateRegistry'
@@ -68,6 +70,69 @@ export class CloudFormationTemplateRegistry extends WatchedFiles<CloudFormation.
     }
 }
 
+/**
+ * The purpose of this class is to resolve a {@link CloudFormationTemplateRegistry}
+ * instance once the given setup is complete.
+ *
+ * This solves the issue where setup can take a while and if we
+ * block on it the entire extension startup time is increased.
+ */
+export class AsyncCloudFormationTemplateRegistry {
+    /** Setup of the registry can take a while, this property indicates it is done */
+    private isSetup = false
+    /** The message that is shown to the user to indicate the registry is being set up */
+    private setupProgressMessage: Thenable<void> | undefined = undefined
+
+    /**
+     * @param asyncSetupFunc registry setup that will be run async
+     */
+    constructor(
+        private readonly instance: CloudFormationTemplateRegistry,
+        asyncSetupFunc: (instance: CloudFormationTemplateRegistry) => Promise<void>
+    ) {
+        getLogger().info('cfn: starting template registry setup.')
+        asyncSetupFunc(instance).then(() => {
+            this.isSetup = true
+            getLogger().info('cfn: template registry setup successful.')
+        })
+    }
+
+    /**
+     * Returns the initial registry instance if setup has completed, otherwise
+     * returning a temporary instance and showing a progress message to the user
+     * to indicate setup is in progress.
+     */
+    getInstance(): CloudFormationTemplateRegistry {
+        if (this.isSetup) {
+            return this.instance
+        }
+
+        // Show user a message indicating setup is in progress
+        if (this.setupProgressMessage === undefined) {
+            this.setupProgressMessage = vscode.window.withProgress(
+                {
+                    location: vscode.ProgressLocation.Notification,
+                    title: localize('AWS.codelens.waitingForTemplateRegistry', 'Scanning CloudFormation templates...'),
+                    cancellable: true,
+                },
+                async (progress, token) => {
+                    token.onCancellationRequested(() => {
+                        // Allows for new message to be created if templateRegistry variable attempted to be used again
+                        this.setupProgressMessage = undefined
+                    })
+                    getLogger().info('cfn: Waiting for template registry setup to complete.')
+                    while (!this.isSetup) {
+                        await sleep(2000)
+                    }
+                    getLogger().info('cfn: Finished waiting for template registry setup.')
+                }
+            )
+        }
+
+        return new NoopWatcher() as unknown as CloudFormationTemplateRegistry
+    }
+}
+
 /**
  * Gets resources and additional metadata for resources tied to a filepath and handler.
  * Checks all registered templates by default; otherwise can operate on a subset TemplateDatum[]