Fixing issue with status in auto generated request in Azure Fn (#1046)

AutoCollection/AzureFunctionsHook.ts

@@ -1,72 +1,104 @@
 import { Context, HttpRequest } from "../Library/Functions";
 import Logging = require("../Library/Logging");
 import TelemetryClient = require("../Library/TelemetryClient");
-import { CorrelationContextManager } from "./CorrelationContextManager";
+import { CorrelationContext, CorrelationContextManager } from "./CorrelationContextManager";
 
 
 /** Node.js Azure Functions handle incoming HTTP requests before Application Insights SDK is available,
  * this code generate incoming request telemetry and generate correlation context to be used 
  * by outgoing requests and other telemetry, we rely on hooks provided by Azure Functions
 */
-export class AutoCollectAzureFunctions {
+export class AzureFunctionsHook {
     private _client: TelemetryClient;
     private _functionsCoreModule: any;
+    private _autoGenerateIncomingRequests: boolean;
     private _preInvocationHook: any;
 
     constructor(client: TelemetryClient) {
         this._client = client;
+        this._autoGenerateIncomingRequests = false;
         try {
             this._functionsCoreModule = require('@azure/functions-core');
         }
         catch (error) {
-            Logging.info("AutoCollectAzureFunctions failed to load, not running in Azure Functions");
+            Logging.info("AzureFunctionsHook failed to load, not running in Azure Functions");
+            return;
         }
+        this._addPreInvocationHook();
     }
 
     public enable(isEnabled: boolean) {
-        if (this._functionsCoreModule) {
-            if (isEnabled) {
-                this._addPreInvocationHook();
-            } else {
-                this._removePreInvocationHook();
-            }
-        }
+        this._autoGenerateIncomingRequests = isEnabled;
     }
 
     public dispose() {
         this.enable(false);
+        this._removePreInvocationHook();
         this._functionsCoreModule = undefined;
     }
 
     private _addPreInvocationHook() {
-        // Only add hook once
         if (!this._preInvocationHook) {
-            this._preInvocationHook = this._functionsCoreModule.registerHook('preInvocation', (context: any) => {
-                const originalCallback = context.functionCallback;
-                context.functionCallback = async (context: Context, req: HttpRequest) => {
-                    const startTime = Date.now(); // Start trackRequest timer
-                    // Start an AI Correlation Context using the provided Function context
-                    const correlationContext = CorrelationContextManager.startOperation(context, req);
-                    if (correlationContext) {
-                        CorrelationContextManager.wrapCallback(async () => {
-                            originalCallback(context, req);
-                            this._client.trackRequest({
-                                name: context?.req?.method + " " + context.req?.url,
-                                resultCode: context?.res?.status,
-                                success: true,
-                                url: (req as HttpRequest)?.url,
-                                time: new Date(startTime),
-                                duration: Date.now() - startTime,
-                                id: correlationContext.operation?.parentId,
-                            });
-                            this._client.flush();
-                        }, correlationContext)();
-                    }
+            this._preInvocationHook = this._functionsCoreModule.registerHook('preInvocation', async (preInvocationContext: any) => {
+                const originalCallback = preInvocationContext.functionCallback;
+                preInvocationContext.functionCallback = async (ctx: Context, request: HttpRequest) => {
+                    this._propagateContext(ctx, request, originalCallback);
                 };
             });
         }
     }
 
+    private async _propagateContext(ctx: Context, request: HttpRequest, originalCallback: any) {
+        // Update context to use Azure Functions one
+        let extractedContext: CorrelationContext = null;
+        try {
+            // Start an AI Correlation Context using the provided Function context
+            extractedContext = CorrelationContextManager.startOperation(ctx, request);
+        }
+        catch (err) {
+            Logging.warn("Failed to propagate context in Azure Functions", err);
+            originalCallback(ctx, request);
+            return;
+        }
+        if (!extractedContext) {
+            // Correlation Context could be disabled causing this to be null
+            Logging.warn("Failed to create context in Azure Functions");
+            originalCallback(ctx, request);
+            return;
+        }
+
+        CorrelationContextManager.wrapCallback(async () => {
+            const startTime = Date.now(); // Start trackRequest timer
+            originalCallback(ctx, request);
+            try {
+                if (this._autoGenerateIncomingRequests) {
+                    let statusCode = 200; //Default
+                    if (ctx.res) {
+                        if (ctx.res.statusCode) {
+                            statusCode = ctx.res.statusCode;
+                        }
+                        else if (ctx.res.status) {
+                            statusCode = ctx.res.status;
+                        }
+                    }
+                    this._client.trackRequest({
+                        name: request.method + " " + request.url,
+                        resultCode: statusCode,
+                        success: statusCode == 200,
+                        url: request.url,
+                        time: new Date(startTime),
+                        duration: Date.now() - startTime,
+                        id: extractedContext.operation?.parentId,
+                    });
+                    this._client.flush();
+                }
+            }
+            catch (err) {
+                Logging.warn("Error creating automatic incoming request in Azure Functions", err);
+            }
+        }, extractedContext)();
+    }
+
     private _removePreInvocationHook() {
         if (this._preInvocationHook) {
             this._preInvocationHook.dispose();

Declarations/Interfaces.ts

@@ -150,8 +150,8 @@ export interface IBaseConfig {
      */
     enableWebInstrumentation: boolean;
     /**
-     * Enable automatic incoming request tracking and correct correlation when using Azure Functions
-     */
+    * Enable automatic incoming request tracking when running in Azure Functions
+    */
     enableAutoCollectAzureFunctions: boolean;
     /**
     * Application Insights resource connection string for web instrumentation and automatic monitoring

Tests/AutoCollection/AzureFunctionsHook.tests.ts

@@ -1,51 +1,98 @@
 import assert = require("assert");
 import sinon = require("sinon");
 import { TelemetryClient } from "../../applicationinsights";
+import { AzureFunctionsHook } from "../../AutoCollection/AzureFunctionsHook";
+import { CorrelationContextManager } from "../../AutoCollection/CorrelationContextManager";
+import { HttpRequest } from "../../Library/Functions";
+import Logging = require("../../Library/Logging");
 
-import { AutoCollectAzureFunctions } from "../../AutoCollection/AzureFunctionsHook";
 
-const testModule = {
-    registerHook(type: string, hook: any) {
+class TestFunctionCore {
+    public registerCalled: boolean = false;
+    public hookName: string;
 
+    registerHook(name: string, func: any) {
+        this.registerCalled = true;
+        this.hookName = name;
     }
-};
+}
 
-describe("AutoCollection/AutoCollectAzureFunctions", () => {
+describe("AutoCollection/AzureFunctionsHook", () => {
+    let sandbox: sinon.SinonSandbox;
+    let client: TelemetryClient;
 
-    it("constructor", () => {
-        let client = new TelemetryClient("1aa11111-bbbb-1ccc-8ddd-eeeeffff3333");
-        let auto = new AutoCollectAzureFunctions(client);
-        assert.equal(auto["_functionsCoreModule"], undefined, "Module is not available so it should be undefined unless running in AzFn env");
+    before(() => {
+        client = new TelemetryClient("1aa11111-bbbb-1ccc-8ddd-eeeeffff3333");
+        sandbox = sinon.sandbox.create();
     });
 
-    it("enable", () => {
-        let client = new TelemetryClient("1aa11111-bbbb-1ccc-8ddd-eeeeffff3333");
-        let auto = new AutoCollectAzureFunctions(client);
-        auto["_functionsCoreModule"] = testModule;
-        const addStub = sinon.stub(auto, "_addPreInvocationHook");
-        const removeStub = sinon.stub(auto, "_removePreInvocationHook");
-        auto.enable(true);
-        assert.ok(removeStub.notCalled);
-        assert.ok(addStub.calledOnce);
+    afterEach(() => {
+        sandbox.restore();
+        CorrelationContextManager.enable(false);
     });
 
-    it("disable", () => {
-        let client = new TelemetryClient("1aa11111-bbbb-1ccc-8ddd-eeeeffff3333");
-        let auto = new AutoCollectAzureFunctions(client);
-        auto["_functionsCoreModule"] = testModule;
-        const addStub = sinon.stub(auto, "_addPreInvocationHook");
-        const removeStub = sinon.stub(auto, "_removePreInvocationHook");
-        auto.enable(false);
-        assert.ok(removeStub.calledOnce);
-        assert.ok(addStub.notCalled);
+
+    it("Hook not added if not running in Azure Functions", () => {
+        const spy = sandbox.spy(Logging, "info");
+        let hook = new AzureFunctionsHook(client);
+        assert.equal(hook["_functionsCoreModule"], undefined);
+        assert.ok(spy.called);
+        assert.equal(spy.args[0][0], "AzureFunctionsHook failed to load, not running in Azure Functions");
+    });
+
+    it("Hook added if running in Azure Functions", () => {
+        let hook = new AzureFunctionsHook(client);
+        let testCore = new TestFunctionCore();
+        hook["_functionsCoreModule"] = testCore;
+        hook["_addPreInvocationHook"]();
+        assert.ok(testCore.registerCalled);
+        assert.equal(testCore.hookName, "preInvocation");
+    });
+
+    it("enable/disable", () => {
+        let hook = new AzureFunctionsHook(client);
+        hook.enable(true);
+        assert.equal(hook["_autoGenerateIncomingRequests"], true);
+        hook.enable(false);
+        assert.equal(hook["_autoGenerateIncomingRequests"], false);
     });
 
-    it("_addPreInvocationHook", () => {
-        let client = new TelemetryClient("1aa11111-bbbb-1ccc-8ddd-eeeeffff3333");
-        let auto = new AutoCollectAzureFunctions(client);
-        const registerHook = sinon.stub(testModule, "registerHook");
-        auto["_functionsCoreModule"] = testModule;
-        auto["_addPreInvocationHook"]();
-        assert.ok(registerHook.calledOnce);
+    it("Context propagation", () => {
+        CorrelationContextManager.enable(true);
+        let hook = new AzureFunctionsHook(client);
+        hook.enable(true);
+        let flushStub = sandbox.stub(hook["_client"], "flush");
+        let trackRequestSpy = sandbox.stub(hook["_client"], "trackRequest");
+        let contextSpy = sandbox.spy(CorrelationContextManager, "wrapCallback");
+        let ctx = {
+            res: { "status": 400 },
+            traceContext: {
+                traceparent: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01",
+                tracestate: "",
+                attributes: {}
+            }
+        };
+        let request: HttpRequest = {
+            method: "HEAD",
+            url: "test.com",
+            headers: { "": "" }
+        };
+        let originalCallbackCalled = false;
+        let originalCallback = () => { originalCallbackCalled = true };
+        hook["_propagateContext"](ctx as any, request, originalCallback);
+        assert.ok(contextSpy.called);
+        assert.ok(originalCallbackCalled);
+        assert.ok(flushStub.called);
+        assert.ok(trackRequestSpy.called);
+        let propagatedContext = contextSpy.args[0][1];
+        assert.equal(propagatedContext.operation.id, "0af7651916cd43dd8448eb211c80319c");
+        assert.equal(propagatedContext.operation.name, "HEAD /");
+        assert.equal(propagatedContext.operation.parentId, "|0af7651916cd43dd8448eb211c80319c.b7ad6b7169203331.");
+        let incomingRequest = trackRequestSpy.args[0][0];
+        assert.equal(incomingRequest.id, "|0af7651916cd43dd8448eb211c80319c.b7ad6b7169203331.");
+        assert.equal(incomingRequest.name, "HEAD test.com");
+        assert.equal(incomingRequest.resultCode, 400);
+        assert.equal(incomingRequest.success, false);
+        assert.equal(incomingRequest.url, "test.com");
     });
 });

applicationinsights.ts

@@ -13,7 +13,7 @@ import QuickPulseClient = require("./Library/QuickPulseStateManager");
 import { IncomingMessage } from "http";
 import { SpanContext } from "@opentelemetry/api";
 import { AutoCollectNativePerformance, IDisabledExtendedMetrics } from "./AutoCollection/NativePerformance";
-import { AutoCollectAzureFunctions } from "./AutoCollection/AzureFunctionsHook";
+import { AzureFunctionsHook } from "./AutoCollection/AzureFunctionsHook";
 
 // We export these imports so that SDK users may use these classes directly.
 // They're exposed using "export import" so that types are passed along as expected
@@ -85,7 +85,7 @@ let _webSnippet: WebSnippet;
 let _nativePerformance: AutoCollectNativePerformance;
 let _serverRequests: AutoCollectHttpRequests;
 let _clientRequests: AutoCollectHttpDependencies;
-let _azureFunctions: AutoCollectAzureFunctions;
+let _azureFunctions: AzureFunctionsHook;
 
 let _isStarted = false;
 
@@ -122,7 +122,7 @@ export function setup(setupString?: string) {
         if (!_nativePerformance) {
             _nativePerformance = new AutoCollectNativePerformance(defaultClient);
         }
-        _azureFunctions = new AutoCollectAzureFunctions(defaultClient);
+        _azureFunctions = new AzureFunctionsHook(defaultClient);
     } else {
         Logging.info("The default client is already setup");
     }
@@ -177,7 +177,7 @@ function _initializeConfig() {
     _forceClsHooked = defaultClient.config.enableUseAsyncHooks !== undefined ? defaultClient.config.enableUseAsyncHooks : _forceClsHooked;
     _isSnippetInjection = defaultClient.config.enableWebInstrumentation !== undefined ? defaultClient.config.enableWebInstrumentation : _isSnippetInjection;
     _isSnippetInjection = defaultClient.config.enableAutoWebSnippetInjection === true ? true : _isSnippetInjection;
-    _isAzureFunctions = defaultClient.config.enableAutoCollectAzureFunctions !== undefined ? defaultClient.config.enableAutoCollectAzureFunctions : _isPerformance;
+    _isAzureFunctions = defaultClient.config.enableAutoCollectAzureFunctions !== undefined ? defaultClient.config.enableAutoCollectAzureFunctions : _isAzureFunctions;
     const extendedMetricsConfig = AutoCollectNativePerformance.parseEnabled(defaultClient.config.enableAutoCollectExtendedMetrics, defaultClient.config);
     _isNativePerformance = extendedMetricsConfig.isEnabled;
     _disabledExtendedMetrics = extendedMetricsConfig.disabledMetrics;