diff --git a/src/AzureAppConfigurationImpl.ts b/src/AzureAppConfigurationImpl.ts index 08f3bb3e..bde471e6 100644 --- a/src/AzureAppConfigurationImpl.ts +++ b/src/AzureAppConfigurationImpl.ts @@ -75,9 +75,12 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration { #featureFlagRefreshInterval: number = DEFAULT_REFRESH_INTERVAL_IN_MS; #featureFlagRefreshTimer: RefreshTimer; - // selectors + // Selectors #featureFlagSelectors: PagedSettingSelector[] = []; + // Load balancing + #lastSuccessfulEndpoint: string = ""; + constructor( clientManager: ConfigurationClientManager, options: AzureAppConfigurationOptions | undefined, @@ -202,7 +205,21 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration { } async #executeWithFailoverPolicy(funcToExecute: (client: AppConfigurationClient) => Promise): Promise { - const clientWrappers = await this.#clientManager.getClients(); + let clientWrappers = await this.#clientManager.getClients(); + if (this.#options?.loadBalancingEnabled && this.#lastSuccessfulEndpoint !== "" && clientWrappers.length > 1) { + let nextClientIndex = 0; + // Iterate through clients to find the index of the client with the last successful endpoint + for (const clientWrapper of clientWrappers) { + nextClientIndex++; + if (clientWrapper.endpoint === this.#lastSuccessfulEndpoint) { + break; + } + } + // If we found the last successful client, rotate the list so that the next client is at the beginning + if (nextClientIndex < clientWrappers.length) { + clientWrappers = [...clientWrappers.slice(nextClientIndex), ...clientWrappers.slice(0, nextClientIndex)]; + } + } let successful: boolean; for (const clientWrapper of clientWrappers) { @@ -210,6 +227,7 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration { try { const result = await funcToExecute(clientWrapper.client); this.#isFailoverRequest = false; + this.#lastSuccessfulEndpoint = clientWrapper.endpoint; successful = true; clientWrapper.updateBackoffStatus(successful); return result; diff --git a/src/AzureAppConfigurationOptions.ts b/src/AzureAppConfigurationOptions.ts index 4aa3f99d..56b47b50 100644 --- a/src/AzureAppConfigurationOptions.ts +++ b/src/AzureAppConfigurationOptions.ts @@ -55,4 +55,12 @@ export interface AzureAppConfigurationOptions { * If not specified, the default value is true. */ replicaDiscoveryEnabled?: boolean; + + /** + * Specifies whether to enable load balance or not. + * + * @remarks + * If not specified, the default value is false. + */ + loadBalancingEnabled?: boolean; } diff --git a/src/requestTracing/constants.ts b/src/requestTracing/constants.ts index 60dbb81a..8b39b636 100644 --- a/src/requestTracing/constants.ts +++ b/src/requestTracing/constants.ts @@ -44,6 +44,9 @@ export enum RequestType { WATCH = "Watch" } +export const FEATURES_KEY = "Features"; + // Tag names export const FAILOVER_REQUEST_TAG = "Failover"; export const KEY_VAULT_CONFIGURED_TAG = "UsesKeyVault"; +export const LOAD_BALANCE_CONFIGURED_TAG = "LB"; diff --git a/src/requestTracing/utils.ts b/src/requestTracing/utils.ts index 8a2fdbc4..d48bd5ec 100644 --- a/src/requestTracing/utils.ts +++ b/src/requestTracing/utils.ts @@ -20,7 +20,9 @@ import { RequestType, SERVICE_FABRIC_ENV_VAR, CORRELATION_CONTEXT_HEADER_NAME, - FAILOVER_REQUEST_TAG + FAILOVER_REQUEST_TAG, + FEATURES_KEY, + LOAD_BALANCE_CONFIGURED_TAG } from "./constants"; // Utils @@ -84,6 +86,9 @@ export function createCorrelationContextHeader(options: AzureAppConfigurationOpt keyValues.set(REQUEST_TYPE_KEY, isInitialLoadCompleted ? RequestType.WATCH : RequestType.STARTUP); keyValues.set(HOST_TYPE_KEY, getHostType()); keyValues.set(ENV_KEY, isDevEnvironment() ? DEV_ENV_VAL : undefined); + if (options?.loadBalancingEnabled) { + keyValues.set(FEATURES_KEY, LOAD_BALANCE_CONFIGURED_TAG); + } const tags: string[] = []; if (options?.keyVaultOptions) { diff --git a/test/failover.test.ts b/test/failover.test.ts index c97a127a..2671a1f2 100644 --- a/test/failover.test.ts +++ b/test/failover.test.ts @@ -35,7 +35,7 @@ describe("failover", function () { it("should failover to replica and load key values from config store", async () => { const isFailoverable = true; - mockConfigurationManagerGetClients(isFailoverable, mockedKVs); + mockConfigurationManagerGetClients([], isFailoverable, mockedKVs); const connectionString = createMockedConnectionString(); // replicaDiscoveryEnabled is default to true @@ -47,7 +47,7 @@ describe("failover", function () { it("should failover to replica and load feature flags from config store", async () => { const isFailoverable = true; - mockConfigurationManagerGetClients(isFailoverable, mockedFeatureFlags); + mockConfigurationManagerGetClients([], isFailoverable, mockedFeatureFlags); const connectionString = createMockedConnectionString(); // replicaDiscoveryEnabled is default to true @@ -66,7 +66,7 @@ describe("failover", function () { it("should throw error when all clients failed", async () => { const isFailoverable = false; - mockConfigurationManagerGetClients(isFailoverable); + mockConfigurationManagerGetClients([], isFailoverable); const connectionString = createMockedConnectionString(); return expect(load(connectionString)).eventually.rejectedWith("All clients failed to get configuration settings."); diff --git a/test/loadBalance.test.ts b/test/loadBalance.test.ts new file mode 100644 index 00000000..f0e04b01 --- /dev/null +++ b/test/loadBalance.test.ts @@ -0,0 +1,96 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +import * as chai from "chai"; +import * as chaiAsPromised from "chai-as-promised"; +chai.use(chaiAsPromised); +const expect = chai.expect; +import { load } from "./exportedApi.js"; +import { restoreMocks, createMockedConnectionString, sleepInMs, createMockedEndpoint, mockConfigurationManagerGetClients, mockAppConfigurationClientLoadBalanceMode } from "./utils/testHelper.js"; +import { AppConfigurationClient } from "@azure/app-configuration"; +import { ConfigurationClientWrapper } from "../src/ConfigurationClientWrapper.js"; + +const fakeEndpoint_1 = createMockedEndpoint("fake_1"); +const fakeEndpoint_2 = createMockedEndpoint("fake_2"); +const fakeClientWrapper_1 = new ConfigurationClientWrapper(fakeEndpoint_1, new AppConfigurationClient(createMockedConnectionString(fakeEndpoint_1))); +const fakeClientWrapper_2 = new ConfigurationClientWrapper(fakeEndpoint_2, new AppConfigurationClient(createMockedConnectionString(fakeEndpoint_2))); +const clientRequestCounter_1 = {count: 0}; +const clientRequestCounter_2 = {count: 0}; + +describe("load balance", function () { + this.timeout(10000); + + beforeEach(() => { + }); + + afterEach(() => { + restoreMocks(); + }); + + it("should load balance the request when loadBalancingEnabled", async () => { + mockConfigurationManagerGetClients([fakeClientWrapper_1, fakeClientWrapper_2], false); + mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_1, clientRequestCounter_1); + mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_2, clientRequestCounter_2); + + const connectionString = createMockedConnectionString(); + const settings = await load(connectionString, { + loadBalancingEnabled: true, + featureFlagOptions: { + enabled: true, + selectors: [{ + keyFilter: "*" + }], + refresh: { + enabled: true, + refreshIntervalInMs: 2000 // 2 seconds for quick test. + } + } + }); + // one request for key values, one request for feature flags + expect(clientRequestCounter_1.count).eq(1); + expect(clientRequestCounter_2.count).eq(1); + + await sleepInMs(2 * 1000 + 1); + await settings.refresh(); + // refresh request for feature flags + expect(clientRequestCounter_1.count).eq(2); + expect(clientRequestCounter_2.count).eq(1); + + await sleepInMs(2 * 1000 + 1); + await settings.refresh(); + expect(clientRequestCounter_1.count).eq(2); + expect(clientRequestCounter_2.count).eq(2); + }); + + it("should not load balance the request when loadBalance disabled", async () => { + clientRequestCounter_1.count = 0; + clientRequestCounter_2.count = 0; + mockConfigurationManagerGetClients([fakeClientWrapper_1, fakeClientWrapper_2], false); + mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_1, clientRequestCounter_1); + mockAppConfigurationClientLoadBalanceMode(fakeClientWrapper_2, clientRequestCounter_2); + + const connectionString = createMockedConnectionString(); + // loadBalancingEnabled is default to false + const settings = await load(connectionString, { + featureFlagOptions: { + enabled: true, + selectors: [{ + keyFilter: "*" + }], + refresh: { + enabled: true, + refreshIntervalInMs: 2000 // 2 seconds for quick test. + } + } + }); + // one request for key values, one request for feature flags + expect(clientRequestCounter_1.count).eq(2); + expect(clientRequestCounter_2.count).eq(0); + + await sleepInMs(2 * 1000 + 1); + await settings.refresh(); + // refresh request for feature flags + expect(clientRequestCounter_1.count).eq(3); + expect(clientRequestCounter_2.count).eq(0); + }); +}); diff --git a/test/utils/testHelper.ts b/test/utils/testHelper.ts index 261b9b57..9284f871 100644 --- a/test/utils/testHelper.ts +++ b/test/utils/testHelper.ts @@ -100,9 +100,21 @@ function mockAppConfigurationClientListConfigurationSettings(...pages: Configura }); } -function mockConfigurationManagerGetClients(isFailoverable: boolean, ...pages: ConfigurationSetting[][]) { +function mockAppConfigurationClientLoadBalanceMode(clientWrapper: ConfigurationClientWrapper, countObject: { count: number }) { + const emptyPages: ConfigurationSetting[][] = []; + sinon.stub(clientWrapper.client, "listConfigurationSettings").callsFake((listOptions) => { + countObject.count += 1; + const kvs = _filterKVs(emptyPages.flat(), listOptions); + return getMockedIterator(emptyPages, kvs, listOptions); + }); +} + +function mockConfigurationManagerGetClients(fakeClientWrappers: ConfigurationClientWrapper[], isFailoverable: boolean, ...pages: ConfigurationSetting[][]) { // Stub the getClients method on the class prototype sinon.stub(ConfigurationClientManager.prototype, "getClients").callsFake(async () => { + if (fakeClientWrappers?.length > 0) { + return fakeClientWrappers; + } const clients: ConfigurationClientWrapper[] = []; const fakeEndpoint = createMockedEndpoint("fake"); const fakeStaticClientWrapper = new ConfigurationClientWrapper(fakeEndpoint, new AppConfigurationClient(createMockedConnectionString(fakeEndpoint))); @@ -230,6 +242,7 @@ export { sinon, mockAppConfigurationClientListConfigurationSettings, mockAppConfigurationClientGetConfigurationSetting, + mockAppConfigurationClientLoadBalanceMode, mockConfigurationManagerGetClients, mockSecretClientGetSecret, restoreMocks,