Skip to content

Commit 211511b

Browse files
satyakighsatyakigh
committed
Refactor schema workflows to avoid extra operations and duplicate schema load calls
1 parent 98c4bf2 commit 211511b

16 files changed

+872
-628
lines changed

‎src/datastore/DataStore.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ export enum StoreName {
1414
public_schemas = 'public_schemas',
1515
sam_schemas = 'sam_schemas',
1616
private_schemas = 'private_schemas',
17-
combined_schemas = 'combined_schemas',
1817
}
1918

2019
export interface DataStore {

‎src/handlers/Initialize.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ export function initializedHandler(workspace: LspWorkspace, components: ServerCo
1010
components.settingsManager
1111
.syncConfiguration()
1212
.then(() => {
13+
components.schemaRetriever.initialize();
1314
return components.cfnLintService.initialize();
1415
})
1516
.then(async () => {

‎src/schema/CombinedSchemas.ts‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
1-
import { LoggerFactory } from '../telemetry/LoggerFactory';
21
import { PrivateSchemas, PrivateSchemasType } from './PrivateSchemas';
32
import { RegionalSchemas, RegionalSchemasType } from './RegionalSchemas';
43
import { ResourceSchema } from './ResourceSchema';
54
import { SamSchemas, SamSchemasType } from './SamSchemas';
65

76
export class CombinedSchemas {
8-
private static readonly log = LoggerFactory.getLogger('CombinedSchemas');
97
readonly numSchemas: number;
108
readonly schemas: Map<string, ResourceSchema>;
119

12-
constructor(regionalSchemas?: RegionalSchemas, privateSchemas?: PrivateSchemas, samSchemas?: SamSchemas) {
10+
constructor(
11+
readonly regionalSchemas?: RegionalSchemas,
12+
readonly privateSchemas?: PrivateSchemas,
13+
readonly samSchemas?: SamSchemas,
14+
) {
1315
this.schemas = new Map<string, ResourceSchema>([
1416
...(privateSchemas?.schemas ?? []),
1517
...(regionalSchemas?.schemas ?? []),
@@ -22,14 +24,10 @@ export class CombinedSchemas {
2224
regionalSchemas?: RegionalSchemasType,
2325
privateSchemas?: PrivateSchemasType,
2426
samSchemas?: SamSchemasType,
25-
) {
27+
): CombinedSchemas {
2628
const regionalSchema = regionalSchemas === undefined ? undefined : RegionalSchemas.from(regionalSchemas);
2729
const privateSchema = privateSchemas === undefined ? undefined : PrivateSchemas.from(privateSchemas);
2830
const samSchema = samSchemas === undefined ? undefined : SamSchemas.from(samSchemas);
29-
30-
CombinedSchemas.log.info(
31-
`Combined schemas from public=${regionalSchemas?.schemas.length}, private=${privateSchema?.schemas.size}, SAM=${samSchema?.schemas.size}`,
32-
);
3331
return new CombinedSchemas(regionalSchema, privateSchema, samSchema);
3432
}
3533
}

‎src/schema/GetSamSchemaTask.ts‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,20 @@
1-
import { Logger } from 'pino';
21
import { DataStore } from '../datastore/DataStore';
2+
import { LoggerFactory } from '../telemetry/LoggerFactory';
33
import { Measure } from '../telemetry/TelemetryDecorator';
44
import { downloadJson } from '../utils/RemoteDownload';
55
import { GetSchemaTask } from './GetSchemaTask';
66
import { SamSchemas, SamSchemasType, SamStoreKey } from './SamSchemas';
77
import { CloudFormationResourceSchema, SamSchema, SamSchemaTransformer } from './SamSchemaTransformer';
88

99
export class GetSamSchemaTask extends GetSchemaTask {
10+
private readonly logger = LoggerFactory.getLogger(GetSamSchemaTask);
11+
1012
constructor(private readonly getSamSchemas: () => Promise<Map<string, CloudFormationResourceSchema>>) {
1113
super();
1214
}
1315

1416
@Measure({ name: 'getSchemas' })
15-
protected override async runImpl(dataStore: DataStore, logger?: Logger): Promise<void> {
17+
protected override async runImpl(dataStore: DataStore): Promise<void> {
1618
try {
1719
const resourceSchemas = await this.getSamSchemas();
1820

@@ -32,9 +34,9 @@ export class GetSamSchemaTask extends GetSchemaTask {
3234

3335
await dataStore.put(SamStoreKey, samSchemasData);
3436

35-
logger?.info(`${resourceSchemas.size} SAM schemas downloaded and stored`);
37+
this.logger.info(`${resourceSchemas.size} SAM schemas downloaded and stored`);
3638
} catch (error) {
37-
logger?.error(error, 'Failed to download SAM schema');
39+
this.logger.error(error, 'Failed to download SAM schema');
3840
throw error;
3941
}
4042
}

‎src/schema/GetSchemaTask.ts‎

Lines changed: 28 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,27 @@
11
import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation';
2-
import { Logger } from 'pino';
32
import { AwsCredentials } from '../auth/AwsCredentials';
43
import { DataStore } from '../datastore/DataStore';
54
import { CfnService } from '../services/CfnService';
5+
import { LoggerFactory } from '../telemetry/LoggerFactory';
66
import { ScopedTelemetry } from '../telemetry/ScopedTelemetry';
77
import { Measure, Telemetry } from '../telemetry/TelemetryDecorator';
88
import { AwsRegion } from '../utils/Region';
99
import { downloadFile } from '../utils/RemoteDownload';
10-
import { PrivateSchemas, PrivateSchemasType } from './PrivateSchemas';
10+
import { PrivateSchemas, PrivateSchemasType, PrivateStoreKey } from './PrivateSchemas';
1111
import { RegionalSchemas, RegionalSchemasType, SchemaFileType } from './RegionalSchemas';
1212
import { cfnResourceSchemaLink, unZipFile } from './RemoteSchemaHelper';
1313

1414
export abstract class GetSchemaTask {
15-
protected abstract runImpl(dataStore: DataStore, logger?: Logger): Promise<void>;
15+
protected abstract runImpl(dataStore: DataStore): Promise<void>;
1616

17-
async run(dataStore: DataStore, logger?: Logger) {
18-
await this.runImpl(dataStore, logger);
17+
async run(dataStore: DataStore) {
18+
await this.runImpl(dataStore);
1919
}
2020
}
2121

2222
export class GetPublicSchemaTask extends GetSchemaTask {
23+
private readonly logger = LoggerFactory.getLogger(GetPublicSchemaTask);
24+
2325
@Telemetry()
2426
private readonly telemetry!: ScopedTelemetry;
2527

@@ -35,9 +37,20 @@ export class GetPublicSchemaTask extends GetSchemaTask {
3537
}
3638

3739
@Measure({ name: 'getSchemas' })
38-
protected override async runImpl(dataStore: DataStore, logger?: Logger) {
40+
protected override async runImpl(dataStore: DataStore) {
41+
this.telemetry.count(`getSchemas.maxAttempt.fault`, 0, {
42+
attributes: {
43+
region: this.region,
44+
},
45+
});
46+
3947
if (this.attempts >= GetPublicSchemaTask.MaxAttempts) {
40-
logger?.error(`Reached max attempts for retrieving schemas for ${this.region} without success`);
48+
this.telemetry.count(`getSchemas.maxAttempt.fault`, 1, {
49+
attributes: {
50+
region: this.region,
51+
},
52+
});
53+
this.logger.error(`Reached max attempts for retrieving schemas for ${this.region} without success`);
4154
return;
4255
}
4356

@@ -53,44 +66,35 @@ export class GetPublicSchemaTask extends GetSchemaTask {
5366
};
5467

5568
await dataStore.put<RegionalSchemasType>(this.region, value);
56-
logger?.info(`${schemas.length} public schemas retrieved for ${this.region}`);
69+
this.logger.info(`${schemas.length} public schemas retrieved for ${this.region}`);
5770
}
5871
}
5972

6073
export class GetPrivateSchemasTask extends GetSchemaTask {
61-
private readonly processedProfiles = new Set<string>();
74+
private readonly logger = LoggerFactory.getLogger(GetPrivateSchemasTask);
6275

63-
constructor(
64-
private readonly getSchemas: () => Promise<DescribeTypeOutput[]>,
65-
private readonly getProfile: () => string,
66-
) {
76+
constructor(private readonly getSchemas: () => Promise<DescribeTypeOutput[]>) {
6777
super();
6878
}
6979

7080
@Measure({ name: 'getSchemas' })
71-
protected override async runImpl(dataStore: DataStore, logger?: Logger) {
81+
protected override async runImpl(dataStore: DataStore) {
7282
try {
73-
const profile = this.getProfile();
74-
if (this.processedProfiles.has(profile)) {
75-
return;
76-
}
77-
7883
const schemas: DescribeTypeOutput[] = await this.getSchemas();
7984

8085
const value: PrivateSchemasType = {
8186
version: PrivateSchemas.V1,
82-
identifier: profile,
87+
identifier: PrivateStoreKey,
8388
schemas: schemas,
8489
firstCreatedMs: Date.now(),
8590
lastModifiedMs: Date.now(),
8691
};
8792

88-
await dataStore.put<PrivateSchemasType>(profile, value);
93+
await dataStore.put<PrivateSchemasType>(PrivateStoreKey, value);
8994

90-
this.processedProfiles.add(profile);
91-
logger?.info(`${schemas.length} private schemas retrieved`);
95+
this.logger.info(`${schemas.length} private schemas retrieved`);
9296
} catch (error) {
93-
logger?.error(error, `Failed to get private schemas`);
97+
this.logger.error(error, 'Failed to get private schemas');
9498
throw error;
9599
}
96100
}

‎src/schema/GetSchemaTaskManager.ts‎

Lines changed: 36 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -1,132 +1,73 @@
11
import { DescribeTypeOutput } from '@aws-sdk/client-cloudformation';
2-
import { SettingsConfigurable, ISettingsSubscriber, SettingsSubscription } from '../settings/ISettingsSubscriber';
3-
import { DefaultSettings } from '../settings/Settings';
42
import { LoggerFactory } from '../telemetry/LoggerFactory';
5-
import { Closeable } from '../utils/Closeable';
6-
import { AwsRegion } from '../utils/Region';
3+
import { AwsRegion, getRegion } from '../utils/Region';
74
import { GetSamSchemaTask } from './GetSamSchemaTask';
85
import { GetPrivateSchemasTask, GetPublicSchemaTask } from './GetSchemaTask';
96
import { SchemaFileType } from './RegionalSchemas';
107
import { CloudFormationResourceSchema } from './SamSchemaTransformer';
118
import { SchemaStore } from './SchemaStore';
129

13-
const TenSeconds = 10 * 1000;
14-
const OneHour = 60 * 60 * 1000;
15-
16-
export class GetSchemaTaskManager implements SettingsConfigurable, Closeable {
10+
export class GetSchemaTaskManager {
11+
private readonly processedRegions = new Set<AwsRegion>();
1712
private readonly tasks: GetPublicSchemaTask[] = [];
1813
private readonly privateTask: GetPrivateSchemasTask;
1914
private readonly samTask: GetSamSchemaTask;
20-
private settingsSubscription?: SettingsSubscription;
2115
private readonly log = LoggerFactory.getLogger(GetSchemaTaskManager);
22-
23-
private isRunning: boolean = false;
24-
25-
private readonly timeout: NodeJS.Timeout;
26-
private readonly interval: NodeJS.Timeout;
16+
private isRunning = false;
2717

2818
constructor(
2919
private readonly schemas: SchemaStore,
3020
private readonly getPublicSchemas: (region: AwsRegion) => Promise<SchemaFileType[]>,
3121
getPrivateResources: () => Promise<DescribeTypeOutput[]>,
3222
getSamSchemas: () => Promise<Map<string, CloudFormationResourceSchema>>,
33-
private profile: string = DefaultSettings.profile.profile,
34-
private readonly onSchemaUpdate: (region?: string, profile?: string) => void,
3523
) {
36-
this.privateTask = new GetPrivateSchemasTask(getPrivateResources, () => this.profile);
24+
this.privateTask = new GetPrivateSchemasTask(getPrivateResources);
3725
this.samTask = new GetSamSchemaTask(getSamSchemas);
38-
39-
this.timeout = setTimeout(() => {
40-
// Wait before trying to call CFN APIs so that credentials have time to update
41-
this.runPrivateTask();
42-
}, TenSeconds);
43-
44-
this.interval = setInterval(() => {
45-
// Keep private schemas up to date with credential changes if profile has not already ben loaded
46-
this.runPrivateTask();
47-
}, OneHour);
4826
}
4927

50-
configure(settingsManager: ISettingsSubscriber): void {
51-
// Clean up existing subscription if present
52-
if (this.settingsSubscription) {
53-
this.settingsSubscription.unsubscribe();
54-
}
55-
56-
// Set initial settings
57-
this.profile = settingsManager.getCurrentSettings().profile.profile;
28+
addTask(reg: string, regionFirstCreatedMs?: number) {
29+
const region = getRegion(reg);
5830

59-
// Subscribe to profile settings changes
60-
this.settingsSubscription = settingsManager.subscribe('profile', (newProfileSettings) => {
61-
this.onSettingsChanged(newProfileSettings.profile);
62-
});
63-
}
31+
if (!this.processedRegions.has(region)) {
32+
this.tasks.push(new GetPublicSchemaTask(region, this.getPublicSchemas, regionFirstCreatedMs));
33+
this.processedRegions.add(region);
34+
}
6435

65-
private onSettingsChanged(newProfile: string): void {
66-
this.profile = newProfile;
36+
if (!this.isRunning) {
37+
this.runNextTask();
38+
}
6739
}
6840

69-
addTask(region: AwsRegion, regionFirstCreatedMs?: number) {
70-
if (!this.currentRegionalTasks().has(region)) {
71-
this.tasks.push(new GetPublicSchemaTask(region, this.getPublicSchemas, regionFirstCreatedMs));
41+
private runNextTask() {
42+
const task = this.tasks.shift();
43+
if (!task) {
44+
this.isRunning = false;
45+
return;
7246
}
73-
this.startProcessing();
47+
48+
this.isRunning = true;
49+
task.run(this.schemas.publicSchemas)
50+
.catch((err) => {
51+
this.log.error(err);
52+
this.tasks.push(task);
53+
})
54+
.finally(() => {
55+
this.isRunning = false;
56+
this.runNextTask();
57+
});
7458
}
7559

7660
runPrivateTask() {
7761
this.privateTask
78-
.run(this.schemas.privateSchemas, this.log)
79-
.then(() => {
80-
this.onSchemaUpdate(undefined, this.profile);
81-
})
82-
.catch(() => {});
62+
.run(this.schemas.privateSchemas)
63+
.then(() => this.schemas.invalidate())
64+
.catch(this.log.error);
8365
}
8466

8567
runSamTask() {
8668
this.samTask
87-
.run(this.schemas.samSchemas, this.log)
88-
.then(() => {
89-
this.onSchemaUpdate(); // No params = SAM update
90-
})
91-
.catch(() => {});
92-
}
93-
94-
public currentRegionalTasks() {
95-
return new Set(this.tasks.map((task) => task.region));
96-
}
97-
98-
private startProcessing() {
99-
if (!this.isRunning && this.tasks.length > 0) {
100-
this.isRunning = true;
101-
this.run();
102-
}
103-
}
104-
105-
private run() {
106-
const task = this.tasks.shift();
107-
if (task) {
108-
task.run(this.schemas.publicSchemas, this.log)
109-
.then(() => {
110-
this.onSchemaUpdate(task.region);
111-
})
112-
.catch(() => {
113-
this.tasks.push(task);
114-
})
115-
.finally(() => {
116-
this.isRunning = false;
117-
this.startProcessing();
118-
});
119-
}
120-
}
121-
122-
public close() {
123-
// Unsubscribe from settings changes
124-
if (this.settingsSubscription) {
125-
this.settingsSubscription.unsubscribe();
126-
this.settingsSubscription = undefined;
127-
}
128-
129-
clearTimeout(this.timeout);
130-
clearInterval(this.interval);
69+
.run(this.schemas.samSchemas)
70+
.then(() => this.schemas.invalidate())
71+
.catch(this.log.error);
13172
}
13273
}

‎src/schema/PrivateSchemas.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ export type PrivateSchemasType = {
99
lastModifiedMs: number;
1010
};
1111

12+
export const PrivateStoreKey = 'PrivateSchemas';
1213
export class PrivateSchemas {
1314
static readonly V1 = 'v1';
1415

‎src/schema/SamSchemas.ts‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,7 @@ export class SamSchemas {
1919
readonly lastModifiedMs: number;
2020
readonly schemas: Map<string, ResourceSchema>;
2121

22-
constructor(
23-
version: string,
24-
schemas: { name: string; content: string; createdMs: number }[],
25-
firstCreatedMs: number,
26-
lastModifiedMs: number,
27-
) {
22+
constructor(version: string, schemas: SchemaFileType[], firstCreatedMs: number, lastModifiedMs: number) {
2823
this.version = version;
2924
this.firstCreatedMs = firstCreatedMs;
3025
this.lastModifiedMs = lastModifiedMs;

0 commit comments

Comments
 (0)