load balance support

linglingye001 · linglingye001 · commit e4ca372253a0 · 2024-11-19T15:57:00.000+08:00
diff --git 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,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;
diff --git 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/ConfigurationClientWrapper.ts b/src/ConfigurationClientWrapper.ts
@@ -12,7 +12,7 @@ export class ConfigurationClientWrapper {
     endpoint: string;
     client: AppConfigurationClient;
     backoffEndTime: number = 0; // Timestamp
-    #failedAttempts: number = 0;
+    failedAttempts: number = 0;
 
     constructor(endpoint: string, client: AppConfigurationClient) {
         this.endpoint = endpoint;
@@ -21,11 +21,17 @@ export class ConfigurationClientWrapper {
 
     updateBackoffStatus(successfull: boolean) {
         if (successfull) {
-            this.#failedAttempts = 0;
+            if (this.failedAttempts > 0) {
+                this.failedAttempts = 0;
+            }
+            this.failedAttempts -= 1;
             this.backoffEndTime = Date.now();
         } else {
-            this.#failedAttempts += 1;
-            this.backoffEndTime = Date.now() + calculateBackoffDuration(this.#failedAttempts);
+            if (this.failedAttempts < 0) {
+                this.failedAttempts = 0;
+            }
+            this.failedAttempts += 1;
+            this.backoffEndTime = Date.now() + calculateBackoffDuration(this.failedAttempts);
         }
     }
 }
diff --git 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
@@ -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
@@ -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
@@ -0,0 +1,113 @@
+// 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 { mockAppConfigurationClientListConfigurationSettings, restoreMocks, createMockedConnectionString, sleepInMs, createMockedFeatureFlag, createMockedEndpoint, mockConfigurationManagerGetClients } from "./utils/testHelper.js";
+import { AppConfigurationClient } from "@azure/app-configuration";
+import { ConfigurationClientWrapper } from "../src/ConfigurationClientWrapper.js";
+
+describe("load balance", function () {
+    this.timeout(10000);
+
+    beforeEach(() => {
+    });
+
+    afterEach(() => {
+        restoreMocks();
+    });
+
+    it("should load balance the request when loadBalancingEnabled", async () => {
+        // mock multiple pages of feature flags
+        const page1 = [
+            createMockedFeatureFlag("Alpha_1", { enabled: true }),
+            createMockedFeatureFlag("Alpha_2", { enabled: true }),
+        ];
+        const page2 = [
+            createMockedFeatureFlag("Beta_1", { enabled: true }),
+            createMockedFeatureFlag("Beta_2", { enabled: true }),
+        ];
+        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 mockedClientWrappers = [fakeClientWrapper_1, fakeClientWrapper_2];
+        mockConfigurationManagerGetClients(mockedClientWrappers, false);
+        mockAppConfigurationClientListConfigurationSettings(page1, page2);
+
+        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(fakeClientWrapper_1.failedAttempts).eq(-1);
+        expect(fakeClientWrapper_2.failedAttempts).eq(-1);
+
+        await sleepInMs(2 * 1000 + 1);
+        await settings.refresh();
+        // refresh request for feature flags
+        expect(fakeClientWrapper_1.failedAttempts).eq(-2);
+        expect(fakeClientWrapper_2.failedAttempts).eq(-1);
+
+        await sleepInMs(2 * 1000 + 1);
+        await settings.refresh();
+        expect(fakeClientWrapper_1.failedAttempts).eq(-2);
+        expect(fakeClientWrapper_2.failedAttempts).eq(-2);
+    });
+
+    it("should load balance the request when loadBalancingEnabled", async () => {
+        // mock multiple pages of feature flags
+        const page1 = [
+            createMockedFeatureFlag("Alpha_1", { enabled: true }),
+            createMockedFeatureFlag("Alpha_2", { enabled: true }),
+        ];
+        const page2 = [
+            createMockedFeatureFlag("Beta_1", { enabled: true }),
+            createMockedFeatureFlag("Beta_2", { enabled: true }),
+        ];
+        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 mockedClientWrappers = [fakeClientWrapper_1, fakeClientWrapper_2];
+        mockConfigurationManagerGetClients(mockedClientWrappers, false);
+        mockAppConfigurationClientListConfigurationSettings(page1, page2);
+
+        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(fakeClientWrapper_1.failedAttempts).eq(-2);
+        expect(fakeClientWrapper_2.failedAttempts).eq(0);
+
+        await sleepInMs(2 * 1000 + 1);
+        await settings.refresh();
+        // refresh request for feature flags
+        expect(fakeClientWrapper_1.failedAttempts).eq(-3);
+        expect(fakeClientWrapper_2.failedAttempts).eq(0);
+    });
+});
diff --git a/test/utils/testHelper.ts b/test/utils/testHelper.ts
@@ -100,9 +100,12 @@ function mockAppConfigurationClientListConfigurationSettings(...pages: Configura
     });
 }
 
-function mockConfigurationManagerGetClients(isFailoverable: boolean, ...pages: ConfigurationSetting[][]) {
+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)));

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,9 @@ export enum RequestType {`
`44`	`44`	`WATCH = "Watch"`
`45`	`45`	`}`
`46`	`46`
	`47`	`+export const FEATURES_KEY = "Features";`
	`48`	`+`
`47`	`49`	`// Tag names`
`48`	`50`	`export const FAILOVER_REQUEST_TAG = "Failover";`
`49`	`51`	`export const KEY_VAULT_CONFIGURED_TAG = "UsesKeyVault";`
	`52`	`+export const LOAD_BALANCE_CONFIGURED_TAG = "LB";`