perf(startup): avoid templateRegistry in awsFiletypes.ts #3962

justinmk3 · web-flow · commit 0c26dc36906b · 2023-10-26T12:43:11.000-07:00
Problem: The AWS Documents handler (`awsFiletypes.ts`) is triggered frequently. It calls `globals.templateRegistry` which triggers an expensive "Scanning CloudFormation templates..." process. #3510 Solution: Extract the validation logic out of `CloudFormationTemplateRegistry.process()` so that `awsFiletypes.ts` can use it without requesting the full registry.
diff --git a/src/shared/awsFiletypes.ts b/src/shared/awsFiletypes.ts
@@ -7,13 +7,13 @@ import * as vscode from 'vscode'
 import * as path from 'path'
 import * as constants from '../shared/constants'
 import * as aslFormats from '../stepFunctions/constants/aslFormats'
-import * as pathutil from '../shared/utilities/pathUtils'
 import * as fsutil from '../shared/filesystemUtilities'
 import * as sysutil from '../shared/systemUtilities'
 import * as collectionUtil from '../shared/utilities/collectionUtils'
 import globals from './extensionGlobals'
 import { telemetry } from './telemetry/telemetry'
 import { AwsFiletype } from './telemetry/telemetry'
+import { CloudFormation } from './cloudformation/cloudformation'
 
 /** AWS filetypes: vscode language ids */
 export const awsFiletypeLangIds = {
@@ -67,9 +67,12 @@ export function activate(): void {
         vscode.workspace.onDidOpenTextDocument(async (doc: vscode.TextDocument) => {
             const isAwsFileExt = isAwsFiletype(doc)
             const isSchemaHandled = globals.schemaService.isMapped(doc.uri)
-            const isCfnTemplate = !!(await globals.templateRegistry).items.find(
-                t => pathutil.normalize(t.path) === pathutil.normalize(doc.fileName)
-            )
+            const cfnTemplate =
+                CloudFormation.isValidFilename(doc.uri) && doc.languageId === 'yaml'
+                    ? await CloudFormation.tryLoad(doc.uri)
+                    : undefined
+            const isCfnTemplate = cfnTemplate?.template !== undefined
+
             if (!isAwsFileExt && !isSchemaHandled && !isCfnTemplate) {
                 return
             }
@@ -78,11 +81,12 @@ export function activate(): void {
             let fileExt: string | undefined = path.extname(doc.fileName)
             fileExt = fileExt ? fileExt : undefined // Telemetry client will fail on empty string.
 
-            // TODO: ask templateRegistry if this is SAM or CFN.
             // TODO: ask schemaService for the precise filetype.
             let telemKind = isAwsConfig(doc.fileName) ? 'awsCredentials' : langidToAwsFiletype(doc.languageId)
-            if (telemKind === 'other') {
-                telemKind = isCfnTemplate ? 'cloudformationSam' : isSchemaHandled ? 'cloudformation' : 'other'
+            if (isCfnTemplate) {
+                telemKind = cfnTemplate.kind === 'sam' ? 'cloudformationSam' : 'cloudformation'
+            } else if (telemKind === 'other') {
+                telemKind = isSchemaHandled ? 'cloudformation' : 'other'
             }
 
             // HACK: for "~/.aws/foo" vscode sometimes _only_ emits "~/.aws/foo.git".
diff --git a/src/shared/cloudformation/activation.ts b/src/shared/cloudformation/activation.ts
@@ -10,23 +10,11 @@ import { isToolkitActive, localize } from '../utilities/vsCodeUtils'
 import { AsyncCloudFormationTemplateRegistry, CloudFormationTemplateRegistry } from '../fs/templateRegistry'
 import { getIdeProperties } from '../extensionUtilities'
 import { NoopWatcher } from '../fs/watchedFiles'
-import { createStarterTemplateFile } from './cloudformation'
+import { CloudFormation, createStarterTemplateFile } from './cloudformation'
 import { Commands } from '../vscode/commands2'
 import globals from '../extensionGlobals'
 import { SamCliSettings } from '../sam/cli/samCliSettings'
 
-export const templateFileGlobPattern = '**/*.{yaml,yml}'
-
-/**
- * Match any file path that contains a .aws-sam folder. The way this works is:
- * match anything that starts  with a '/' or '\', then '.aws-sam', then either
- * a '/' or '\' followed by any number of characters or end of a string (so it
- * matches both /.aws-sam or /.aws-sam/<any number of characters>)
- */
-export const templateFileExcludePattern = /.*[/\\]\.aws-sam([/\\].*|$)/
-
-export const devfileExcludePattern = /.*devfile\.(yaml|yml)/i
-
 /**
  * Creates a CloudFormationTemplateRegistry which retains the state of CloudFormation templates in a workspace.
  * This also assigns a FileSystemWatcher which will update the registry on any change to tracked templates.
@@ -68,9 +56,9 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi
  */
 function setTemplateRegistryInGlobals(registry: CloudFormationTemplateRegistry) {
     const registrySetupFunc = async (registry: CloudFormationTemplateRegistry) => {
-        await registry.addExcludedPattern(devfileExcludePattern)
-        await registry.addExcludedPattern(templateFileExcludePattern)
-        await registry.addWatchPatterns([templateFileGlobPattern])
+        await registry.addExcludedPattern(CloudFormation.devfileExcludePattern)
+        await registry.addExcludedPattern(CloudFormation.templateFileExcludePattern)
+        await registry.addWatchPatterns([CloudFormation.templateFileGlobPattern])
         await registry.watchUntitledFiles()
 
         return registry
diff --git a/src/shared/cloudformation/cloudformation.ts b/src/shared/cloudformation/cloudformation.ts
@@ -13,12 +13,23 @@ import { SystemUtilities } from '../systemUtilities'
 import { getLogger } from '../logger'
 import { lambdaPackageTypeImage } from '../constants'
 import { isCloud9 } from '../extensionUtilities'
+import { isUntitledScheme, normalizeVSCodeUri } from '../utilities/vsCodeUtils'
 
 export namespace CloudFormation {
     export const SERVERLESS_API_TYPE = 'AWS::Serverless::Api' // eslint-disable-line @typescript-eslint/naming-convention
     export const SERVERLESS_FUNCTION_TYPE = 'AWS::Serverless::Function' // eslint-disable-line @typescript-eslint/naming-convention
     export const LAMBDA_FUNCTION_TYPE = 'AWS::Lambda::Function' // eslint-disable-line @typescript-eslint/naming-convention
 
+    export const templateFileGlobPattern = '**/*.{yaml,yml}'
+    export const devfileExcludePattern = /.*devfile\.(yaml|yml)/i
+    /**
+     * Match any file path that contains a .aws-sam folder. The way this works is:
+     * match anything that starts  with a '/' or '\', then '.aws-sam', then either
+     * a '/' or '\' followed by any number of characters or end of a string (so it
+     * matches both /.aws-sam or /.aws-sam/<any number of characters>)
+     */
+    export const templateFileExcludePattern = /.*[/\\]\.aws-sam([/\\].*|$)/
+
     export function isZipLambdaResource(
         resource?: ZipResourceProperties | ImageResourceProperties
     ): resource is ZipResourceProperties {
@@ -364,6 +375,22 @@ export namespace CloudFormation {
         [key: string]: Resource | undefined
     }
 
+    /** Returns true if the given name or path is a valid CloudFormation or SAM filename. */
+    export function isValidFilename(filename: string | vscode.Uri): boolean {
+        filename = typeof filename === 'string' ? filename : filename.fsPath
+        filename = filename.trim()
+        if (!filename.endsWith('.yml') && !filename.endsWith('.yaml')) {
+            return false
+        }
+        // Note: intentionally _not_ checking `templateFileExcludePattern` here, because while excluding
+        // template files in .aws-sam/ is relevant for the workspace scan, it's irrelevant if such
+        // a file was opened explicitly by the user.
+        if (filename.match(devfileExcludePattern)) {
+            return false
+        }
+        return true
+    }
+
     export async function load(filename: string, validate: boolean = true): Promise<Template> {
         if (!(await SystemUtilities.fileExists(filename))) {
             throw new Error(`Template file not found: ${filename}`)
@@ -384,6 +411,51 @@ export namespace CloudFormation {
         return template
     }
 
+    /**
+     * Returns a `CloudFormation.Template` if the given file (on disk) or `contents` is a valid
+     * CloudFormation document, or `{ template: undefined, kind: undefined }` if the file is
+     * invalid.
+     */
+    export async function tryLoad(
+        uri: vscode.Uri,
+        contents?: string
+    ): Promise<{ template: CloudFormation.Template | undefined; kind: 'sam' | 'cfn' | undefined }> {
+        const rv: {
+            template: CloudFormation.Template | undefined
+            kind: 'sam' | 'cfn' | undefined
+        } = { template: undefined, kind: undefined }
+        try {
+            if (isUntitledScheme(uri)) {
+                if (!contents) {
+                    // this error technically just throw us into the catch so the error message isn't used
+                    throw new Error('Contents must be defined for untitled uris')
+                }
+                rv.template = await CloudFormation.loadByContents(contents, false)
+            } else {
+                rv.template = await CloudFormation.load(normalizeVSCodeUri(uri), false)
+            }
+        } catch (e) {
+            return {
+                template: undefined,
+                kind: undefined,
+            }
+        }
+
+        // Check if the template is a SAM template, using the same heuristic as the cfn-lint team:
+        // https://github.com/aws-cloudformation/aws-cfn-lint-visual-studio-code/blob/629de0bac4f36cfc6534e409a6f6766a2240992f/client/src/yaml-support/yaml-schema.ts#L39-L51
+        if (rv.template.AWSTemplateFormatVersion || rv.template.Resources) {
+            rv.kind =
+                rv.template.Transform && rv.template.Transform.toString().startsWith('AWS::Serverless') ? 'sam' : 'cfn'
+
+            return rv
+        }
+
+        return {
+            template: undefined,
+            kind: undefined,
+        }
+    }
+
     export async function save(template: Template, filename: string): Promise<void> {
         const templateAsYaml: string = yaml.dump(template)
 
diff --git a/src/shared/fs/templateRegistry.ts b/src/shared/fs/templateRegistry.ts
@@ -14,50 +14,30 @@ import { getLambdaDetails } from '../../lambda/utils'
 import { 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'
 import { PerfLog } from '../logger/logger'
 
 export class CloudFormationTemplateRegistry extends WatchedFiles<CloudFormation.Template> {
     protected name: string = 'CloudFormationTemplateRegistry'
+
     protected async process(uri: vscode.Uri, contents?: string): Promise<CloudFormation.Template | undefined> {
         // P0: Assume all template.yaml/yml files are CFN templates and assign correct JSON schema.
         // P1: Alter registry functionality to search ALL YAML files and apply JSON schemas + add to registry based on validity
 
-        let template: CloudFormation.Template | undefined
-        const path = normalizeVSCodeUri(uri)
-        try {
-            if (isUntitledScheme(uri)) {
-                if (!contents) {
-                    // this error technically just throw us into the catch so the error message isn't used
-                    throw new Error('Contents must be defined for untitled uris')
-                }
-                template = await CloudFormation.loadByContents(contents, false)
-            } else {
-                template = await CloudFormation.load(path, false)
-            }
-        } catch (e) {
+        const r = await CloudFormation.tryLoad(uri, contents)
+        if (r.kind === undefined) {
             globals.schemaService.registerMapping({ uri, type: 'yaml', schema: undefined })
             return undefined
         }
 
-        // same heuristic used by cfn-lint team:
-        // https://github.com/aws-cloudformation/aws-cfn-lint-visual-studio-code/blob/629de0bac4f36cfc6534e409a6f6766a2240992f/client/src/yaml-support/yaml-schema.ts#L39-L51
-        if (template.AWSTemplateFormatVersion || template.Resources) {
-            if (template.Transform && template.Transform.toString().startsWith('AWS::Serverless')) {
-                // apply serverless schema
-                globals.schemaService.registerMapping({ uri, type: 'yaml', schema: 'sam' })
-            } else {
-                // apply cfn schema
-                globals.schemaService.registerMapping({ uri, type: 'yaml', schema: 'cfn' })
-            }
-
-            return template
+        if (r.kind === 'sam') {
+            globals.schemaService.registerMapping({ uri, type: 'yaml', schema: 'sam' })
+        } else if (r.kind === 'cfn') {
+            globals.schemaService.registerMapping({ uri, type: 'yaml', schema: 'cfn' })
         }
 
-        globals.schemaService.registerMapping({ uri, type: 'yaml', schema: undefined })
-        return undefined
+        return r.template
     }
 
     // handles delete case
@@ -123,7 +103,7 @@ export class AsyncCloudFormationTemplateRegistry {
                     location: vscode.ProgressLocation.Notification,
                     title: localize(
                         'AWS.codelens.waitingForTemplateRegistry',
-                        'Scanning CloudFormation templates... (except paths configured in [search.exclude](command:workbench.action.openSettings?"@id:search.exclude"))'
+                        'Scanning CloudFormation templates (except [search.exclude](command:workbench.action.openSettings?"@id:search.exclude") paths)...'
                     ),
                     cancellable: true,
                 },
diff --git a/src/test/shared/awsFiletypes.test.ts b/src/test/shared/awsFiletypes.test.ts
@@ -39,7 +39,6 @@ describe('awsFiletypes', function () {
     })
 
     it('emit telemetry when opened by user', async function () {
-        await (await globals.templateRegistry).addItem(cfnUri!)
         await vscode.commands.executeCommand('vscode.open', cfnUri)
         await vscode.commands.executeCommand('vscode.open', awsConfigUri)
         await vscode.workspace.openTextDocument({
@@ -61,16 +60,13 @@ describe('awsFiletypes', function () {
 
         assert(r, 'did not emit expected telemetry')
         assert(r.length === 3, 'emitted file_editAwsFile too many times')
-        const m1filetype = r[0].Metadata?.find(o => o.Key === 'awsFiletype')?.Value
-        const m2filetype = r[1].Metadata?.find(o => o.Key === 'awsFiletype')?.Value
-        const m3filetype = r[2].Metadata?.find(o => o.Key === 'awsFiletype')?.Value
-        assert.strictEqual(m1filetype, 'cloudformationSam')
-        assert.strictEqual(m2filetype, 'awsCredentials')
-        assert.strictEqual(m3filetype, 'ssmDocument')
+        const metrics = r.map(o => o.Metadata?.find(o => o.Key === 'awsFiletype')?.Value)
+        // The order is arbitrary (decided by vscode event system).
+        metrics.sort()
+        assert.deepStrictEqual(metrics, ['awsCredentials', 'cloudformationSam', 'ssmDocument'])
     })
 
     it('emit telemetry exactly once per filetype in a given flush window', async function () {
-        await (await globals.templateRegistry).addItem(cfnUri!)
         await vscode.commands.executeCommand('vscode.open', cfnUri)
         async function getMetrics() {
             return await waitUntil(
diff --git a/src/test/shared/cloudformation/cloudformation.test.ts b/src/test/shared/cloudformation/cloudformation.test.ts
@@ -34,6 +34,18 @@ describe('CloudFormation', function () {
         await fs.remove(filename)
     })
 
+    it('isValidFilename()', async function () {
+        assert.deepStrictEqual(CloudFormation.isValidFilename('/foo/bar.yaml'), true)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('/foo/template.yaml'), true)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('template.yaml'), true)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('template.yml'), true)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('/.aws-sam/template.yaml'), true)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('devfile.yml'), false)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('devfile.yaml'), false)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('template.yml.bk'), false)
+        assert.deepStrictEqual(CloudFormation.isValidFilename('template.txt'), false)
+    })
+
     describe('load', async function () {
         it('can successfully load a file', async function () {
             const yamlStr = makeSampleSamTemplateYaml(true)
diff --git a/src/test/shared/debug/launchConfiguration.test.ts b/src/test/shared/debug/launchConfiguration.test.ts
@@ -17,8 +17,8 @@ import { AwsSamDebuggerConfiguration } from '../../../shared/sam/debugger/awsSam
 import { AwsSamDebugConfigurationValidator } from '../../../shared/sam/debugger/awsSamDebugConfigurationValidator'
 import * as pathutils from '../../../shared/utilities/pathUtils'
 import * as testutil from '../../testUtil'
-import { templateFileGlobPattern } from '../../../shared/cloudformation/activation'
 import globals from '../../../shared/extensionGlobals'
+import { CloudFormation } from '../../../shared/cloudformation/cloudformation'
 
 const samDebugConfiguration: AwsSamDebuggerConfiguration = {
     type: 'aws-sam',
@@ -106,7 +106,7 @@ describe('LaunchConfiguration', function () {
 
     beforeEach(async function () {
         const registry = await globals.templateRegistry
-        await registry.addWatchPatterns([templateFileGlobPattern])
+        await registry.addWatchPatterns([CloudFormation.templateFileGlobPattern])
 
         // TODO: remove mocks in favor of testing src/testFixtures/ data.
         mockConfigSource = mock()
