[SVLS-7305] Instrument Azure Event Hubs

.github/workflows/serverless.yml

@@ -87,7 +87,30 @@ jobs:
         with:
           suffix: plugins-${{ github.job }}-${{ matrix.node-version }}
       - uses: codecov/codecov-action@5a1091511ad55cbe89839c7260b706298ca349f7 # v5.5.1
-
+  azure-event-hubs:
+    runs-on: ubuntu-latest
+    services:
+        azurite:
+          image: mcr.microsoft.com/azure-storage/azurite:3.35.0
+          ports:
+            - "127.0.0.1:10000:10000"
+            - "127.0.0.1:10001:10001"
+            - "127.0.0.1:10002:10002"
+        azureeventhubsemulator:
+          image: mcr.microsoft.com/azure-messaging/eventhubs-emulator:2.1.0
+          ports:
+            - "127.0.0.1:5673:5672"
+            - "127.0.0.1:9092:9092"
+          env:
+            BLOB_SERVER: azurite
+            METADATA_SERVER: azurite
+            ACCEPT_EULA: "Y"
+    env:
+      PLUGINS: azure-event-hubs
+      SERVICES: azurite,azureeventhubsemulator
+    steps:
+        - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+        - uses: ./.github/actions/plugins/test
   azure-functions:
     runs-on: ubuntu-latest
     services:
@@ -97,6 +120,15 @@ jobs:
               - "127.0.0.1:10000:10000"
               - "127.0.0.1:10001:10001"
               - "127.0.0.1:10002:10002"
+        azureeventhubsemulator:
+          image: mcr.microsoft.com/azure-messaging/eventhubs-emulator:2.1.0
+          ports:
+            - "127.0.0.1:5673:5672"
+            - "127.0.0.1:9092:9092"
+          env:
+            BLOB_SERVER: azurite
+            METADATA_SERVER: azurite
+            ACCEPT_EULA: "Y"
         azureservicebusemulator:
           image: mcr.microsoft.com/azure-messaging/servicebus-emulator:1.1.2
           ports:
@@ -115,10 +147,16 @@ jobs:
               MSSQL_SA_PASSWORD: "Localtestpass1!"
     env:
       PLUGINS: azure-functions
-      SERVICES: azureservicebusemulator,azuresqledge
+      SERVICES: azureservicebusemulator,azuresqledge,azurite,azureeventhubsemulator
     steps:
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
-      - uses: ./.github/actions/node/active-lts
+      - run: |
+          docker cp \
+          ${{ github.workspace }}/packages/datadog-plugin-azure-functions/test/fixtures/eventhub-emulator-config.json \
+          ${{ job.services.azureeventhubsemulator.id }}:/Eventhubs_Emulator/ConfigFiles/Config.json
+      - run : |
+          docker restart ${{ job.services.azureeventhubsemulator.id }}
+      - uses: ./.github/actions/node/newest-maintenance-lts
       - uses: ./.github/actions/install
       - run: npm install -g [email protected]
       - run: echo "$(dirname $(which func))" >> $GITHUB_PATH

.gitignore

@@ -125,3 +125,6 @@ packages/dd-trace/test/plugins/versions/node_modules
 packages/dd-trace/test/plugins/versions/yarn.lock
 !packages/dd-trace/**/telemetry/logs
 packages/datadog-plugin-azure-functions/test/integration-test/fixtures/node_modules
+__azurite_db_queue__.json
+__azurite_db_queue_extent__.json
+__queuestorage__/AzuriteConfig

docker-compose.yml

@@ -4,25 +4,36 @@ services:
     ports:
       - "127.0.0.1:3000-3002:3000-3002"
   azurite:
-    image: mcr.microsoft.com/azure-storage/azurite:3.34.0
+    image: mcr.microsoft.com/azure-storage/azurite:3.35.0
     ports:
         - "127.0.0.1:10000:10000"
         - "127.0.0.1:10001:10001"
         - "127.0.0.1:10002:10002"
