diff --git a/packages/datadog-instrumentations/src/azure-event-hubs.js b/packages/datadog-instrumentations/src/azure-event-hubs.js index 2d339bcc7f1..7c7000c2105 100644 --- a/packages/datadog-instrumentations/src/azure-event-hubs.js +++ b/packages/datadog-instrumentations/src/azure-event-hubs.js @@ -6,6 +6,7 @@ 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({ @@ -13,6 +14,11 @@ addHook({ 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) @@ -20,7 +26,7 @@ addHook({ 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) diff --git a/packages/datadog-plugin-azure-event-hubs/src/producer.js b/packages/datadog-plugin-azure-event-hubs/src/producer.js index 44e38aba99d..32a8da43f00 100644 --- a/packages/datadog-plugin-azure-event-hubs/src/producer.js +++ b/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) { diff --git a/packages/datadog-plugin-azure-event-hubs/test/integration-test/client.spec.js b/packages/datadog-plugin-azure-event-hubs/test/integration-test/client.spec.js index 37bee92dce1..8c9ca99aa34 100644 --- a/packages/datadog-plugin-azure-event-hubs/test/integration-test/client.spec.js +++ b/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