azure: Add in-memory cache for Azure location/provider API calls

azure/index.d.ts

@@ -201,6 +201,41 @@ export declare class LocationListStep<T extends ILocationWizardContext> extends
     public static getQuickPickDescription?: (location: AzExtLocation) => string | undefined;
 }
 
+/**
+ * A simple cache for location data that deduplicates in-flight requests.
+ *
+ * Currently backed by an in-memory Map. Designed so the backing store can be
+ * swapped to persistent storage (e.g. `vscode.Memento` / globalState) in the
+ * future to survive across VS Code restarts with a longer TTL (e.g. 7 days).
+ */
+export declare class LocationCache {
+    /**
+     * @param ttlMs Optional time-to-live in milliseconds. When omitted, entries
+     * never expire (suitable for in-memory caches that reset on extension
+     * deactivation). Set this when switching to persistent storage.
+     */
+    constructor(ttlMs?: number);
+    /**
+     * Get a value from the cache, or fetch it if missing/expired.
+     * Concurrent calls with the same key share a single in-flight request.
+     */
+    getOrLoad<T>(key: string, loader: () => Promise<T>): Promise<T>;
+    /** Remove all entries (e.g. on sign-out or subscription change) */
+    clear(): void;
+}
+
+/**
+ * Module-level cache for subscription locations, shared across all wizard instances
+ * within the same extension activation. Keyed by `subscriptionId|includeExtended`.
+ */
+export declare const subscriptionLocationsCache: LocationCache;
+
+/**
+ * Module-level cache for provider resource-type locations, shared across all wizard instances
+ * within the same extension activation. Keyed by `subscriptionId|provider|resourceType`.
+ */
+export declare const providerLocationsCache: LocationCache;
+
 /**
  * Checks to see if providers (i.e. 'Microsoft.Web') are registered and registers them if they're not
  */

azure/src/index.ts

@@ -16,6 +16,7 @@ export * from './utils/parseAzureResourceId';
 export * from './utils/setupAzureLogger';
 export * from './utils/uiUtils';
 export * from './wizard/LocationListStep';
+export { LocationCache, subscriptionLocationsCache, providerLocationsCache } from './wizard/LocationCache';
 export * from './wizard/ResourceGroupCreateStep';
 export * from './wizard/ResourceGroupListStep';
 export * from './wizard/ResourceGroupNameStep';

azure/src/wizard/LocationCache.ts

@@ -0,0 +1,82 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+/**
+ * A simple cache for location data that deduplicates in-flight requests.
+ *
+ * Currently backed by an in-memory Map. Designed so the backing store can be
+ * swapped to persistent storage (e.g. `vscode.Memento` / globalState) in the
+ * future to survive across VS Code restarts with a longer TTL (e.g. 7 days).
+ */
+
+interface CacheEntry<T> {
+    data: T;
+    /** Unix timestamp (ms) when this entry was stored */
+    storedAt: number;
+}
+
+export class LocationCache {
+    private readonly cache = new Map<string, CacheEntry<unknown>>();
+
+    /**
+     * In-flight promises keyed the same way as the cache.
+     * Ensures concurrent callers share the same request instead of firing duplicates.
+     */
+    private readonly inflight = new Map<string, Promise<unknown>>();
+
+    /**
+     * @param ttlMs Optional time-to-live in milliseconds. When omitted, entries
+     * never expire (suitable for in-memory caches that reset on extension
+     * deactivation). Set this when switching to persistent storage.
+     */
+    constructor(private readonly ttlMs?: number) { }
+
+    /**
+     * Get a value from the cache, or fetch it if missing/expired.
+     * Concurrent calls with the same key share a single in-flight request.
+     */
+    async getOrLoad<T>(key: string, loader: () => Promise<T>): Promise<T> {
+        const cached = this.cache.get(key) as CacheEntry<T> | undefined;
+        if (cached && !this.isExpired(cached)) {
+            return cached.data;
+        }
+
+        // Check for an in-flight request we can piggy-back on
+        const existing = this.inflight.get(key) as Promise<T> | undefined;
+        if (existing) {
+            return existing;
+        }
+
+        const promise = loader().then(data => {
+            this.cache.set(key, { data, storedAt: Date.now() });
+            this.inflight.delete(key);
+            return data;
+        }).catch(err => {
+            this.inflight.delete(key);
+            throw err;
+        });
+
+        this.inflight.set(key, promise);
+        return promise;
+    }
+
+    /** Remove all entries (e.g. on sign-out or subscription change) */
+    clear(): void {
+        this.cache.clear();
+        // Don't clear inflight — let pending requests finish naturally
+    }
+
+    private isExpired(entry: CacheEntry<unknown>): boolean {
+        return this.ttlMs !== undefined && (Date.now() - entry.storedAt) > this.ttlMs;
+    }
+}
+
+/**
+ * Module-level caches shared across all wizard instances within the same
+ * extension activation. Keyed by subscription ID (+ provider info for
+ * provider locations).
+ */
+export const subscriptionLocationsCache = new LocationCache();
+export const providerLocationsCache = new LocationCache();