+  azureeventhubsemulator:
+    image: mcr.microsoft.com/azure-messaging/eventhubs-emulator:2.1.0
+    volumes:
+       - "./packages/datadog-plugin-azure-functions/test/fixtures/eventhub-emulator-config.json:/Eventhubs_Emulator/ConfigFiles/Config.json"
+    ports:
+      - "127.0.0.1:5673:5672"
+      - "127.0.0.1:9092:9092"
+    environment:
+      BLOB_SERVER: azurite
+      METADATA_SERVER: azurite
+      ACCEPT_EULA: "Y"
   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"
+      - "127.0.0.1:5672:5672"
+      - "127.0.0.1:5300:5300"
     environment:
-        ACCEPT_EULA: "Y"
-        MSSQL_SA_PASSWORD: "Localtestpass1!"
-        SQL_SERVER: azuresqledge
+      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"
+      - "127.0.0.1:1433:1433"
     environment:
         ACCEPT_EULA: "Y"
         MSSQL_SA_PASSWORD: "Localtestpass1!"

docs/test.ts

@@ -308,6 +308,7 @@ tracer.use('amqplib');
 tracer.use('avsc');
 tracer.use('aws-sdk');
 tracer.use('aws-sdk', awsSdkOptions);
+tracer.use('azure-event-hubs')
 tracer.use('azure-functions');
 tracer.use('bunyan');
 tracer.use('couchbase');

index.d.ts

