feat(apim): add inferred span support for Azure API Managment (#7496)

* initial

* add apim to the inferrered proxy test

* rename method

* udpate method name in index.js

* clean up some comments and remove serverless type

* fix addRequestTags type

* remove object.fromentries

* remove unexpected import

* re add object.fromentries

* remove unwanted imports

* fix bug found by codex

* add test for bug caught by codex

Author: jcstorms1

packages/datadog-plugin-azure-functions/src/index.js

@@ -1,7 +1,6 @@
 'use strict'
 
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
-const serverless = require('../../dd-trace/src/plugins/util/serverless')
 const web = require('../../dd-trace/src/plugins/util/web')
 
 const triggerMap = {
@@ -26,19 +25,51 @@ class AzureFunctionsPlugin extends TracingPlugin {
   bindStart (ctx) {
     const meta = getMetaForTrigger(ctx)
     const triggerType = triggerMap[ctx.methodName]
+    const isHttpTrigger = triggerType === 'Http'
     const isMessagingService = (triggerType === 'ServiceBus' || triggerType === 'EventHubs')
-    const childOf = isMessagingService ? null : extractTraceContext(this._tracer, ctx)
-    const span = this.startSpan(this.operationName(), {
-      childOf,
-      service: this.serviceName(),
-      type: 'serverless',
-      meta,
-    }, ctx)
-
-    if (isMessagingService) {
-      setSpanLinks(triggerType, this.tracer, span, ctx)
-    }
 
+    let span
+
+    if (isHttpTrigger) {
+      const { httpRequest } = ctx
+      const path = (new URL(httpRequest.url)).pathname
+      const req = {
+        method: httpRequest.method,
+        headers: Object.fromEntries(httpRequest.headers),
+        url: path,
+      }
+      // Patch the request to create web context
+      const webContext = web.patch(req)
+      webContext.config = this.config
+      webContext.tracer = this.tracer
+      webContext.paths = [path]
+      // Creates a standard span and an inferred proxy span if headers are present
+      span = web.startServerlessSpanWithInferredProxy(
+        this.tracer,
+        this.config,
+        this.operationName(),
+        req,
+        ctx
+      )
+
+      span._integrationName = 'azure-functions'
+      span.context()._tags.component = 'azure-functions'
+      span.addTags(meta)
+      webContext.span = span
+      webContext.azureFunctionCtx = ctx
+      ctx.webContext = webContext
+    } else {
+      // For non-HTTP triggers, use standard flow
+      span = this.startSpan(this.operationName(), {
+        service: this.serviceName(),
+        type: 'serverless',
+        meta,
+      }, ctx)
+
+      if (isMessagingService) {
+        setSpanLinks(triggerType, this.tracer, span, ctx)
+      }
+    }
     ctx.span = span
     return ctx.currentStore
   }
@@ -48,24 +79,16 @@ class AzureFunctionsPlugin extends TracingPlugin {
     ctx.currentStore.span.setTag('error.message', ctx.error)
   }
 
-  asyncEnd (ctx) {
-    const { httpRequest, methodName, result = {} } = ctx
-    if (triggerMap[methodName] === 'Http') {
-      // If the method is an HTTP trigger, we need to patch the request and finish the span
-      const path = (new URL(httpRequest.url)).pathname
-      const req = {
-        method: httpRequest.method,
-        headers: Object.fromEntries(httpRequest.headers),
-        url: path,
+  asyncStart (ctx) {
+    const { methodName, result = {}, webContext } = ctx
+    const triggerType = triggerMap[methodName]
+
+    // For HTTP triggers, use web utilities to finish all spans (including inferred proxy)
+    if (triggerType === 'Http') {
+      if (webContext) {
+        webContext.res = { statusCode: result.status }
+        web.finishAll(webContext, 'serverless')
       }
-      const context = web.patch(req)
-      context.config = this.config
-      context.paths = [path]
-      context.res = { statusCode: result.status }
-      context.span = ctx.currentStore.span
-
-      serverless.finishSpan(context)
-    // Fallback for other trigger types
     } else {
       super.finish()
     }
@@ -80,6 +103,7 @@ function getMetaForTrigger ({ functionName, methodName, invocationContext }) {
   let meta = {
     'aas.function.name': functionName,
     'aas.function.trigger': mapTriggerTag(methodName),
+    'span.type': 'serverless',
   }
 
   if (triggerMap[methodName] === 'ServiceBus') {
@@ -112,14 +136,6 @@ function mapTriggerTag (methodName) {
   return triggerMap[methodName] || 'Unknown'
 }
 
-function extractTraceContext (tracer, ctx) {
-  if (triggerMap[ctx.methodName] === 'Http') {
-    return tracer.extract('http_headers', Object.fromEntries(ctx.httpRequest.headers))
-  }
-  // Returning null indicates that the span is a root span
-  return null
-}
-
 // message & messages & batch with cardinality of 1 == applicationProperties
 // messages with cardinality of many == applicationPropertiesArray
 function setSpanLinks (triggerType, tracer, span, ctx) {

packages/datadog-plugin-azure-functions/test/integration-test/http-test/client.spec.js

@@ -50,7 +50,14 @@ describe('esm', () => {
         assert.strictEqual(payload.length, 1)
         assert.ok(Array.isArray(payload[0]))
         assert.strictEqual(payload[0].length, 1)
-        assert.strictEqual(payload[0][0].name, 'azure.functions.invoke')
+
+        const span = payload[0][0]
+
+        assert.strictEqual(span.name, 'azure.functions.invoke')
+        assert.strictEqual(span.meta['_dd.integration'], 'azure-functions')
+        assert.strictEqual(span.meta.component, 'azure-functions')
+        assert.strictEqual(span.meta['http.route'], '/api/httptest')
+        assert.strictEqual(span.resource, 'GET /api/httptest')
       })
     }).timeout(60_000)
 

packages/dd-trace/src/plugins/util/inferred_proxy.js

@@ -14,12 +14,17 @@ const PROXY_HEADER_PATH = 'x-dd-proxy-path'
 const PROXY_HEADER_HTTPMETHOD = 'x-dd-proxy-httpmethod'
 const PROXY_HEADER_DOMAIN = 'x-dd-proxy-domain-name'
 const PROXY_HEADER_STAGE = 'x-dd-proxy-stage'
+const PROXY_HEADER_REGION = 'x-dd-proxy-region'
 
 const supportedProxies = {
   'aws-apigateway': {
     spanName: 'aws.apigateway',
     component: 'aws-apigateway',
   },
+  'azure-apim': {
+    spanName: 'azure.apim',
+    component: 'azure-apim',
+  },
 }
 
 function createInferredProxySpan (headers, childOf, tracer, reqCtx, traceCtx, config, startSpanHelper) {
@@ -53,6 +58,7 @@ function createInferredProxySpan (headers, childOf, tracer, reqCtx, traceCtx, co
       [HTTP_METHOD]: proxyContext.method,
       [HTTP_URL]: proxyContext.domainName + proxyContext.path,
       stage: proxyContext.stage,
+      region: proxyContext.region,
     },
   }, traceCtx, config)
 
@@ -91,6 +97,7 @@ function extractInferredProxyContext (headers) {
     stage: headers[PROXY_HEADER_STAGE],
     domainName: headers[PROXY_HEADER_DOMAIN],
     proxySystemName: headers[PROXY_HEADER_SYSTEM],
+    region: headers[PROXY_HEADER_REGION],
   }
 }
 

packages/dd-trace/src/plugins/util/serverless.js

@@ -119,7 +119,7 @@ const web = {
       context.span.context()._name = name
       span = context.span
     } else {
-      span = web.startChildSpan(tracer, config, name, req, traceCtx)
+      span = web.startServerlessSpanWithInferredProxy(tracer, config, name, req, traceCtx)
     }
 
     context.tracer = tracer
@@ -275,7 +275,7 @@ const web = {
     return context.middleware.at(-1)
   },
 
-  startChildSpan (tracer, config, name, req, traceCtx) {
+  startServerlessSpanWithInferredProxy (tracer, config, name, req, traceCtx) {
     const headers = req.headers
     const reqCtx = contexts.get(req)
     const { storage } = require('../../../../datadog-core')
@@ -338,12 +338,12 @@ const web = {
     }
   },
 
-  finishSpan (context) {
+  finishSpan (context, spanType) {
     const { req, res } = context
 
     if (context.finished && !req.stream) return
 
-    addRequestTags(context, this.TYPE)
+    addRequestTags(context, spanType)
     addResponseTags(context)
 
     context.config.hooks.request(context.span, req, res)
@@ -353,14 +353,14 @@ const web = {
     context.finished = true
   },
 
-  finishAll (context) {
+  finishAll (context, spanType) {
     for (const beforeEnd of context.beforeEnd) {
       beforeEnd()
     }
 
     web.finishMiddleware(context)
 
-    web.finishSpan(context)
+    web.finishSpan(context, spanType)
 
     finishInferredProxySpan(context)
   },
@@ -457,12 +457,13 @@ function reactivate (req, fn) {
 function addRequestTags (context, spanType) {
   const { req, span, inferredProxySpan, config } = context
   const url = extractURL(req)
+  const type = spanType ?? WEB
 
   span.addTags({
     [HTTP_URL]: obfuscateQs(config, url),
     [HTTP_METHOD]: req.method,
     [SPAN_KIND]: SERVER,
-    [SPAN_TYPE]: spanType,
+    [SPAN_TYPE]: type,
     [HTTP_USERAGENT]: req.headers['user-agent'],
   })