azure/src/wizard/LocationListStep.ts

@@ -12,6 +12,7 @@ import { createResourcesClient, createSubscriptionsClient } from '../clients';
 import { resourcesProvider } from '../constants';
 import { ext } from '../extensionVariables';
 import { uiUtils } from '../utils/uiUtils';
+import { providerLocationsCache, subscriptionLocationsCache } from './LocationCache';
 
 /* eslint-disable @typescript-eslint/naming-convention */
 interface ILocationWizardContextInternal extends types.ILocationWizardContext {
@@ -254,19 +255,28 @@ export class LocationListStep<T extends ILocationWizardContextInternal> extends
 }
 
 async function getAllLocations(wizardContext: types.ILocationWizardContext): Promise<types.AzExtLocation[]> {
-    const client = await createSubscriptionsClient(wizardContext);
-    const locations = await uiUtils.listAllIterator<Location>(client.subscriptions.listLocations(wizardContext.subscriptionId, { includeExtendedLocations: wizardContext.includeExtendedLocations }));
-    return locations.filter((l): l is types.AzExtLocation => !!(l.id && l.name && l.displayName));
+    const includeExtended = !!wizardContext.includeExtendedLocations;
+    const cacheKey = `${wizardContext.subscriptionId}|${includeExtended}`;
+
+    return subscriptionLocationsCache.getOrLoad(cacheKey, async () => {
+        const client = await createSubscriptionsClient(wizardContext);
+        const locations = await uiUtils.listAllIterator<Location>(client.subscriptions.listLocations(wizardContext.subscriptionId, { includeExtendedLocations: includeExtended }));
+        return locations.filter((l): l is types.AzExtLocation => !!(l.id && l.name && l.displayName));
+    });
 }
 
 async function getProviderLocations(wizardContext: types.ILocationWizardContext, provider: string, resourceType: string): Promise<string[]> {
-    const rgClient = await createResourcesClient(wizardContext);
-    const providerData = await rgClient.providers.get(provider);
-    const resourceTypeData = providerData.resourceTypes?.find(rt => rt.resourceType?.toLowerCase() === resourceType.toLowerCase());
-    if (!resourceTypeData) {
-        throw new ProviderResourceTypeNotFoundError(providerData, resourceType);
-    }
-    return nonNullProp(resourceTypeData, 'locations');
+    const cacheKey = `${wizardContext.subscriptionId}|${provider.toLowerCase()}|${resourceType.toLowerCase()}`;
+
+    return providerLocationsCache.getOrLoad(cacheKey, async () => {
+        const rgClient = await createResourcesClient(wizardContext);
+        const providerData = await rgClient.providers.get(provider);
+        const resourceTypeData = providerData.resourceTypes?.find(rt => rt.resourceType?.toLowerCase() === resourceType.toLowerCase());
+        if (!resourceTypeData) {
+            throw new ProviderResourceTypeNotFoundError(providerData, resourceType);
+        }
+        return nonNullProp(resourceTypeData, 'locations');
+    });
 }
 
 function compareLocation(l1: types.AzExtLocation, l2: types.AzExtLocation): number {

azure/test/LocationCache.test.ts

@@ -0,0 +1,154 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as assert from 'assert';
+import { LocationCache } from '../src/wizard/LocationCache';
+
+suite('LocationCache', () => {
+    let cache: LocationCache;
+
+    setup(() => {
+        cache = new LocationCache();
+    });
+
+    test('returns data from loader on cache miss', async () => {
+        const result = await cache.getOrLoad('key1', async () => ['eastus', 'westus']);
+        assert.deepStrictEqual(result, ['eastus', 'westus']);
+    });
+
+    test('returns cached data on subsequent calls without calling loader again', async () => {
+        let callCount = 0;
+        const loader = async () => {
+            callCount++;
+            return ['eastus'];
+        };
+
+        const result1 = await cache.getOrLoad('key1', loader);
+        const result2 = await cache.getOrLoad('key1', loader);
+
+        assert.deepStrictEqual(result1, ['eastus']);
+        assert.deepStrictEqual(result2, ['eastus']);
+        assert.strictEqual(callCount, 1, 'loader should only be called once');
+    });
+
+    test('uses separate entries for different keys', async () => {
+        await cache.getOrLoad('sub1|false', async () => ['eastus']);
+        await cache.getOrLoad('sub2|false', async () => ['westus']);
+
+        const result1 = await cache.getOrLoad('sub1|false', async () => { throw new Error('should not call'); });
+        const result2 = await cache.getOrLoad('sub2|false', async () => { throw new Error('should not call'); });
+
+        assert.deepStrictEqual(result1, ['eastus']);
+        assert.deepStrictEqual(result2, ['westus']);
+    });
+
+    test('deduplicates concurrent in-flight requests for the same key', async () => {
+        let callCount = 0;
+        let resolve: (value: string[]) => void;
+        const loader = () => {
+            callCount++;
+            return new Promise<string[]>(r => { resolve = r; });
+        };
+
+        const p1 = cache.getOrLoad('key1', loader);
+        const p2 = cache.getOrLoad('key1', loader);
+
+        // Both should be waiting on the same promise
+        assert.strictEqual(callCount, 1, 'loader should only be called once for concurrent requests');
+
+        resolve!(['eastus']);
+        const [result1, result2] = await Promise.all([p1, p2]);
+
+        assert.deepStrictEqual(result1, ['eastus']);
+        assert.deepStrictEqual(result2, ['eastus']);
+    });
+
+    test('clear removes all cached entries', async () => {
+        let callCount = 0;
+        const loader = async () => {
+            callCount++;
+            return ['eastus'];
+        };
+
+        await cache.getOrLoad('key1', loader);
+        assert.strictEqual(callCount, 1);
+
+        cache.clear();
+
+        await cache.getOrLoad('key1', loader);
+        assert.strictEqual(callCount, 2, 'loader should be called again after clear');
+    });
+
+    test('expired entries are refreshed', async () => {
+        // Use a very short TTL
+        const shortCache = new LocationCache(1);
+        let callCount = 0;
+        const loader = async () => {
+            callCount++;
+            return [`result-${callCount}`];
+        };
+
+        const result1 = await shortCache.getOrLoad('key1', loader);
+        assert.deepStrictEqual(result1, ['result-1']);
+
+        // Wait for expiry
+        await new Promise(r => setTimeout(r, 10));
+
+        const result2 = await shortCache.getOrLoad('key1', loader);
+        assert.deepStrictEqual(result2, ['result-2']);
+        assert.strictEqual(callCount, 2, 'loader should be called again after expiry');
+    });
+
+    test('entries without TTL never expire', async () => {
+        // Default constructor has no TTL
+        let callCount = 0;
+
+        await cache.getOrLoad('key1', async () => { callCount++; return ['eastus']; });
+        await cache.getOrLoad('key1', async () => { callCount++; return ['westus']; });
+
+        assert.strictEqual(callCount, 1);
+    });
+
+    test('loader error does not poison the cache', async () => {
+        let shouldFail = true;
+        const loader = async () => {
+            if (shouldFail) {
+                throw new Error('network error');
+            }
+            return ['eastus'];
+        };
+
+        await assert.rejects(() => cache.getOrLoad('key1', loader), /network error/);
+
+        shouldFail = false;
+        const result = await cache.getOrLoad('key1', loader);
+        assert.deepStrictEqual(result, ['eastus']);
+    });
+
+    test('loader error is propagated to all concurrent waiters', async () => {
+        let reject: (err: Error) => void;
+        const loader = () => new Promise<string[]>((_, r) => { reject = r; });
+
+        const p1 = cache.getOrLoad('key1', loader);
+        const p2 = cache.getOrLoad('key1', loader);
+
+        reject!(new Error('boom'));
+
+        await assert.rejects(() => p1, /boom/);
+        await assert.rejects(() => p2, /boom/);
+    });
+
+    test('after error, a new loader call succeeds', async () => {
+        let callCount = 0;
+        const failLoader = async () => { callCount++; throw new Error('fail'); };
+        const okLoader = async () => { callCount++; return ['eastus']; };
+
+        await assert.rejects(() => cache.getOrLoad('key1', failLoader), /fail/);
+        const result = await cache.getOrLoad('key1', okLoader);
+
+        assert.deepStrictEqual(result, ['eastus']);
+        assert.strictEqual(callCount, 2);
+    });
+});