[Azure Monitor OpenTelemetry] Add Azure Functions Correlation Hook (Azure#26313)

Functionality for @azure/monitor-opentelemetry package
Adding Azure Fn code to handle auto correlation when running in that
environment

Author: hectorhdzg

common/config/rush/pnpm-lock.yaml

@@ -86,6 +86,7 @@
     "typescript": "~5.0.0"
   },
   "dependencies": {
+    "@azure/functions": "^3.2.0",
     "@azure/monitor-opentelemetry-exporter": "1.0.0-beta.11",
     "@azure/opentelemetry-instrumentation-azure-sdk": "^1.0.0-beta.3",
     "@opentelemetry/api": "^1.4.1",

sdk/monitor/monitor-opentelemetry/package.json

@@ -0,0 +1,70 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import { Context as AzureFnContext } from "@azure/functions";
+import {
+  context,
+  propagation,
+  ROOT_CONTEXT,
+  Context as OpenTelemetryContext,
+} from "@opentelemetry/api";
+import { Logger } from "../shared/logging";
+
+export class AzureFunctionsHook {
+  private _functionsCoreModule: any;
+  private _preInvocationHook: any;
+
+  constructor() {
+    try {
+      // TODO: Add types files when publicly available
+      this._functionsCoreModule = require("@azure/functions-core");
+      // Only v3 of Azure Functions library is supported right now. See matrix of versions here:
+      // https://github.com/Azure/azure-functions-nodejs-library
+      const funcProgModel = this._functionsCoreModule.getProgrammingModel();
+      if (funcProgModel.name === "@azure/functions" && funcProgModel.version.startsWith("3.")) {
+        this._addPreInvocationHook();
+      } else {
+        Logger.getInstance().debug(
+          `AzureFunctionsHook does not support model "${funcProgModel.name}" version "${funcProgModel.version}"`
+        );
+      }
+    } catch (error) {
+      Logger.getInstance().debug(
+        "@azure/functions-core failed to load, not running in Azure Functions"
+      );
+    }
+  }
+
+  public shutdown() {
+    if (this._preInvocationHook) {
+      this._preInvocationHook.dispose();
+      this._preInvocationHook = undefined;
+    }
+    this._functionsCoreModule = undefined;
+  }
+
+  private _addPreInvocationHook() {
+    if (!this._preInvocationHook) {
+      this._preInvocationHook = this._functionsCoreModule.registerHook(
+        "preInvocation",
+        async (preInvocationContext: any) => {
+          const ctx: AzureFnContext = <AzureFnContext>preInvocationContext.invocationContext;
+          // Update context to use Azure Functions one
+          let extractedContext: OpenTelemetryContext | any = null;
+          try {
+            if (ctx.traceContext) {
+              extractedContext = propagation.extract(ROOT_CONTEXT, ctx.traceContext);
+            }
+            const currentContext = extractedContext || context.active();
+            preInvocationContext.functionCallback = context.bind(
+              currentContext,
+              preInvocationContext.functionCallback
+            );
+          } catch (err) {
+            Logger.getInstance().error("Failed to propagate context in Azure Functions", err);
+          }
+        }
+      );
+    }
+  }
+}

sdk/monitor/monitor-opentelemetry/src/traces/azureFnHook.ts

@@ -31,6 +31,7 @@ import { AzureMonitorOpenTelemetryConfig } from "../shared/config";
 import { MetricHandler } from "../metrics/handler";
 import { ignoreOutgoingRequestHook } from "../utils/common";
 import { AzureMonitorSpanProcessor } from "./spanProcessor";
+import { AzureFunctionsHook } from "./azureFnHook";
 
 /**
  * Azure Monitor OpenTelemetry Trace Handler
@@ -50,6 +51,7 @@ export class TraceHandler {
   private _redis4Instrumentation?: Instrumentation;
   private _config: AzureMonitorOpenTelemetryConfig;
   private _metricHandler?: MetricHandler;
+  private _azureFunctionsHook: AzureFunctionsHook;
 
   /**
    * Initializes a new instance of the TraceHandler class.
@@ -83,6 +85,7 @@ export class TraceHandler {
       const azureSpanProcessor = new AzureMonitorSpanProcessor(this._metricHandler);
       this._tracerProvider.addSpanProcessor(azureSpanProcessor);
     }
+    this._azureFunctionsHook = new AzureFunctionsHook();
     this._initializeInstrumentations();
   }
 
@@ -105,6 +108,7 @@ export class TraceHandler {
    */
   public async shutdown(): Promise<void> {
     await this._tracerProvider.shutdown();
+    this._azureFunctionsHook.shutdown();
   }
 
   /**

sdk/monitor/monitor-opentelemetry/src/traces/handler.ts

@@ -6,10 +6,10 @@ import sinon from "sinon";
 import { trace, context, isValidTraceId, isValidSpanId } from "@opentelemetry/api";
 import { LogRecord as APILogRecord } from "@opentelemetry/api-logs";
 import { ExportResultCode } from "@opentelemetry/core";
-import { LogHandler } from "../../../src/logs";
-import { MetricHandler } from "../../../src/metrics";
-import { TraceHandler } from "../../../src/traces";
-import { AzureMonitorOpenTelemetryConfig } from "../../../src/shared";
+import { LogHandler } from "../../../../src/logs";
+import { MetricHandler } from "../../../../src/metrics";
+import { TraceHandler } from "../../../../src/traces";
+import { AzureMonitorOpenTelemetryConfig } from "../../../../src/shared";
 
 describe("LogHandler", () => {
   let sandbox: sinon.SinonSandbox;

sdk/monitor/monitor-opentelemetry/test/internal/unit/logHandler.test.ts renamed to sdk/monitor/monitor-opentelemetry/test/internal/unit/logs/logHandler.test.ts

@@ -2,8 +2,8 @@
 // Licensed under the MIT license.
 
 import * as assert from "assert";
-import { MetricHandler } from "../../../src/metrics";
-import { AzureMonitorOpenTelemetryConfig } from "../../../src/shared";
+import { MetricHandler } from "../../../../src/metrics";
+import { AzureMonitorOpenTelemetryConfig } from "../../../../src/shared";
 
 describe("MetricHandler", () => {
   let _config: AzureMonitorOpenTelemetryConfig;

sdk/monitor/monitor-opentelemetry/test/internal/unit/metricHandler.test.ts renamed to sdk/monitor/monitor-opentelemetry/test/internal/unit/metrics/metricHandler.test.ts

@@ -12,8 +12,8 @@ import {
 import { ExportResultCode } from "@opentelemetry/core";
 import { LoggerProvider, LogRecord, Logger } from "@opentelemetry/sdk-logs";
 import { Resource } from "@opentelemetry/resources";
-import { StandardMetrics } from "../../../src/metrics/standardMetrics";
-import { AzureMonitorOpenTelemetryConfig } from "../../../src/shared";
+import { StandardMetrics } from "../../../../src/metrics/standardMetrics";
+import { AzureMonitorOpenTelemetryConfig } from "../../../../src/shared";
 
 describe("#StandardMetricsHandler", () => {
   let exportStub: sinon.SinonStub;

sdk/monitor/monitor-opentelemetry/test/internal/unit/standadMetrics.test.ts renamed to sdk/monitor/monitor-opentelemetry/test/internal/unit/metrics/standadMetrics.test.ts

@@ -0,0 +1,132 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+import * as assert from "assert";
+import * as sinon from "sinon";
+import { AzureFunctionsHook } from "../../../../src/traces/azureFnHook";
+import { TraceHandler } from "../../../../src/traces";
+import { Logger } from "../../../../src/shared/logging";
+import { AzureMonitorOpenTelemetryConfig } from "../../../../src/shared";
+
+describe("Library/AzureFunctionsHook", () => {
+  let sandbox: sinon.SinonSandbox;
+
+  before(() => {
+    sandbox = sinon.createSandbox();
+  });
+
+  afterEach(() => {
+    sandbox.restore();
+  });
+
+  it("Hook not added if not running in Azure Functions", () => {
+    const spy = sandbox.spy(Logger.getInstance(), "debug");
+    let hook = new AzureFunctionsHook();
+    assert.equal(hook["_functionsCoreModule"], undefined);
+    assert.ok(spy.called);
+    assert.equal(
+      spy.args[0][0],
+      "@azure/functions-core failed to load, not running in Azure Functions"
+    );
+  });
+
+  describe("AutoCollection/AzureFunctionsHook load fake Azure Functions Core", () => {
+    let originalRequire: any;
+
+    before(() => {
+      var Module = require("module");
+      originalRequire = Module.prototype.require;
+    });
+
+    afterEach(() => {
+      var Module = require("module");
+      Module.prototype.require = originalRequire;
+    });
+
+    it("Hook not added if using not supported programming model", () => {
+      var Module = require("module");
+      var preInvocationCalled = false;
+      Module.prototype.require = function () {
+        if (arguments[0] === "@azure/functions-core") {
+          return {
+            registerHook(name: string) {
+              if (name === "preInvocation") {
+                preInvocationCalled = true;
+              }
+            },
+            getProgrammingModel() {
+              return {
+                name: "@azure/functions",
+                version: "2.x",
+              };
+            },
+          };
+        }
+        return originalRequire.apply(this, arguments);
+      };
+      let azureFnHook = new AzureFunctionsHook();
+      assert.ok(azureFnHook, "azureFnHook");
+      assert.ok(!preInvocationCalled, "preInvocationCalled");
+    });
+
+    it("Pre Invokation Hook added if running in Azure Functions and context is propagated", () => {
+      let Module = require("module");
+      let preInvocationCalled = false;
+      let config = new AzureMonitorOpenTelemetryConfig();
+      config.azureMonitorExporterConfig.connectionString =
+        "InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333;";
+      let traceHandler = new TraceHandler(config);
+
+      Module.prototype.require = function () {
+        if (arguments[0] === "@azure/functions-core") {
+          return {
+            registerHook(name: string, callback: any) {
+              if (name === "preInvocation") {
+                preInvocationCalled = true;
+                let ctx = {
+                  res: { status: 400 },
+                  invocationId: "testinvocationId",
+                  traceContext: {
+                    traceparent: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+                    tracestate: "",
+                    attributes: {},
+                  },
+                };
+                let preInvocationContext: any = {
+                  inputs: [],
+                  functionCallback: () => {
+                    let span = traceHandler.getTracer().startSpan("test");
+                    // Context should be propagated here
+                    assert.equal(
+                      (span as any)["_spanContext"]["traceId"],
+                      "0af7651916cd43dd8448eb211c80319c"
+                    );
+                    assert.ok((span as any)["_spanContext"]["spanId"]);
+                  },
+                  hookData: {},
+                  appHookData: {},
+                  invocationContext: ctx,
+                };
+                // Azure Functions should call preinvocation callback
+                callback(preInvocationContext);
+                // Azure Functions should call customer function callback
+                preInvocationContext.functionCallback(null, null);
+              }
+            },
+
+            getProgrammingModel() {
+              return {
+                name: "@azure/functions",
+                version: "3.x",
+              };
+            },
+          };
+        }
+        return originalRequire.apply(this, arguments);
+      };
+      let azureFnHook = new AzureFunctionsHook();
+      assert.ok(azureFnHook, "azureFnHook");
+      assert.ok(preInvocationCalled, "preInvocationCalled");
+    });
+  });
+});

sdk/monitor/monitor-opentelemetry/test/internal/unit/traces/azureFnHook..test.ts

@@ -5,9 +5,9 @@ import * as assert from "assert";
 import * as sinon from "sinon";
 import { ExportResultCode } from "@opentelemetry/core";
 
-import { TraceHandler } from "../../../src/traces";
-import { MetricHandler } from "../../../src/metrics";
-import { AzureMonitorOpenTelemetryConfig } from "../../../src/shared";
+import { TraceHandler } from "../../../../src/traces";
+import { MetricHandler } from "../../../../src/metrics";
+import { AzureMonitorOpenTelemetryConfig } from "../../../../src/shared";
 import { HttpInstrumentationConfig } from "@opentelemetry/instrumentation-http";
 import { ReadableSpan, SpanProcessor } from "@opentelemetry/sdk-trace-base";
 import { Span } from "@opentelemetry/api";