@@ -166,6 +166,7 @@ interface Plugins {
   "apollo": tracer.plugins.apollo;
   "avsc": tracer.plugins.avsc;
   "aws-sdk": tracer.plugins.aws_sdk;
+  "azure-event-hubs": tracer.plugins.azure_event_hubs;
   "azure-functions": tracer.plugins.azure_functions;
   "azure-service-bus": tracer.plugins.azure_service_bus;
   "bunyan": tracer.plugins.bunyan;
@@ -1409,6 +1410,12 @@ declare namespace tracer {
       [key: string]: boolean | Object | undefined;
     }
 
+    /**
+     * This plugin automatically instruments the
+     * @azure/event-hubs module
+     */
+    interface azure_event_hubs extends Integration {}
+
     /**
      * This plugin automatically instruments the
      * @azure/functions module.
@@ -1420,6 +1427,7 @@ declare namespace tracer {
      * @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-event-hubs.js

@@ -0,0 +1,37 @@
+'use strict'
+
+const {
+  addHook
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+const dc = require('dc-polyfill')
+
+const producerCh = dc.tracingChannel('apm:azure-event-hubs:send')
+
+addHook({
+  name: '@azure/event-hubs',
+  versions: ['>=6.0.0']
+}, obj => {
+  const EventHubProducerClient = obj.EventHubProducerClient
+  shimmer.wrap(EventHubProducerClient.prototype, 'createBatch',
+    createBatch => async function () {
+      const batch = await createBatch.apply(this, arguments)
+      shimmer.wrap(batch, 'tryAdd',
+        tryAdd => function (eventData) {
+          const config = this._context.config
+          const functionName = tryAdd.name
+          return producerCh.tracePromise(
+            tryAdd,
+            { functionName, eventData, batch: this, config },
+            this, ...arguments)
+        })
+      return batch
+    })
+  shimmer.wrap(EventHubProducerClient.prototype, 'sendBatch',
+    sendBatch => function (eventData) {
+      const config = this._context.config
+      const functionName = sendBatch.name
+      return producerCh.tracePromise(sendBatch, { functionName, eventData, config }, this, ...arguments)
+    })
+  return obj
+})

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

@@ -23,6 +23,9 @@ addHook({ name: '@azure/functions', versions: ['>=4'], patchDefault: false }, (a
   shimmer.wrap(app, 'serviceBusQueue', wrapHandler)
   shimmer.wrap(app, 'serviceBusTopic', wrapHandler)
 
+  // Event Hub triggers
+  shimmer.wrap(app, 'eventHub', wrapHandler)
+
   return azureFunction
 })
 

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

@@ -5,6 +5,7 @@ module.exports = {
   '@apollo/gateway': () => require('../apollo'),
   'apollo-server-core': () => require('../apollo-server-core'),
   '@aws-sdk/smithy-client': () => require('../aws-sdk'),
+  '@azure/event-hubs': () => require('../azure-event-hubs'),
   '@azure/functions': () => require('../azure-functions'),
   '@azure/service-bus': () => require('../azure-service-bus'),
   '@cucumber/cucumber': () => require('../cucumber'),

packages/datadog-plugin-azure-event-hubs/src/index.js

@@ -0,0 +1,15 @@
+'use strict'
+
+const ProducerPlugin = require('./producer')
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
+
+class AzureEventHubsPlugin extends CompositePlugin {
+  static get id () { return 'azure-event-hubs' }
+  static get plugins () {
+    return {
+      producer: ProducerPlugin
+    }
+  }
+}
+
+module.exports = AzureEventHubsPlugin

packages/datadog-plugin-azure-event-hubs/src/producer.js

@@ -0,0 +1,82 @@
+'use strict'
+
+const { getEnvironmentVariable } = require('../../dd-trace/src/config-helper')
+const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
+
+class AzureEventHubsProducerPlugin extends ProducerPlugin {
+  static get id () { return 'azure-event-hubs' }
+  static get operation () { return 'send' }
+  static get prefix () { return 'tracing:apm:azure-event-hubs:send' }
+
+  bindStart (ctx) {
+    // we do not want to make these spans when batch linking is disabled.
+    if (!batchLinksAreEnabled() && ctx.functionName === 'tryAdd') {
+      return ctx.currentStore
+    }
+
+    const qualifiedNamespace = ctx.config.endpoint.replace('sb://', '').replace('/', '')
+    const entityPath = ctx.config.entityPath
+    const span = this.startSpan({
+      resource: entityPath,
+      type: 'messaging',
+      meta: {
+        component: 'azure-event-hubs',
+        'messaging.system': 'eventhubs',
+        'messaging.destination.name': entityPath,
+        'network.destination.name': qualifiedNamespace,
+      }
+    }, ctx)
+
+    if (ctx.functionName === 'tryAdd') {
+      span._spanContext._name = 'azure.eventhubs.create'
+      span.setTag('messaging.operation', 'create')
+
+      if (ctx.eventData.messageID !== undefined) {
+        span.setTag('message.id', ctx.eventData.messageID)
+      }
+
+      if (batchLinksAreEnabled()) {
+        ctx.batch._spanContexts.push(span.context())
+        injectTraceContext(this.tracer, span, ctx.eventData)
+      }
+    }
+
+    if (ctx.functionName === 'sendBatch') {
+      const eventData = ctx.eventData
+      const eventDataLength = eventData.length || eventData._context.connection._eventsCount
+      span.setTag('messaging.operation', 'send')
+      span.setTag('messaging.batch.message_count', eventDataLength)
+
+      if (eventData.constructor.name !== 'EventDataBatchImpl' && Array.isArray(eventData)) {
+        eventData.forEach(event => {
+          injectTraceContext(this.tracer, span, event)
+        })
+      } else {
+        if (batchLinksAreEnabled()) {
+          eventData._spanContexts.forEach(spanContext => {
+            span.addLink(spanContext)
+          })
+        }
+      }
+    }
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    super.finish()
+  }
+}
+
+function injectTraceContext (tracer, span, event) {
+  if (!event.properties) {
+    event.properties = {}
+  }
+  tracer.inject(span, 'text_map', event.properties)
+}
+
+function batchLinksAreEnabled () {
+  const eh = getEnvironmentVariable('DD_TRACE_AZURE_EVENTHUBS_BATCH_LINKS_ENABLED')
+  return eh !== 'false'
+}
+
+module.exports = AzureEventHubsProducerPlugin