Add Support for Tracking Multi iKey Usage & Feature Statsbeat Handler (#1438)

JacksonWeber · web-flow · commit 8033d83d42a7 · 2025-06-27T14:37:04.000-07:00
* Add support for tracking multi ikey usage.

* Add feature statsbeat handler class and support for multi_ikey feature tracking.
diff --git a/src/main.ts b/src/main.ts
@@ -13,11 +13,12 @@ import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
 
 import { AutoCollectLogs } from "./logs/autoCollectLogs";
 import { AutoCollectExceptions } from "./logs/exceptions";
-import { AZURE_MONITOR_STATSBEAT_FEATURES, AzureMonitorOpenTelemetryOptions } from "./types";
+import { AzureMonitorOpenTelemetryOptions } from "./types";
 import { ApplicationInsightsConfig } from "./shared/configuration/config";
 import { LogApi } from "./shim/logsApi";
-import { StatsbeatFeature, StatsbeatInstrumentation } from "./shim/types";
+import { StatsbeatFeature } from "./shim/types";
 import { RequestSpanProcessor } from "./traces/requestProcessor";
+import { StatsbeatFeaturesManager } from "./shared/util/statsbeatFeaturesManager";
 
 let autoCollectLogs: AutoCollectLogs;
 let exceptions: AutoCollectExceptions;
@@ -27,11 +28,10 @@ let exceptions: AutoCollectExceptions;
  * @param options Configuration
  */
 export function useAzureMonitor(options?: AzureMonitorOpenTelemetryOptions) {
-    // Must set statsbeat features before they are read by the distro
-    process.env[AZURE_MONITOR_STATSBEAT_FEATURES] = JSON.stringify({
-        instrumentation: StatsbeatInstrumentation.NONE,
-        feature: StatsbeatFeature.SHIM
-    });
+    // Initialize statsbeat features with default values and enable SHIM feature
+    StatsbeatFeaturesManager.getInstance().initialize();
+    StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.SHIM);
+    
     // Allows for full filtering of dependency/request spans
     options.spanProcessors = [new RequestSpanProcessor(options.enableAutoCollectDependencies, options.enableAutoCollectRequests)];
     distroUseAzureMonitor(options);
diff --git a/src/shared/util/statsbeatFeaturesManager.ts b/src/shared/util/statsbeatFeaturesManager.ts
@@ -0,0 +1,123 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { AZURE_MONITOR_STATSBEAT_FEATURES } from "../../types";
+import { StatsbeatFeature, StatsbeatInstrumentation } from "../../shim/types";
+
+/**
+ * Interface for statsbeat features configuration
+ */
+interface StatsbeatFeaturesConfig {
+    instrumentation: number;
+    feature: number;
+}
+
+/**
+ * Utility class to manage statsbeat features using bitmap flags
+ */
+export class StatsbeatFeaturesManager {
+    private static instance: StatsbeatFeaturesManager;
+
+    /**
+     * Get the singleton instance of StatsbeatFeaturesManager
+     */
+    public static getInstance(): StatsbeatFeaturesManager {
+        if (!StatsbeatFeaturesManager.instance) {
+            StatsbeatFeaturesManager.instance = new StatsbeatFeaturesManager();
+        }
+        return StatsbeatFeaturesManager.instance;
+    }
+
+    /**
+     * Get the current statsbeat features configuration from environment variable
+     */
+    private getCurrentConfig(): StatsbeatFeaturesConfig {
+        const envValue = process.env[AZURE_MONITOR_STATSBEAT_FEATURES];
+        if (envValue) {
+            try {
+                return JSON.parse(envValue);
+            } catch (error) {
+                // If parsing fails, return default values
+                return {
+                    instrumentation: StatsbeatInstrumentation.NONE,
+                    feature: StatsbeatFeature.SHIM
+                };
+            }
+        }
+        return {
+            instrumentation: StatsbeatInstrumentation.NONE,
+            feature: StatsbeatFeature.SHIM
+        };
+    }
+
+    /**
+     * Set the statsbeat features environment variable with updated configuration
+     */
+    private setConfig(config: StatsbeatFeaturesConfig): void {
+        process.env[AZURE_MONITOR_STATSBEAT_FEATURES] = JSON.stringify(config);
+    }
+
+    /**
+     * Enable a specific statsbeat feature by setting the corresponding bit
+     */
+    public enableFeature(feature: StatsbeatFeature): void {
+        const currentConfig = this.getCurrentConfig();
+        currentConfig.feature |= feature; // Use bitwise OR to set the bit
+        this.setConfig(currentConfig);
+    }
+
+    /**
+     * Disable a specific statsbeat feature by clearing the corresponding bit
+     */
+    public disableFeature(feature: StatsbeatFeature): void {
+        const currentConfig = this.getCurrentConfig();
+        currentConfig.feature &= ~feature; // Use bitwise AND with NOT to clear the bit
+        this.setConfig(currentConfig);
+    }
+
+    /**
+     * Check if a specific statsbeat feature is enabled
+     */
+    public isFeatureEnabled(feature: StatsbeatFeature): boolean {
+        const currentConfig = this.getCurrentConfig();
+        return (currentConfig.feature & feature) !== 0;
+    }
+
+    /**
+     * Enable a specific statsbeat instrumentation by setting the corresponding bit
+     */
+    public enableInstrumentation(instrumentation: StatsbeatInstrumentation): void {
+        const currentConfig = this.getCurrentConfig();
+        currentConfig.instrumentation |= instrumentation; // Use bitwise OR to set the bit
+        this.setConfig(currentConfig);
+    }
+
+    /**
+     * Disable a specific statsbeat instrumentation by clearing the corresponding bit
+     */
+    public disableInstrumentation(instrumentation: StatsbeatInstrumentation): void {
+        const currentConfig = this.getCurrentConfig();
+        currentConfig.instrumentation &= ~instrumentation; // Use bitwise AND with NOT to clear the bit
+        this.setConfig(currentConfig);
+    }
+
+    /**
+     * Check if a specific statsbeat instrumentation is enabled
+     */
+    public isInstrumentationEnabled(instrumentation: StatsbeatInstrumentation): boolean {
+        const currentConfig = this.getCurrentConfig();
+        return (currentConfig.instrumentation & instrumentation) !== 0;
+    }
+
+    /**
+     * Initialize the statsbeat features environment variable with default values if not set
+     */
+    public initialize(): void {
+        if (!process.env[AZURE_MONITOR_STATSBEAT_FEATURES]) {
+            this.setConfig({
+                instrumentation: StatsbeatInstrumentation.NONE,
+                feature: StatsbeatFeature.SHIM
+            });
+        }
+    }
+}
diff --git a/src/shim/telemetryClient.ts b/src/shim/telemetryClient.ts
@@ -26,13 +26,15 @@ import { AttributeLogProcessor } from "../shared/util/attributeLogRecordProcesso
 import { LogApi } from "./logsApi";
 import { flushAzureMonitor, shutdownAzureMonitor, useAzureMonitor } from "../main";
 import { AzureMonitorOpenTelemetryOptions } from "../types";
-import { UNSUPPORTED_MSG } from "./types";
+import { UNSUPPORTED_MSG, StatsbeatFeature } from "./types";
+import { StatsbeatFeaturesManager } from "../shared/util/statsbeatFeaturesManager";
 
 /**
  * Application Insights telemetry client provides interface to track telemetry items, register telemetry initializers and
  * and manually trigger immediate sending (flushing)
  */
 export class TelemetryClient {
+    private static _instanceCount = 0;
     public context: Context;
     public commonProperties: { [key: string]: string };
     public config: Config;
@@ -48,6 +50,13 @@ export class TelemetryClient {
      * @param setupString the Connection String or Instrumentation Key to use (read from environment variable if not specified)
      */
     constructor(input?: string) {
+        TelemetryClient._instanceCount++;
+        
+        // Set statsbeat feature if this is the second or subsequent TelemetryClient instance
+        if (TelemetryClient._instanceCount >= 2) {
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.MULTI_IKEY);
+        }
+        
         const config = new Config(input, this._configWarnings);
         this.config = config;
         this.commonProperties = {};
diff --git a/src/shim/types.ts b/src/shim/types.ts
@@ -351,6 +351,8 @@ export enum StatsbeatFeature {
     DISTRO = 8,
     LIVE_METRICS = 16,
     SHIM = 32,
+    CUSTOMER_STATSBEAT = 64,
+    MULTI_IKEY = 128,
 }
 
 /**
diff --git a/test/unitTests/shared/util/statsbeatFeaturesManager.tests.ts b/test/unitTests/shared/util/statsbeatFeaturesManager.tests.ts
@@ -0,0 +1,145 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import * as assert from "assert";
+import { StatsbeatFeaturesManager } from "../../../../src/shared/util/statsbeatFeaturesManager";
+import { StatsbeatFeature, StatsbeatInstrumentation } from "../../../../src/shim/types";
+
+describe("shared/util/StatsbeatFeaturesManager", () => {
+    let originalEnv: NodeJS.ProcessEnv;
+
+    beforeEach(() => {
+        // Save original environment
+        originalEnv = { ...process.env };
+        // Clear the AZURE_MONITOR_STATSBEAT_FEATURES environment variable before each test
+        delete process.env["AZURE_MONITOR_STATSBEAT_FEATURES"];
+    });
+
+    afterEach(() => {
+        // Restore original environment
+        process.env = originalEnv;
+    });
+
+    describe("initialize", () => {
+        it("should initialize environment variable with default values when not set", () => {
+            StatsbeatFeaturesManager.getInstance().initialize();
+            
+            const envValue = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"];
+            assert.ok(envValue, "AZURE_MONITOR_STATSBEAT_FEATURES should be set after initialization");
+            
+            const config = JSON.parse(envValue);
+            assert.strictEqual(config.instrumentation, StatsbeatInstrumentation.NONE, "instrumentation should default to NONE");
+            assert.strictEqual(config.feature, StatsbeatFeature.SHIM, "feature should default to SHIM");
+        });
+
+        it("should not overwrite existing environment variable", () => {
+            const existingValue = JSON.stringify({
+                instrumentation: StatsbeatInstrumentation.MONGODB,
+                feature: StatsbeatFeature.LIVE_METRICS
+            });
+            process.env["AZURE_MONITOR_STATSBEAT_FEATURES"] = existingValue;
+            
+            StatsbeatFeaturesManager.getInstance().initialize();
+            
+            assert.strictEqual(process.env["AZURE_MONITOR_STATSBEAT_FEATURES"], existingValue, "existing value should not be overwritten");
+        });
+    });
+
+    describe("enableFeature", () => {
+        it("should enable MULTI_IKEY feature using bitmap", () => {
+            StatsbeatFeaturesManager.getInstance().initialize();
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.MULTI_IKEY);
+            
+            const envValue = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"];
+            assert.ok(envValue, "environment variable should be set");
+            
+            const config = JSON.parse(envValue);
+            assert.ok((config.feature & StatsbeatFeature.MULTI_IKEY) !== 0, "MULTI_IKEY feature should be enabled");
+            assert.ok((config.feature & StatsbeatFeature.SHIM) !== 0, "SHIM feature should remain enabled");
+        });
+
+        it("should enable CUSTOMER_STATSBEAT feature using bitmap", () => {
+            StatsbeatFeaturesManager.getInstance().initialize();
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.CUSTOMER_STATSBEAT);
+            
+            const envValue = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"];
+            assert.ok(envValue, "environment variable should be set");
+            
+            const config = JSON.parse(envValue);
+            assert.ok((config.feature & StatsbeatFeature.CUSTOMER_STATSBEAT) !== 0, "CUSTOMER_STATSBEAT feature should be enabled");
+            assert.ok((config.feature & StatsbeatFeature.SHIM) !== 0, "SHIM feature should remain enabled");
+        });
+
+        it("should enable multiple features using bitmap", () => {
+            StatsbeatFeaturesManager.getInstance().initialize();
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.MULTI_IKEY);
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.LIVE_METRICS);
+            
+            const envValue = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"];
+            assert.ok(envValue, "environment variable should be set");
+            
+            const config = JSON.parse(envValue);
+            assert.ok((config.feature & StatsbeatFeature.MULTI_IKEY) !== 0, "MULTI_IKEY feature should be enabled");
+            assert.ok((config.feature & StatsbeatFeature.LIVE_METRICS) !== 0, "LIVE_METRICS feature should be enabled");
+            assert.ok((config.feature & StatsbeatFeature.SHIM) !== 0, "SHIM feature should remain enabled");
+        });
+    });
+
+    describe("disableFeature", () => {
+        it("should disable specific feature using bitmap", () => {
+            StatsbeatFeaturesManager.getInstance().initialize();
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.MULTI_IKEY);
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.LIVE_METRICS);
+            
+            // Disable only MULTI_IKEY
+            StatsbeatFeaturesManager.getInstance().disableFeature(StatsbeatFeature.MULTI_IKEY);
+            
+            const envValue = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"];
+            assert.ok(envValue, "environment variable should be set");
+            
+            const config = JSON.parse(envValue);
+            assert.strictEqual((config.feature & StatsbeatFeature.MULTI_IKEY), 0, "MULTI_IKEY feature should be disabled");
+            assert.ok((config.feature & StatsbeatFeature.LIVE_METRICS) !== 0, "LIVE_METRICS feature should remain enabled");
+            assert.ok((config.feature & StatsbeatFeature.SHIM) !== 0, "SHIM feature should remain enabled");
+        });
+    });
+
+    describe("isFeatureEnabled", () => {
+        it("should correctly detect enabled features", () => {
+            StatsbeatFeaturesManager.getInstance().initialize();
+            
+            assert.ok(StatsbeatFeaturesManager.getInstance().isFeatureEnabled(StatsbeatFeature.SHIM), "SHIM should be enabled by default");
+            assert.ok(!StatsbeatFeaturesManager.getInstance().isFeatureEnabled(StatsbeatFeature.MULTI_IKEY), "MULTI_IKEY should not be enabled by default");
+            
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.MULTI_IKEY);
+            assert.ok(StatsbeatFeaturesManager.getInstance().isFeatureEnabled(StatsbeatFeature.MULTI_IKEY), "MULTI_IKEY should be enabled after enableFeature");
+        });
+    });
+
+    describe("instrumentation management", () => {
+        it("should enable and disable instrumentation features", () => {
+            StatsbeatFeaturesManager.getInstance().initialize();
+            
+            assert.ok(!StatsbeatFeaturesManager.getInstance().isInstrumentationEnabled(StatsbeatInstrumentation.MONGODB), "MONGODB should not be enabled by default");
+            
+            StatsbeatFeaturesManager.getInstance().enableInstrumentation(StatsbeatInstrumentation.MONGODB);
+            assert.ok(StatsbeatFeaturesManager.getInstance().isInstrumentationEnabled(StatsbeatInstrumentation.MONGODB), "MONGODB should be enabled after enableInstrumentation");
+            
+            StatsbeatFeaturesManager.getInstance().disableInstrumentation(StatsbeatInstrumentation.MONGODB);
+            assert.ok(!StatsbeatFeaturesManager.getInstance().isInstrumentationEnabled(StatsbeatInstrumentation.MONGODB), "MONGODB should be disabled after disableInstrumentation");
+        });
+    });
+
+    describe("error handling", () => {
+        it("should handle malformed JSON in environment variable", () => {
+            process.env["AZURE_MONITOR_STATSBEAT_FEATURES"] = "invalid json";
+            
+            // Should not throw and should return default values
+            assert.ok(!StatsbeatFeaturesManager.getInstance().isFeatureEnabled(StatsbeatFeature.MULTI_IKEY), "should handle malformed JSON gracefully");
+            
+            // Should be able to enable features despite malformed initial value
+            StatsbeatFeaturesManager.getInstance().enableFeature(StatsbeatFeature.MULTI_IKEY);
+            assert.ok(StatsbeatFeaturesManager.getInstance().isFeatureEnabled(StatsbeatFeature.MULTI_IKEY), "should be able to enable features after handling malformed JSON");
+        });
+    });
+});
diff --git a/test/unitTests/shim/telemetryClient.tests.ts b/test/unitTests/shim/telemetryClient.tests.ts

Original file line number	Diff line number	Diff line change
`@@ -351,6 +351,8 @@ export enum StatsbeatFeature {`
`351`	`351`	`DISTRO = 8,`
`352`	`352`	`LIVE_METRICS = 16,`
`353`	`353`	`SHIM = 32,`
	`354`	`+ CUSTOMER_STATSBEAT = 64,`
	`355`	`+ MULTI_IKEY = 128,`
`354`	`356`	`}`
`355`	`357`
`356`	`358`	`/**`