fix: Stop using Azure's EventHubs batch spanContext map for storing DD span contexts

packages/datadog-instrumentations/src/azure-event-hubs.js

@@ -6,21 +6,27 @@ const {
 const shimmer = require('../../datadog-shimmer')
 const dc = require('dc-polyfill')
 
+const isItDefault = new WeakSet()
 const producerCh = dc.tracingChannel('apm:azure-event-hubs:send')
 
 addHook({
   name: '@azure/event-hubs',
   versions: ['>=6.0.0']
 }, obj => {
   const EventHubProducerClient = obj.EventHubProducerClient
+
+  if (!isItDefault.has(EventHubProducerClient.prototype)) {
+    isItDefault.add(EventHubProducerClient.prototype)
+  }
+
   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(
+          return producerCh.traceSync(
             tryAdd,
             { functionName, eventData, batch: this, config },
             this, ...arguments)

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

@@ -2,7 +2,7 @@
 
 const { getEnvironmentVariable } = require('../../dd-trace/src/config-helper')
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
-
+const spanContexts = new WeakMap()
 class AzureEventHubsProducerPlugin extends ProducerPlugin {
   static get id () { return 'azure-event-hubs' }
   static get operation () { return 'send' }
@@ -36,7 +36,12 @@ class AzureEventHubsProducerPlugin extends ProducerPlugin {
       }
 
       if (batchLinksAreEnabled()) {
-        ctx.batch._spanContexts.push(span.context())
+        const spanContext = spanContexts.get(ctx.batch)
+        if (spanContext) {
+          spanContext.push(span.context())
+        } else {
+          spanContexts.set(ctx.batch, [span.context()])
+        }
         injectTraceContext(this.tracer, span, ctx.eventData)
       }
     }
@@ -53,9 +58,12 @@ class AzureEventHubsProducerPlugin extends ProducerPlugin {
         })
       } else {
         if (batchLinksAreEnabled()) {
-          eventData._spanContexts.forEach(spanContext => {
-            span.addLink(spanContext)
-          })
+          const contexts = spanContexts.get(eventData)
+          if (contexts) {
+            for (const spanContext of contexts) {
+              span.addLink(spanContext)
+            }
+          }
         }
       }
     }
@@ -65,6 +73,10 @@ class AzureEventHubsProducerPlugin extends ProducerPlugin {
   asyncEnd (ctx) {
     super.finish()
   }
+
+  end (ctx) {
+    super.finish()
+  }
 }
 
 function injectTraceContext (tracer, span, event) {

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

@@ -15,7 +15,7 @@ const { withVersions } = require('../../../dd-trace/test/setup/mocha')
 const spawnEnv = { DD_TRACE_FLUSH_INTERVAL: '2000' }
 
 // TODO: Fix this test / esm issue
-describe.skip('esm', () => {
+describe('esm', () => {
   let agent
   let proc