fix(service-bus) Fix synchronous shimming of tryAddMessage (#7098)

* initial

* uncomment test

* Update packages/datadog-plugin-azure-service-bus/test/integration-test/tryAddMessageRegressionTest/client.spec.js

Co-authored-by: Ruben Bridgewater <ruben@bridgewater.de>

* Update packages/datadog-plugin-azure-service-bus/test/integration-test/tryAddMessageRegressionTest/client.spec.js

Co-authored-by: Ruben Bridgewater <ruben@bridgewater.de>

* remove chai import

* Merge branch 'storms/fix-azure-service-bus-try-add-message' of github.com:DataDog/dd-trace-js into storms/fix-azure-service-bus-try-add-message

---------

Co-authored-by: Ruben Bridgewater <ruben@bridgewater.de>

Author: jcstorms1

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

@@ -38,7 +38,7 @@ addHook({ name: '@azure/service-bus', versions: ['>=7.9.2'] }, (obj) => {
           shimmer.wrap(batch, 'tryAddMessage', tryAddMessage => function (msg) {
             const functionName = tryAddMessage.name
             const config = this._context.config
-            return producerCh.tracePromise(
+            return producerCh.traceSync(
               tryAddMessage, { config, functionName, batch, msg }, this, ...arguments)
           })
           return batch

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

@@ -75,6 +75,10 @@ class AzureServiceBusProducerPlugin extends ProducerPlugin {
   asyncEnd (ctx) {
     super.finish(ctx)
   }
+
+  end (ctx) {
+    super.finish(ctx)
+  }
 }
 
 function injectTraceContext (tracer, span, msg) {

…bus/test/integration-test/client.spec.js …ntegration-test/core-test/client.spec.jspackages/datadog-plugin-azure-service-bus/test/integration-test/client.spec.js renamed to packages/datadog-plugin-azure-service-bus/test/integration-test/core-test/client.spec.js packages/datadog-plugin-azure-service-bus/test/integration-test/client.spec.js renamed to packages/datadog-plugin-azure-service-bus/test/integration-test/core-test/client.spec.js

@@ -8,8 +8,9 @@ const {
   sandboxCwd,
   useSandbox,
   spawnPluginIntegrationTestProc
-} = require('../../../../integration-tests/helpers')
-const { withVersions } = require('../../../dd-trace/test/setup/mocha')
+} = require('../../../../../integration-tests/helpers')
+const { withVersions } = require('../../../../dd-trace/test/setup/mocha')
+
 const spawnEnv = { DD_TRACE_FLUSH_INTERVAL: '2000' }
 
 describe('esm', () => {
@@ -18,7 +19,7 @@ describe('esm', () => {
 
   withVersions('azure-service-bus', '@azure/service-bus', version => {
     useSandbox([`'@azure/service-bus@${version}'`], false, [
-      './packages/datadog-plugin-azure-service-bus/test/integration-test/*'])
+      './packages/datadog-plugin-azure-service-bus/test/integration-test/core-test/*'])
 
     beforeEach(async () => {
       agent = await new FakeAgent().start()

…ice-bus/test/integration-test/server.mjs …st/integration-test/core-test/server.mjspackages/datadog-plugin-azure-service-bus/test/integration-test/server.mjs renamed to packages/datadog-plugin-azure-service-bus/test/integration-test/core-test/server.mjs packages/datadog-plugin-azure-service-bus/test/integration-test/server.mjs renamed to packages/datadog-plugin-azure-service-bus/test/integration-test/core-test/server.mjs

@@ -0,0 +1,48 @@
+'use strict'
+
+const {
+  FakeAgent,
+  sandboxCwd,
+  useSandbox,
+  spawnPluginIntegrationTestProc
+} = require('../../../../../integration-tests/helpers')
+const { withVersions } = require('../../../../dd-trace/test/setup/mocha')
+const assert = require('assert')
+
+const spawnEnv = { DD_TRACE_FLUSH_INTERVAL: '2000' }
+
+describe('esm', () => {
+  let agent
+  let proc
+
+  withVersions('azure-service-bus', '@azure/service-bus', version => {
+    useSandbox([`'@azure/service-bus@${version}'`], false, [
+      './packages/datadog-plugin-azure-service-bus/test/integration-test/tryAddMessageRegressionTest/*'])
+
+    beforeEach(async () => {
+      agent = await new FakeAgent().start()
+      process.env.DD_TRACE_DISABLED_PLUGINS = 'amqplib,amqp10,rhea,net'
+    })
+
+    afterEach(async () => {
+      proc && proc.kill()
+      await agent.stop()
+    })
+
+    it('tryAddMessage returns a boolean, not a Promise', async () => {
+      const res = agent.assertMessageReceived(({ headers, payload }) => {
+        assert.ok(Array.isArray(payload))
+        assert.strictEqual(payload.length, 3)
+        // Verify we got the expected spans from the test
+        assert.strictEqual(payload[0][0].name, 'azure.servicebus.create')
+        assert.strictEqual(payload[1][0].name, 'azure.servicebus.create')
+        assert.strictEqual(payload[2][0].name, 'azure.servicebus.send')
+      })
+
+      // This test file will throw an error if tryAddMessage returns a Promise instead of a boolean
+      proc = await spawnPluginIntegrationTestProc(sandboxCwd(), 'server.mjs', agent.port, spawnEnv)
+
+      await res
+    }).timeout(60000)
+  })
+})

packages/datadog-plugin-azure-service-bus/test/integration-test/tryAddMessageRegressionTest/client.spec.js

@@ -0,0 +1,51 @@
+import 'dd-trace/init.js'
+import { ServiceBusClient } from '@azure/service-bus'
+
+const connectionString = 'Endpoint=sb://127.0.0.1;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=SAS_KEY_VALUE;UseDevelopmentEmulator=true;'
+const queueName = 'queue.1'
+
+const client = new ServiceBusClient(connectionString)
+const sender = client.createSender(queueName)
+
+const messages = [
+  { body: 'Test message 1' },
+  { body: 'Test message 2' }
+]
+
+// Test that tryAddMessage returns a boolean, not a Promise
+const batch = await sender.createMessageBatch()
+
+const result1 = batch.tryAddMessage(messages[0])
+const result2 = batch.tryAddMessage(messages[1])
+
+// Verify the return types - throw error if not correct
+if (typeof result1 !== 'boolean') {
+  throw new Error(`tryAddMessage should return a boolean, but returned ${typeof result1}`)
+}
+
+if (result1 instanceof Promise) {
+  throw new Error('tryAddMessage should not return a Promise')
+}
+
+if (typeof result2 !== 'boolean') {
+  throw new Error(`tryAddMessage should return a boolean, but returned ${typeof result2}`)
+}
+
+if (result2 instanceof Promise) {
+  throw new Error('tryAddMessage should not return a Promise')
+}
+
+// Verify the values are correct
+if (result1 !== true) {
+  throw new Error(`Expected first tryAddMessage to return true, got ${result1}`)
+}
+
+if (result2 !== true) {
+  throw new Error(`Expected second tryAddMessage to return true, got ${result2}`)
+}
+
+// Send the batch to complete the operation
+await sender.sendMessages(batch)
+
+await sender.close()
+await client.close()