Load balance support (#135)

* load balance support

* improve test

Author: linglingye001

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,14 +205,29 @@ export class AzureAppConfigurationImpl implements AzureAppConfiguration {
     }
 
     async #executeWithFailoverPolicy(funcToExecute: (client: AppConfigurationClient) => Promise<any>): Promise<any> {
-        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) {
             successful = false;
             try {
                 const result = await funcToExecute(clientWrapper.client);
                 this.#isFailoverRequest = false;
+                this.#lastSuccessfulEndpoint = clientWrapper.endpoint;
                 successful = true;
                 clientWrapper.updateBackoffStatus(successful);
                 return result;

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;
 }

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";

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) {

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.");

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);
+    });
+});

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,