fix: Azure Service Bus batching memory leak (#6917)

* initial

* fix issues with esm

* clean up

* Update packages/datadog-plugin-azure-service-bus/src/producer.js

Co-authored-by: Thomas Watson <thomas.watson@datadoghq.com>

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

Co-authored-by: Thomas Watson <thomas.watson@datadoghq.com>

* Update packages/datadog-plugin-azure-service-bus/src/producer.js

Co-authored-by: Thomas Watson <thomas.watson@datadoghq.com>

---------

Co-authored-by: Thomas Watson <thomas.watson@datadoghq.com>

Author: jcstorms1

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

@@ -8,57 +8,64 @@ const shimmer = require('../../datadog-shimmer')
 const dc = require('dc-polyfill')
 
 const producerCh = dc.tracingChannel('apm:azure-service-bus:send')
+const isItDefault = new WeakSet()
 
-addHook({ name: '@azure/service-bus', versions: ['>=7.9.2'], patchDefault: false }, (obj) => {
+addHook({ name: '@azure/service-bus', versions: ['>=7.9.2'] }, (obj) => {
   const ServiceBusClient = obj.ServiceBusClient
-  let didItShim = false
   shimmer.wrap(ServiceBusClient.prototype, 'createSender',
     createSender => function (queueOrTopicName) {
       const sender = createSender.apply(this, arguments)
-      if (didItShim) return sender
       const senderPrototype = sender.constructor.prototype
       const senderSenderPrototype = sender._sender.constructor.prototype
-      shimmer.wrap(senderPrototype, 'scheduleMessages', scheduleMessages =>
-        function (msg, scheduledEnqueueTimeUtc) {
-          const functionName = scheduleMessages.name
+
+      if (!isItDefault.has(senderPrototype)) {
+        isItDefault.add(senderPrototype)
+
+        shimmer.wrap(senderPrototype, 'scheduleMessages', scheduleMessages =>
+          function (msg, scheduledEnqueueTimeUtc) {
+            const functionName = scheduleMessages.name
+            const config = this._context.config
+            const entityPath = this._entityPath
+            return producerCh.tracePromise(
+              scheduleMessages,
+              { config, entityPath, functionName, msg, scheduledEnqueueTimeUtc },
+              this, ...arguments
+            )
+          })
+
+        shimmer.wrap(senderPrototype, 'createMessageBatch', createMessageBatch => async function () {
+          const batch = await createMessageBatch.apply(this, arguments)
+          shimmer.wrap(batch, 'tryAddMessage', tryAddMessage => function (msg) {
+            const functionName = tryAddMessage.name
+            const config = this._context.config
+            return producerCh.tracePromise(
+              tryAddMessage, { config, functionName, batch, msg }, this, ...arguments)
+          })
+          return batch
+        })
+      }
+
+      if (!isItDefault.has(senderSenderPrototype)) {
+        isItDefault.add(senderSenderPrototype)
+
+        shimmer.wrap(senderSenderPrototype, 'send', send => function (msg) {
+          const functionName = send.name
           const config = this._context.config
-          const entityPath = this._entityPath
+          const entityPath = this.entityPath
           return producerCh.tracePromise(
-            scheduleMessages,
-            { config, entityPath, functionName, msg, scheduledEnqueueTimeUtc },
-            this, ...arguments
+            send, { config, entityPath, functionName, msg }, this, ...arguments
           )
         })
 
-      shimmer.wrap(senderPrototype, 'createMessageBatch', createMessageBatch => async function () {
-        const batch = await createMessageBatch.apply(this, arguments)
-        shimmer.wrap(batch.constructor.prototype, 'tryAddMessage', tryAddMessage => function (msg) {
-          const functionName = tryAddMessage.name
+        shimmer.wrap(senderSenderPrototype, 'sendBatch', sendBatch => function (msg) {
+          const functionName = sendBatch.name
           const config = this._context.config
+          const entityPath = this.entityPath
           return producerCh.tracePromise(
-            tryAddMessage, { config, functionName, batch, msg }, this, ...arguments)
+            sendBatch, { config, entityPath, functionName, msg }, this, ...arguments
+          )
         })
-        return batch
-      })
-
-      shimmer.wrap(senderSenderPrototype, 'send', send => function (msg) {
-        const functionName = send.name
-        const config = this._context.config
-        const entityPath = this.entityPath
-        return producerCh.tracePromise(
-          send, { config, entityPath, functionName, msg }, this, ...arguments
-        )
-      })
-
-      shimmer.wrap(senderSenderPrototype, 'sendBatch', sendBatch => function (msg) {
-        const functionName = sendBatch.name
-        const config = this._context.config
-        const entityPath = this.entityPath
-        return producerCh.tracePromise(
-          sendBatch, { config, entityPath, functionName, msg }, this, ...arguments
-        )
-      })
-      didItShim = true
+      }
       return sender
     })
   return obj

packages/datadog-plugin-azure-service-bus/src/producer.js

@@ -2,6 +2,7 @@
 
 const { getEnvironmentVariable } = require('../../dd-trace/src/config-helper')
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
+const spanContexts = new WeakMap()
 
 class AzureServiceBusProducerPlugin extends ProducerPlugin {
   static get id () { return 'azure-service-bus' }
@@ -36,7 +37,12 @@ class AzureServiceBusProducerPlugin 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.msg)
       }
     }
@@ -47,9 +53,12 @@ class AzureServiceBusProducerPlugin extends ProducerPlugin {
       if (isBatch) {
         span.setTag('messaging.batch.message_count', messages.count)
         if (batchLinksAreEnabled()) {
-          messages._spanContexts.forEach(spanContext => {
-            span.addLink(spanContext)
-          })
+          const contexts = spanContexts.get(messages)
+          if (contexts) {
+            for (const spanContext of contexts) {
+              span.addLink(spanContext)
+            }
+          }
         }
       } else if (Array.isArray(messages)) {
         span.setTag('messaging.batch.message_count', messages.length)
@@ -64,7 +73,7 @@ class AzureServiceBusProducerPlugin extends ProducerPlugin {
   }
 
   asyncEnd (ctx) {
-    super.finish()
+    super.finish(ctx)
   }
 }