diff --git a/src/main.ts b/src/main.ts index b2eda00e8..d280257a9 100644 --- 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 new file mode 100644 index 000000000..60074f7a5 --- /dev/null +++ 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 index 10e09a582..1aed99eb5 100644 --- 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 index 384a7abd3..9ccea373c 100644 --- 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 new file mode 100644 index 000000000..317d51d30 --- /dev/null +++ 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 index 84042bb93..63e27037d 100644 --- a/test/unitTests/shim/telemetryClient.tests.ts +++ b/test/unitTests/shim/telemetryClient.tests.ts @@ -403,4 +403,136 @@ describe("shim/TelemetryClient", () => { assert.ok(stub.calledOnce); }); }); + + describe("Instance count tracking and MULTI_IKEY statsbeat feature", () => { + 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"]; + // Reset the static instance count for testing + (TelemetryClient as any)._instanceCount = 0; + }); + + afterEach(() => { + // Restore original environment + process.env = originalEnv; + }); + + it("should not enable MULTI_IKEY feature when creating first TelemetryClient instance", () => { + const firstClient = new TelemetryClient("InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333"); + + // Check statsbeat features environment variable + const statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + if (statsbeatFeatures) { + const config = JSON.parse(statsbeatFeatures); + // MULTI_IKEY bit should not be set (128) + assert.strictEqual((config.feature & 128), 0, "MULTI_IKEY feature should not be enabled for first instance"); + } + + firstClient.shutdown(); + }); + + it("should enable MULTI_IKEY feature when creating second TelemetryClient instance", () => { + const firstClient = new TelemetryClient("InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333"); + + // First instance should not have MULTI_IKEY feature enabled + let statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + if (statsbeatFeatures) { + const config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 0, "MULTI_IKEY feature should not be enabled for first instance"); + } + + const secondClient = new TelemetryClient("InstrumentationKey=2bb22222-cccc-2ddd-9eee-fffff4444444"); + + // Second instance should have MULTI_IKEY feature enabled + statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + assert.ok(statsbeatFeatures, "AZURE_MONITOR_STATSBEAT_FEATURES should be set"); + const config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 128, "MULTI_IKEY feature should be enabled for second instance"); + + firstClient.shutdown(); + secondClient.shutdown(); + }); + + it("should keep MULTI_IKEY feature enabled when creating additional TelemetryClient instances", () => { + const firstClient = new TelemetryClient("InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333"); + const secondClient = new TelemetryClient("InstrumentationKey=2bb22222-cccc-2ddd-9eee-fffff4444444"); + + let statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + assert.ok(statsbeatFeatures, "AZURE_MONITOR_STATSBEAT_FEATURES should be set after second instance"); + let config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 128, "MULTI_IKEY feature should be enabled after second instance"); + + const thirdClient = new TelemetryClient("InstrumentationKey=3cc33333-dddd-3eee-afff-ggggg5555555"); + + statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + assert.ok(statsbeatFeatures, "AZURE_MONITOR_STATSBEAT_FEATURES should remain set for third instance"); + config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 128, "MULTI_IKEY feature should remain enabled for third instance"); + + firstClient.shutdown(); + secondClient.shutdown(); + thirdClient.shutdown(); + }); + + it("should increment instance count correctly for multiple TelemetryClient instances", () => { + const firstClient = new TelemetryClient("InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333"); + assert.strictEqual((TelemetryClient as any)._instanceCount, 1, "Instance count should be 1 after first client"); + + const secondClient = new TelemetryClient("InstrumentationKey=2bb22222-cccc-2ddd-9eee-fffff4444444"); + assert.strictEqual((TelemetryClient as any)._instanceCount, 2, "Instance count should be 2 after second client"); + + const thirdClient = new TelemetryClient("InstrumentationKey=3cc33333-dddd-3eee-afff-ggggg5555555"); + assert.strictEqual((TelemetryClient as any)._instanceCount, 3, "Instance count should be 3 after third client"); + + firstClient.shutdown(); + secondClient.shutdown(); + thirdClient.shutdown(); + }); + + it("should work with different connection strings", () => { + const firstClient = new TelemetryClient("InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333;IngestionEndpoint=https://eastus-8.in.applicationinsights.azure.com/"); + + let statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + if (statsbeatFeatures) { + const config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 0, "MULTI_IKEY feature should not be enabled for first instance with connection string"); + } + + const secondClient = new TelemetryClient("InstrumentationKey=2bb22222-cccc-2ddd-9eee-fffff4444444;IngestionEndpoint=https://westus-2.in.applicationinsights.azure.com/"); + + statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + assert.ok(statsbeatFeatures, "AZURE_MONITOR_STATSBEAT_FEATURES should be set"); + const config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 128, "MULTI_IKEY feature should be enabled for second instance with different connection string"); + + firstClient.shutdown(); + secondClient.shutdown(); + }); + + it("should work when no connection string is provided", () => { + const firstClient = new TelemetryClient(); + assert.strictEqual((TelemetryClient as any)._instanceCount, 1, "Instance count should be 1 after first client with no connection string"); + + let statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + if (statsbeatFeatures) { + const config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 0, "MULTI_IKEY feature should not be enabled for first instance with no connection string"); + } + + const secondClient = new TelemetryClient(); + assert.strictEqual((TelemetryClient as any)._instanceCount, 2, "Instance count should be 2 after second client with no connection string"); + + statsbeatFeatures = process.env["AZURE_MONITOR_STATSBEAT_FEATURES"]; + assert.ok(statsbeatFeatures, "AZURE_MONITOR_STATSBEAT_FEATURES should be set"); + const config = JSON.parse(statsbeatFeatures); + assert.strictEqual((config.feature & 128), 128, "MULTI_IKEY feature should be enabled for second instance with no connection string"); + + firstClient.shutdown(); + secondClient.shutdown(); + }); + }); }); \ No newline at end of file