[SVLS-7133] Instrument Azure Service Bus (#6067)

* initial

* plugin loading

* full distributed trace

* all span attributes set correctly

* add consumer tests

* update test

* finalize service bus test

* remove qpid from service bus test

* instrument and test topics

* remove localhost from connectionstring

* lint

* try removing localhost

* revert localhost removal

* try to use newer node versions

* revert node versioning for integration tests

* fix azure-functions test

* fix sql server & lint

* try actual name of sql service

* clean up azure functions

* formatting

* lint

Author: jcstorms1

.github/workflows/apm-integrations.yml

@@ -187,8 +187,52 @@ jobs:
 
   azure-functions:
     runs-on: ubuntu-latest
+    services:
+        azureservicebusemulator:
+          image: mcr.microsoft.com/azure-messaging/servicebus-emulator:1.1.2
+          ports:
+              - "127.0.0.1:5672:5672"
+              - "127.0.0.1:5300:5300"
+          env:
+              ACCEPT_EULA: "Y"
+              MSSQL_SA_PASSWORD: "Localtestpass1!"
+              SQL_SERVER: azuresqledge
+        azuresqledge:
+          image: mcr.microsoft.com/azure-sql-edge:1.0.7
+          ports:
+              - "127.0.0.1:1433:1433"
+          env:
+              ACCEPT_EULA: "Y"
+              MSSQL_SA_PASSWORD: "Localtestpass1!"
     env:
       PLUGINS: azure-functions
+      SERVICES: azureservicebusemulator,azuresqledge
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: ./.github/actions/plugins/test
+
+  azure-service-bus:
+    runs-on: ubuntu-latest
+    services:
+        azureservicebusemulator:
+          image: mcr.microsoft.com/azure-messaging/servicebus-emulator:1.1.2
+          ports:
+              - "127.0.0.1:5672:5672"
+              - "127.0.0.1:5300:5300"
+          env:
+              ACCEPT_EULA: "Y"
+              MSSQL_SA_PASSWORD: "Localtestpass1!"
+              SQL_SERVER: azuresqledge
+        azuresqledge:
+          image: mcr.microsoft.com/azure-sql-edge:1.0.7
+          ports:
+              - "127.0.0.1:1433:1433"
+          env:
+              ACCEPT_EULA: "Y"
+              MSSQL_SA_PASSWORD: "Localtestpass1!"
+    env:
+      PLUGINS: azure-service-bus
+      SERVICES: azureservicebusemulator,azuresqledge
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
       - uses: ./.github/actions/plugins/test

docker-compose.yml

@@ -3,6 +3,23 @@ services:
     image: aerospike:ce-6.4.0.3
     ports:
       - "127.0.0.1:3000-3002:3000-3002"
+  azureservicebusemulator:
+    image: mcr.microsoft.com/azure-messaging/servicebus-emulator:1.1.2
+    ports:
+        - "127.0.0.1:5672:5672"
+        - "127.0.0.1:5300:5300"
+    environment:
+        ACCEPT_EULA: "Y"
+        MSSQL_SA_PASSWORD: "Localtestpass1!"
+        SQL_SERVER: azuresqledge
+  azuresqledge:
+    image: mcr.microsoft.com/azure-sql-edge:1.0.7
+    platform: linux/amd64
+    ports:
+        - "127.0.0.1:1433:1433"
+    environment:
+        ACCEPT_EULA: "Y"
+        MSSQL_SA_PASSWORD: "Localtestpass1!"
   couchbase:
     image: ghcr.io/datadog/couchbase-server-sandbox:latest
     ports:

index.d.ts

@@ -167,6 +167,7 @@ interface Plugins {
   "avsc": tracer.plugins.avsc;
   "aws-sdk": tracer.plugins.aws_sdk;
   "azure-functions": tracer.plugins.azure_functions;
+  "azure-service-bus": tracer.plugins.azure_service_bus;
   "bunyan": tracer.plugins.bunyan;
   "cassandra-driver": tracer.plugins.cassandra_driver;
   "child_process": tracer.plugins.child_process;
@@ -1381,6 +1382,11 @@ declare namespace tracer {
     */
     interface azure_functions extends Instrumentation {}
 
+    /**
+     * This plugin automatically instruments the
+     * @azure/service-bus module
+     */
+    interface azure_service_bus extends Integration {}
     /**
      * This plugin patches the [bunyan](https://github.com/trentm/node-bunyan)
      * to automatically inject trace identifiers in log records when the

packages/datadog-instrumentations/src/azure-functions.js

@@ -11,13 +11,18 @@ const azureFunctionsChannel = dc.tracingChannel('datadog:azure:functions:invoke'
 addHook({ name: '@azure/functions', versions: ['>=4'] }, azureFunction => {
   const { app } = azureFunction
 
+  // Http triggers
   shimmer.wrap(app, 'deleteRequest', wrapHandler)
   shimmer.wrap(app, 'http', wrapHandler)
   shimmer.wrap(app, 'get', wrapHandler)
   shimmer.wrap(app, 'patch', wrapHandler)
   shimmer.wrap(app, 'post', wrapHandler)
   shimmer.wrap(app, 'put', wrapHandler)
 
+  // Service Bus triggers
+  shimmer.wrap(app, 'serviceBusQueue', wrapHandler)
+  shimmer.wrap(app, 'serviceBusTopic', wrapHandler)
+
   return azureFunction
 })
 

packages/datadog-instrumentations/src/azure-service-bus.js

@@ -0,0 +1,38 @@
+'use strict'
+
+const {
+  channel,
+  addHook
+} = require('./helpers/instrument')
+
+const shimmer = require('../../datadog-shimmer')
+
+const producerStartCh = channel('apm:azure-service-bus:send:start')
+const producerErrorCh = channel('apm:azure-service-bus:send:error')
+const producerFinishCh = channel('apm:azure-service-bus:send:finish')
+
+addHook({ name: '@azure/service-bus', versions: ['>=7.9.2'] }, (obj) => {
+  const ServiceBusClient = obj.ServiceBusClient
+  shimmer.wrap(ServiceBusClient.prototype, 'createSender', createSender => function (queueOrTopicName) {
+    const sender = createSender.apply(this, arguments)
+    shimmer.wrap(sender._sender, 'send', send => function (msg) {
+      const ctx = { sender, msg }
+      return producerStartCh.runStores(ctx, () => {
+        return send.apply(this, arguments)
+          .then(
+            response => {
+              producerFinishCh.publish(ctx)
+            },
+            error => {
+              ctx.error = error
+              producerErrorCh.publish(ctx)
+              producerFinishCh.publish(ctx)
+              throw error
+            }
+          )
+      })
+    })
+    return sender
+  })
+  return obj
+})

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -6,6 +6,7 @@ module.exports = {
   'apollo-server-core': () => require('../apollo-server-core'),
   '@aws-sdk/smithy-client': () => require('../aws-sdk'),
   '@azure/functions': () => require('../azure-functions'),
+  '@azure/service-bus': () => require('../azure-service-bus'),
   '@cucumber/cucumber': () => require('../cucumber'),
   '@playwright/test': () => require('../playwright'),
   '@elastic/elasticsearch': () => require('../elasticsearch'),

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

@@ -1,7 +1,6 @@
 'use strict'
 
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
-const { storage } = require('../../datadog-core')
 const serverless = require('../../dd-trace/src/plugins/util/serverless')
 const web = require('../../dd-trace/src/plugins/util/web')
 
@@ -11,7 +10,9 @@ const triggerMap = {
   get: 'Http',
   patch: 'Http',
   post: 'Http',
-  put: 'Http'
+  put: 'Http',
+  serviceBusQueue: 'ServiceBus',
+  serviceBusTopic: 'ServiceBus',
 }
 
 class AzureFunctionsPlugin extends TracingPlugin {
@@ -22,24 +23,16 @@ class AzureFunctionsPlugin extends TracingPlugin {
   static get prefix () { return 'tracing:datadog:azure:functions:invoke' }
 
   bindStart (ctx) {
-    const { functionName, methodName, httpRequest } = ctx
-    const store = storage('legacy').getStore()
-    // httpRequest.headers is a map
-    const childOf = this._tracer.extract('http_headers', Object.fromEntries(httpRequest.headers))
+    const childOf = extractTraceContext(this._tracer, ctx)
+    const meta = getMetaForTrigger(ctx)
     const span = this.startSpan(this.operationName(), {
       childOf,
       service: this.serviceName(),
       type: 'serverless',
-      meta: {
-        'aas.function.name': functionName,
-        'aas.function.trigger': mapTriggerTag(methodName)
-      }
-    }, false)
+      meta,
+    }, ctx)
 
     ctx.span = span
-    ctx.parentStore = store
-    ctx.currentStore = { ...store, span }
-
     return ctx.currentStore
   }
 
@@ -49,30 +42,66 @@ class AzureFunctionsPlugin extends TracingPlugin {
   }
 
   asyncEnd (ctx) {
-    const { httpRequest, result = {} } = ctx
-    const path = (new URL(httpRequest.url)).pathname
-    const req = {
-      method: httpRequest.method,
-      headers: Object.fromEntries(httpRequest.headers),
-      url: path
-    }
-
-    const context = web.patch(req)
-    context.config = this.config
-    context.paths = [path]
-    context.res = { statusCode: result.status }
-    context.span = ctx.currentStore.span
+    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
+      }
+      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)
+      serverless.finishSpan(context)
+    // Fallback for other trigger types
+    } else {
+      super.finish()
+    }
   }
 
   configure (config) {
     return super.configure(web.normalizeConfig(config))
   }
 }
 
+function getMetaForTrigger ({ functionName, methodName, invocationContext }) {
+  let meta = {
+    'aas.function.name': functionName,
+    'aas.function.trigger': mapTriggerTag(methodName)
+  }
+
+  if (triggerMap[methodName] === 'ServiceBus') {
+    const triggerEntity = invocationContext.options.trigger.queueName || invocationContext.options.trigger.topicName
+    meta = {
+      ...meta,
+      'messaging.message_id': invocationContext.triggerMetadata.messageId,
+      'messaging.operation': 'receive',
+      'messaging.system': 'servicebus',
+      'messaging.destination.name': triggerEntity,
+      'resource.name': `ServiceBus ${functionName}`,
+      'span.kind': 'consumer'
+    }
+  }
+
+  return meta
+}
+
 function mapTriggerTag (methodName) {
   return triggerMap[methodName] || 'Unknown'
 }
 
+function extractTraceContext (tracer, ctx) {
+  switch (String(triggerMap[ctx.methodName])) {
+    case 'Http':
+      return tracer.extract('http_headers', Object.fromEntries(ctx.httpRequest.headers))
+    case 'ServiceBus':
+      return tracer.extract('text_map', ctx.invocationContext.triggerMetadata.applicationProperties)
+  }
+}
+
 module.exports = AzureFunctionsPlugin

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

@@ -20,8 +20,13 @@ describe('esm', () => {
   withVersions('azure-functions', '@azure/functions', NODE_MAJOR < 20 ? '<4.7.3' : '*', version => {
     before(async function () {
       this.timeout(50000)
-      sandbox = await createSandbox([`@azure/functions@${version}`, 'azure-functions-core-tools@4'], false,
-        ['./packages/datadog-plugin-azure-functions/test/integration-test/fixtures/*'])
+      sandbox = await createSandbox([
+        `@azure/functions@${version}`,
+        'azure-functions-core-tools@4',
+        '@azure/service-bus@7.9.2'
+      ],
+      false,
+      ['./packages/datadog-plugin-azure-functions/test/integration-test/fixtures/*'])
     })
 
     after(async function () {
@@ -54,15 +59,51 @@ describe('esm', () => {
       })
     }).timeout(50000)
 
-    it('propagates context to child requests', async () => {
+    it('propagates context to child http requests', async () => {
       const envArgs = {
         PATH: `${sandbox.folder}/node_modules/azure-functions-core-tools/bin:${process.env.PATH}`
       }
       proc = await spawnPluginIntegrationTestProc(sandbox.folder, 'func', ['start'], agent.port, undefined, envArgs)
 
       return curlAndAssertMessage(agent, 'http://127.0.0.1:7071/api/httptest2', ({ headers, payload }) => {
         assert.strictEqual(payload.length, 2)
-        assert.strictEqual(payload[0][0].parent_id, payload[1][1].span_id)
+        assert.strictEqual(payload[1][0].span_id, payload[1][1].parent_id)
+      })
+    }).timeout(50000)
+
+    it('propagates context through a service bus queue', async () => {
+      const envArgs = {
+        PATH: `${sandbox.folder}/node_modules/azure-functions-core-tools/bin:${process.env.PATH}`
+      }
+      proc = await spawnPluginIntegrationTestProc(sandbox.folder, 'func', ['start'], agent.port, undefined, envArgs)
+
+      return curlAndAssertMessage(agent, 'http://127.0.0.1:7071/api/httptest3', ({ headers, payload }) => {
+        assert.strictEqual(payload.length, 3)
+        assert.strictEqual(payload[1][1].span_id, payload[2][0].parent_id)
+        assert.strictEqual(payload[2][0].name, 'azure.functions.invoke')
+        assert.strictEqual(payload[2][0].resource, 'ServiceBus queueTest')
+        assert.strictEqual(payload[2][0].meta['messaging.destination.name'], 'queue.1')
+        assert.strictEqual(payload[2][0].meta['messaging.operation'], 'receive')
+        assert.strictEqual(payload[2][0].meta['messaging.system'], 'servicebus')
+        assert.strictEqual(payload[2][0].meta['span.kind'], 'consumer')
+      })
+    }).timeout(50000)
+
+    it('propagates context through a service bus topic', async () => {
+      const envArgs = {
+        PATH: `${sandbox.folder}/node_modules/azure-functions-core-tools/bin:${process.env.PATH}`
+      }
+      proc = await spawnPluginIntegrationTestProc(sandbox.folder, 'func', ['start'], agent.port, undefined, envArgs)
+
+      return curlAndAssertMessage(agent, 'http://127.0.0.1:7071/api/httptest4', ({ headers, payload }) => {
+        assert.strictEqual(payload.length, 3)
+        assert.strictEqual(payload[1][1].span_id, payload[2][0].parent_id)
+        assert.strictEqual(payload[2][0].name, 'azure.functions.invoke')
+        assert.strictEqual(payload[2][0].resource, 'ServiceBus topicTest')
+        assert.strictEqual(payload[2][0].meta['messaging.destination.name'], 'topic.1')
+        assert.strictEqual(payload[2][0].meta['messaging.operation'], 'receive')
+        assert.strictEqual(payload[2][0].meta['messaging.system'], 'servicebus')
+        assert.strictEqual(payload[2][0].meta['span.kind'], 'consumer')
       })
     }).timeout(50000)
   })

packages/datadog-plugin-azure-functions/test/integration-test/fixtures/local.settings.json

@@ -3,6 +3,7 @@
   "Values": {
     "FUNCTIONS_WORKER_RUNTIME": "node",
     "AzureWebJobsFeatureFlags": "EnableWorkerIndexing",
-    "AzureWebJobsStorage": ""
+    "AzureWebJobsStorage": "UseDevelopmentStorage=true",
+    "MyServiceBus": "Endpoint=sb://127.0.0.1;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=SAS_KEY_VALUE;UseDevelopmentEmulator=true;"
   }
 }

packages/datadog-plugin-azure-functions/test/integration-test/fixtures/package.json

@@ -5,11 +5,5 @@
   "main": "src/functions/server.mjs",
   "scripts": {
     "start": "func start"
-  },
-  "dependencies": {
-    "@azure/functions": "^4.6.0"
-  },
-  "devDependencies": {
-    "azure-functions-core-tools": "^4.x"
   }
 }