Skip to content

Commit ceef756

Browse files
linglingye001linglingye001
authored
Optimization for concurrent Key Vault Ref resolution - deduplicate requests to the same secret (#328)
* Added map to track pending key vault requests keyed by secret source id * dedup key vault ref in preload stage * update * update * resolved comments * update * update
1 parent 4f6d8b3 commit ceef756

7 files changed

Lines changed: 309 additions & 65 deletions

File tree

‎src/appConfigurationImpl.ts‎

Lines changed: 14 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,6 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
111111
#secretRefreshEnabled: boolean = false;
112112
#secretReferences: ConfigurationSetting[] = []; // cached key vault references
113113
#secretRefreshTimer: RefreshTimer | undefined = undefined;
114-
#resolveSecretsInParallel: boolean = false;
115114

116115
/**
117116
* Selectors of key-values obtained from @see AzureAppConfigurationOptions.selectors
@@ -203,7 +202,6 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
203202
this.#secretRefreshEnabled = true;
204203
this.#secretRefreshTimer = new RefreshTimer(secretRefreshIntervalInMs);
205204
}
206-
this.#resolveSecretsInParallel = options.keyVaultOptions.parallelSecretResolutionEnabled ?? false;
207205
}
208206
this.#adapters.push(new AzureKeyVaultKeyValueAdapter(options?.keyVaultOptions, this.#secretRefreshTimer));
209207
this.#adapters.push(new JsonKeyValueAdapter());
@@ -559,22 +557,19 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
559557
this.#aiConfigurationTracing.reset();
560558
}
561559

560+
for (const adapter of this.#adapters) {
561+
await adapter.preload?.(loadedSettings.filter(s => adapter.canProcess(s))); // dedup and warm the secret cache
562+
}
563+
562564
for (const setting of loadedSettings) {
563565
if (isSecretReference(setting)) {
564566
this.#secretReferences.push(setting); // cache secret references for resolve/refresh secret separately
565-
continue;
566567
}
567568
// adapt configuration settings to key-values
568569
const [key, value] = await this.#processKeyValue(setting);
569570
keyValues.push([key, value]);
570571
}
571572

572-
if (this.#secretReferences.length > 0) {
573-
await this.#resolveSecretReferences(this.#secretReferences, (key, value) => {
574-
keyValues.push([key, value]);
575-
});
576-
}
577-
578573
this.#clearLoadedKeyValues(); // clear existing key-values in case of configuration setting deletion
579574
for (const [k, v] of keyValues) {
580575
this.#configMap.set(k, v); // reset the configuration
@@ -672,7 +667,7 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
672667

673668
if (needRefresh) {
674669
for (const adapter of this.#adapters) {
675-
await adapter.onChangeDetected();
670+
await adapter.onChangeDetected?.();
676671
}
677672
await this.#loadSelectedKeyValues();
678673

@@ -716,9 +711,16 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
716711
return Promise.resolve(false);
717712
}
718713

719-
await this.#resolveSecretReferences(this.#secretReferences, (key, value) => {
714+
const keyVaultRefAdapter = this.#adapters.find(adapter => adapter instanceof AzureKeyVaultKeyValueAdapter) as AzureKeyVaultKeyValueAdapter | undefined;
715+
if (keyVaultRefAdapter) {
716+
// dedup and warm the secret cache
717+
await keyVaultRefAdapter.preload(this.#secretReferences);
718+
}
719+
720+
for (const setting of this.#secretReferences) {
721+
const [key, value] = await this.#processKeyValue(setting);
720722
this.#configMap.set(key, value);
721-
});
723+
}
722724

723725
this.#secretRefreshTimer.reset();
724726
return Promise.resolve(true);
@@ -894,27 +896,6 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
894896
throw new Error(ErrorMessages.ALL_FALLBACK_CLIENTS_FAILED);
895897
}
896898

897-
async #resolveSecretReferences(secretReferences: ConfigurationSetting[], resultHandler: (key: string, value: unknown) => void): Promise<void> {
898-
if (this.#resolveSecretsInParallel) {
899-
const secretResolutionPromises: Promise<void>[] = [];
900-
for (const setting of secretReferences) {
901-
const secretResolutionPromise = this.#processKeyValue(setting)
902-
.then(([key, value]) => {
903-
resultHandler(key, value);
904-
});
905-
secretResolutionPromises.push(secretResolutionPromise);
906-
}
907-
908-
// Wait for all secret resolution promises to be resolved
909-
await Promise.all(secretResolutionPromises);
910-
} else {
911-
for (const setting of secretReferences) {
912-
const [key, value] = await this.#processKeyValue(setting);
913-
resultHandler(key, value);
914-
}
915-
}
916-
}
917-
918899
async #processKeyValue(setting: ConfigurationSetting<string>): Promise<[string, unknown]> {
919900
this.#setAIConfigurationTracing(setting);
920901

‎src/common/errorMessages.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import { MIN_REFRESH_INTERVAL_IN_MS } from "../refresh/refreshOptions.js";
55
import { MIN_SECRET_REFRESH_INTERVAL_IN_MS } from "../keyvault/keyVaultOptions.js";
6+
import { ConfigurationSetting } from "@azure/app-configuration";
67

78
export const enum ErrorMessages {
89
INVALID_WATCHED_SETTINGS_KEY = "The characters '*' and ',' are not supported in key of watched settings.",
@@ -26,3 +27,14 @@ export const enum KeyVaultReferenceErrorMessages {
2627
KEY_VAULT_OPTIONS_UNDEFINED = "Failed to process the Key Vault reference because Key Vault options are not configured.",
2728
KEY_VAULT_REFERENCE_UNRESOLVABLE = "Failed to resolve the key vault reference. No key vault secret client, credential or secret resolver callback is available to resolve the secret."
2829
}
30+
31+
export function buildKeyVaultReferenceErrorMessage(message: string, secretIdentifier?: string, setting?: ConfigurationSetting): string {
32+
let errorMessage = message;
33+
if (secretIdentifier) {
34+
errorMessage += ` SecretIdentifier: '${secretIdentifier}'`;
35+
}
36+
if (setting) {
37+
errorMessage += ` Key: '${setting.key}' Label: '${setting.label ?? ""}' ETag: '${setting.etag ?? ""}'`;
38+
}
39+
return errorMessage;
40+
}

‎src/jsonKeyValueAdapter.ts‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,7 @@ export class JsonKeyValueAdapter implements IKeyValueAdapter {
4343
return [setting.key, parsedValue];
4444
}
4545

46-
async onChangeDetected(): Promise<void> {
47-
return;
48-
}
49-
50-
#tryParseJson(value: string): { success: true; result: unknown } | { success: false } {
46+
#tryParseJson(value: string): { success: true; result: unknown } | { success: false } {
5147
try {
5248
return { success: true, result: JSON.parse(value) };
5349
} catch (error) {

‎src/keyValueAdapter.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,13 @@ export interface IKeyValueAdapter {
1414
*/
1515
processKeyValue(setting: ConfigurationSetting): Promise<[string, unknown]>;
1616

17+
/**
18+
* This method deduplicates and warms up Key Vault secret requests so that processKeyValue only reads from cache.
19+
*/
20+
preload?(settings: ConfigurationSetting[]): Promise<void>;
21+
1722
/**
1823
* This method is called when a change is detected in the configuration setting.
1924
*/
20-
onChangeDetected(): Promise<void>;
25+
onChangeDetected?(): Promise<void>;
2126
}

