Adding max age metrics for schemas (#197)

Author: satyakigh

.github/workflows/alpha-release.yml

@@ -2,8 +2,6 @@ name: Release Alpha
 run-name: Release Alpha ${{ github.actor }} ${{ github.event_name }}
 
 on:
-  schedule:
-    - cron: '0 8 * * 1,3'  # Monday, Wednesday 8am
   workflow_dispatch:
     inputs:
       ref:

src/featureFlag/FeatureFlagProvider.ts

@@ -1,6 +1,6 @@
 import { readFileSync, writeFileSync } from 'fs';
 import { join } from 'path';
-import axios from 'axios';
+import { downloadJson } from '../schema/RemoteSchemaHelper';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { Measure } from '../telemetry/TelemetryDecorator';
 import { Closeable } from '../utils/Closeable';
@@ -58,12 +58,9 @@ export class FeatureFlagProvider implements Closeable {
 
     @Measure({ name: 'getFromOnline' })
     private async getFromOnline(env: string): Promise<unknown> {
-        const response = await axios<unknown>({
-            method: 'get',
-            url: `https://raw.githubusercontent.com/aws-cloudformation/cloudformation-languageserver/refs/head/main/assets/featureFlag/${env.toLowerCase()}.json`,
-        });
-
-        return response.data;
+        return await downloadJson(
+            `https://raw.githubusercontent.com/aws-cloudformation/cloudformation-languageserver/refs/head/main/assets/featureFlag/${env.toLowerCase()}.json`,
+        );
     }
 
     private log() {

src/schema/GetSamSchemaTask.ts

@@ -1,24 +1,23 @@
 import { DataStore } from '../datastore/DataStore';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
 import { Measure } from '../telemetry/TelemetryDecorator';
-import { extractErrorMessage } from '../utils/Errors';
-import { downloadFile } from './RemoteSchemaHelper';
-import { SamSchemas, SamSchemasType } from './SamSchemas';
+import { GetSchemaTask } from './GetSchemaTask';
+import { downloadJson } from './RemoteSchemaHelper';
+import { SamSchemas, SamSchemasType, SamStoreKey } from './SamSchemas';
 import { SamSchemaTransformer, SamSchema } from './SamSchemaTransformer';
 
 const logger = LoggerFactory.getLogger('GetSamSchemaTask');
 
-export class GetSamSchemaTask {
+export class GetSamSchemaTask extends GetSchemaTask {
     private static readonly SAM_SCHEMA_URL =
         'https://raw.githubusercontent.com/aws/serverless-application-model/refs/heads/main/schema_source/sam.schema.json';
 
-    @Measure({ name: 'getSamSchema' })
-    async run(dataStore: DataStore): Promise<void> {
+    @Measure({ name: 'getSchemas' })
+    override async runImpl(dataStore: DataStore): Promise<void> {
         try {
             logger.info('Downloading SAM schema');
 
-            const schemaBuffer = await downloadFile(GetSamSchemaTask.SAM_SCHEMA_URL);
-            const samSchema = JSON.parse(schemaBuffer.toString()) as Record<string, unknown>;
+            const samSchema = await downloadJson<Record<string, unknown>>(GetSamSchemaTask.SAM_SCHEMA_URL);
 
             const resourceSchemas = SamSchemaTransformer.transformSamSchema(samSchema as unknown as SamSchema);
 
@@ -36,11 +35,11 @@ export class GetSamSchemaTask {
                 lastModifiedMs: Date.now(),
             };
 
-            await dataStore.put('sam-schemas', samSchemasData);
+            await dataStore.put(SamStoreKey, samSchemasData);
 
             logger.info(`Downloaded and stored ${resourceSchemas.size} SAM resource schemas`);
         } catch (error) {
-            logger.error({ error: extractErrorMessage(error) }, 'Failed to download SAM schema');
+            logger.error(error, 'Failed to download SAM schema');
             throw error;
         }
     }

src/schema/GetSchemaTask.ts

@@ -8,7 +8,7 @@ import { PrivateSchemas, PrivateSchemasType } from './PrivateSchemas';
 import { RegionalSchemas, RegionalSchemasType, SchemaFileType } from './RegionalSchemas';
 import { cfnResourceSchemaLink, downloadFile, unZipFile } from './RemoteSchemaHelper';
 
-abstract class GetSchemaTask {
+export abstract class GetSchemaTask {
     protected abstract runImpl(dataStore: DataStore, logger?: Logger): Promise<void>;
 
     async run(dataStore: DataStore, logger?: Logger) {

src/schema/GetSchemaTaskManager.ts

@@ -36,7 +36,6 @@ export class GetSchemaTaskManager implements SettingsConfigurable, Closeable {
         this.timeout = setTimeout(() => {
             // Wait before trying to call CFN APIs so that credentials have time to update
             this.runPrivateTask();
-            void this.runSamTask();
         }, TenSeconds);
 
         this.interval = setInterval(() => {
@@ -78,13 +77,11 @@ export class GetSchemaTaskManager implements SettingsConfigurable, Closeable {
             .catch(() => {});
     }
 
-    private async runSamTask(): Promise<void> {
-        try {
-            await this.samTask.run(this.schemas.publicSchemas);
-            this.schemas.invalidateCombinedSchemas();
-        } catch (error) {
-            this.log.error({ error }, 'Failed to run SAM schema task');
-        }
+    runSamTask() {
+        this.samTask
+            .run(this.schemas.samSchemas, this.log)
+            .then(() => this.schemas.invalidateCombinedSchemas())
+            .catch(() => {});
     }
 
     public currentRegionalTasks() {

src/schema/RemoteSchemaHelper.ts

@@ -82,3 +82,12 @@ export async function unZipFile(buffer: Promise<Buffer>): Promise<SchemaFileType
         });
     });
 }
+
+export async function downloadJson<T = unknown>(url: string): Promise<T> {
+    const response = await axios<T>({
+        method: 'get',
+        url: url,
+    });
+
+    return response.data;
+}

src/schema/SamSchemas.ts

@@ -7,6 +7,8 @@ export type SamSchemasType = {
     lastModifiedMs: number;
 };
 
+export const SamStoreKey = 'SamSchemas';
+
 export class SamSchemas {
     static readonly V1 = 'v1';
 

src/schema/SchemaRetriever.ts

@@ -2,11 +2,13 @@ import { DateTime } from 'luxon';
 import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
 import { DefaultSettings, ProfileSettings } from '../settings/Settings';
 import { LoggerFactory } from '../telemetry/LoggerFactory';
+import { TelemetryService } from '../telemetry/TelemetryService';
 import { Closeable } from '../utils/Closeable';
 import { AwsRegion, getRegion } from '../utils/Region';
 import { CombinedSchemas } from './CombinedSchemas';
 import { GetSchemaTaskManager } from './GetSchemaTaskManager';
 import { RegionalSchemasType } from './RegionalSchemas';
+import { SamSchemasType, SamStoreKey } from './SamSchemas';
 import { SchemaStore } from './SchemaStore';
 
 const StaleDaysThreshold = 7;
@@ -16,11 +18,18 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable {
     private settingsSubscription?: SettingsSubscription;
     private settings: ProfileSettings = DefaultSettings.profile;
     private readonly log = LoggerFactory.getLogger(SchemaRetriever);
+    private readonly telemetry = TelemetryService.instance.get('SchemaRetriever');
 
     constructor(
         private readonly schemaTaskManager: GetSchemaTaskManager,
         private readonly schemaStore: SchemaStore,
-    ) {}
+    ) {
+        this.telemetry.registerGaugeProvider('schema.public.maxAge', () => this.getPublicSchemaMaxAge(), {
+            unit: 'ms',
+        });
+
+        this.telemetry.registerGaugeProvider('schema.sam.maxAge', () => this.getSamSchemaAge(), { unit: 'ms' });
+    }
 
     configure(settingsManager: ISettingsSubscriber): void {
         // Clean up existing subscription if present
@@ -35,6 +44,7 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable {
         // Initialize schemas with current region
         this.getRegionalSchemasIfMissing([this.settings.region]);
         this.getRegionalSchemasIfStale();
+        this.getSamSchemasIfMissingOrStale();
 
         // Subscribe to profile settings changes
         this.settingsSubscription = settingsManager.subscribe('profile', (newProfileSettings) => {
@@ -125,4 +135,42 @@ export class SchemaRetriever implements SettingsConfigurable, Closeable {
             }
         }
     }
+
+    private getSamSchemasIfMissingOrStale() {
+        const existingValue = this.schemaStore.samSchemas.get<SamSchemasType>(SamStoreKey);
+
+        if (existingValue === undefined) {
+            this.schemaTaskManager.runSamTask();
+            return;
+        }
+
+        const now = DateTime.now();
+        const lastModified = DateTime.fromMillis(existingValue.lastModifiedMs);
+        const isStale = now.diff(lastModified, 'days').days >= StaleDaysThreshold;
+
+        if (isStale) {
+            this.schemaTaskManager.runSamTask();
+        }
+    }
+
+    private getPublicSchemaMaxAge(): number {
+        let maxAge = 0;
+        for (const key of this.schemaStore.publicSchemas.keys(50)) {
+            const region = getRegion(key);
+            const existingValue = this.getRegionalSchemasFromStore(region);
+            if (existingValue) {
+                const age = DateTime.now().diff(DateTime.fromMillis(existingValue.lastModifiedMs)).milliseconds;
+                maxAge = Math.max(maxAge, age);
+            }
+        }
+        return maxAge;
+    }
+
+    private getSamSchemaAge(): number {
+        const existingValue = this.schemaStore.samSchemas.get<SamSchemasType>(SamStoreKey);
+        if (!existingValue) {
+            return 0;
+        }
+        return DateTime.now().diff(DateTime.fromMillis(existingValue.lastModifiedMs)).milliseconds;
+    }
 }

src/schema/SchemaStore.ts

@@ -3,7 +3,7 @@ import { AwsRegion } from '../utils/Region';
 import { CombinedSchemas } from './CombinedSchemas';
 import { PrivateSchemasType } from './PrivateSchemas';
 import { RegionalSchemasType } from './RegionalSchemas';
-import { SamSchemasType } from './SamSchemas';
+import { SamSchemasType, SamStoreKey } from './SamSchemas';
 
 export class SchemaStore {
     public readonly publicSchemas = this.dataStoreFactory.get('public_schemas', Persistence.local);
@@ -20,7 +20,7 @@ export class SchemaStore {
         if (!cached) {
             const regionalSchemas = this.publicSchemas.get<RegionalSchemasType>(region);
             const privateSchemas = this.privateSchemas.get<PrivateSchemasType>(profile);
-            const samSchemas = this.samSchemas.get<SamSchemasType>('sam-schemas');
+            const samSchemas = this.samSchemas.get<SamSchemasType>(SamStoreKey);
 
             cached = CombinedSchemas.from(regionalSchemas, privateSchemas, samSchemas);
             void this.combinedSchemas.put(cacheKey, cached);

src/services/AwsClient.ts

@@ -9,7 +9,6 @@ type IamClientConfig = {
     region: string;
     credentials: IamCredentials;
     customUserAgent: string;
-    endpoint?: string;
 };
 
 export class AwsClient {
@@ -18,9 +17,11 @@ export class AwsClient {
         private readonly cloudformationEndpoint?: string,
     ) {}
 
-    // By default, clients will retry on throttling exceptions 3 times
     public getCloudFormationClient() {
-        return new CloudFormationClient(this.iamClientConfig());
+        return new CloudFormationClient({
+            ...this.iamClientConfig(),
+            endpoint: this.cloudformationEndpoint,
+        });
     }
 
     public getCloudControlClient() {
@@ -38,7 +39,6 @@ export class AwsClient {
                 region: credential.region,
                 credentials: credential,
                 customUserAgent: `${ExtensionId}/${ExtensionVersion}`,
-                endpoint: this.cloudformationEndpoint,
             };
         } catch {
             throw new Error('AWS credentials not configured. Authentication required for online features.');