‎src/keyvault/keyVaultKeyValueAdapter.ts‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,8 @@ import { AzureKeyVaultSecretProvider } from "./keyVaultSecretProvider.js";
77
import { KeyVaultOptions } from "./keyVaultOptions.js";
88
import { RefreshTimer } from "../refresh/refreshTimer.js";
99
import { ArgumentError, KeyVaultReferenceError } from "../common/errors.js";
10-
import { KeyVaultReferenceErrorMessages } from "../common/errorMessages.js";
10+
import { KeyVaultReferenceErrorMessages, buildKeyVaultReferenceErrorMessage } from "../common/errorMessages.js";
1111
import { KeyVaultSecretIdentifier, parseKeyVaultSecretIdentifier } from "@azure/keyvault-secrets";
12-
import { isRestError } from "@azure/core-rest-pipeline";
13-
import { AuthenticationError } from "@azure/identity";
1412

1513
export class AzureKeyVaultKeyValueAdapter implements IKeyValueAdapter {
1614
#keyVaultOptions: KeyVaultOptions | undefined;
@@ -35,17 +33,44 @@ export class AzureKeyVaultKeyValueAdapter implements IKeyValueAdapter {
3533
parseSecretReference(setting).value.secretId
3634
);
3735
} catch (error) {
38-
throw new KeyVaultReferenceError(buildKeyVaultReferenceErrorMessage("Invalid Key Vault reference.", setting), { cause: error });
36+
throw new KeyVaultReferenceError(buildKeyVaultReferenceErrorMessage("Invalid Key Vault reference.", undefined, setting), { cause: error });
3937
}
4038

41-
try {
42-
const secretValue = await this.#keyVaultSecretProvider.getSecretValue(secretIdentifier);
43-
return [setting.key, secretValue];
44-
} catch (error) {
45-
if (isRestError(error) || error instanceof AuthenticationError) {
46-
throw new KeyVaultReferenceError(buildKeyVaultReferenceErrorMessage("Failed to resolve Key Vault reference.", setting, secretIdentifier.sourceId), { cause: error });
39+
const secretValue = await this.#keyVaultSecretProvider.getSecretValue(secretIdentifier);
40+
return [setting.key, secretValue];
41+
}
42+
43+
async preload(settings: ConfigurationSetting[]): Promise<void> {
44+
if (!this.#keyVaultOptions) {
45+
return; // no-op when keyVaultOptions is not configured
46+
}
47+
// Deduplicate references by secret identifier (sourceId).
48+
const uniqueSecrets = new Map<string, KeyVaultSecretIdentifier>();
49+
for (const setting of settings) {
50+
if (!this.canProcess(setting)) {
51+
continue;
52+
}
53+
let secretIdentifier: KeyVaultSecretIdentifier;
54+
try {
55+
secretIdentifier = parseKeyVaultSecretIdentifier(
56+
parseSecretReference(setting).value.secretId
57+
);
58+
} catch (error) {
59+
throw new KeyVaultReferenceError(buildKeyVaultReferenceErrorMessage("Invalid Key Vault reference.", undefined, setting), { cause: error });
60+
}
61+
if (!uniqueSecrets.has(secretIdentifier.sourceId)) {
62+
uniqueSecrets.set(secretIdentifier.sourceId, secretIdentifier);
63+
}
64+
}
65+
66+
// Resolve failures surface as KeyVaultReferenceError from the provider, identified by secret identifier.
67+
const uniqueSecretIdentifiers = [...uniqueSecrets.values()];
68+
if (this.#keyVaultOptions.parallelSecretResolutionEnabled) {
69+
await Promise.all(uniqueSecretIdentifiers.map(secretIdentifier => this.#keyVaultSecretProvider.loadSecretValue(secretIdentifier)));
70+
} else {
71+
for (const secretIdentifier of uniqueSecretIdentifiers) {
72+
await this.#keyVaultSecretProvider.loadSecretValue(secretIdentifier);
4773
}
48-
throw error;
4974
}
5075
}
5176

@@ -54,7 +79,3 @@ export class AzureKeyVaultKeyValueAdapter implements IKeyValueAdapter {
5479
return;
5580
}
5681
}
57-
58-
function buildKeyVaultReferenceErrorMessage(message: string, setting: ConfigurationSetting, secretIdentifier?: string ): string {
59-
return `${message} Key: '${setting.key}' Label: '${setting.label ?? ""}' ETag: '${setting.etag ?? ""}' ${secretIdentifier ? ` SecretIdentifier: '${secretIdentifier}'` : ""}`;
60-
}

‎src/keyvault/keyVaultSecretProvider.ts‎

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@
33

44
import { KeyVaultOptions, MIN_SECRET_REFRESH_INTERVAL_IN_MS } from "./keyVaultOptions.js";
55
import { RefreshTimer } from "../refresh/refreshTimer.js";
6-
import { ArgumentError } from "../common/errors.js";
6+
import { ArgumentError, KeyVaultReferenceError } from "../common/errors.js";
77
import { SecretClient, KeyVaultSecretIdentifier } from "@azure/keyvault-secrets";
8-
import { KeyVaultReferenceErrorMessages } from "../common/errorMessages.js";
8+
import { KeyVaultReferenceErrorMessages, buildKeyVaultReferenceErrorMessage } from "../common/errorMessages.js";
9+
import { isRestError } from "@azure/core-rest-pipeline";
10+
import { AuthenticationError } from "@azure/identity";
911

1012
export class AzureKeyVaultSecretProvider {
1113
#keyVaultOptions: KeyVaultOptions | undefined;
@@ -33,21 +35,34 @@ export class AzureKeyVaultSecretProvider {
3335
}
3436
}
3537

36-
async getSecretValue(secretIdentifier: KeyVaultSecretIdentifier): Promise<unknown> {
38+
async loadSecretValue(secretIdentifier: KeyVaultSecretIdentifier): Promise<unknown> {
3739
const identifierKey = secretIdentifier.sourceId;
38-
39-
// If the refresh interval is not expired, return the cached value if available.
40-
if (this.#cachedSecretValues.has(identifierKey) &&
41-
(!this.#secretRefreshTimer || !this.#secretRefreshTimer.canRefresh())) {
42-
return this.#cachedSecretValues.get(identifierKey);
40+
const shouldRefresh = this.#secretRefreshTimer?.canRefresh() ?? false;
41+
if (this.#cachedSecretValues.has(identifierKey) && !shouldRefresh) {
42+
return this.#cachedSecretValues.get(identifierKey); // already cached and still fresh
43+
}
44+
let secretValue: unknown;
45+
try {
46+
secretValue = await this.#getSecretValueFromKeyVault(secretIdentifier);
47+
} catch (error) {
48+
if (isRestError(error) || error instanceof AuthenticationError) {
49+
throw new KeyVaultReferenceError(buildKeyVaultReferenceErrorMessage("Failed to resolve Key Vault reference.", secretIdentifier.sourceId), { cause: error });
50+
}
51+
throw error;
4352
}
44-
45-
// Fallback to fetching the secret value from Key Vault.
46-
const secretValue = await this.#getSecretValueFromKeyVault(secretIdentifier);
4753
this.#cachedSecretValues.set(identifierKey, secretValue);
4854
return secretValue;
4955
}
5056

57+
// Serves the secret value from cache when available; otherwise loads it from Key Vault on demand.
58+
async getSecretValue(secretIdentifier: KeyVaultSecretIdentifier): Promise<unknown> {
59+
const identifierKey = secretIdentifier.sourceId;
60+
if (this.#cachedSecretValues.has(identifierKey)) {
61+
return this.#cachedSecretValues.get(identifierKey); // serve from cache
62+
}
63+
return this.loadSecretValue(secretIdentifier); // fallback: load on demand
64+
}
65+
5166
clearCache(): void {
5267
if (this.#minSecretRefreshTimer.canRefresh()) {
5368
this.#cachedSecretValues.clear();

0 commit comments

Comments
 (0)