Fix a bug with tryAddMessage being changed to async which causes messages to be dropped

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/test/integration-test/client.spec.js

@@ -150,6 +150,22 @@ describe('esm', () => {
 
       await res
     }).timeout(60000)
+
+    it('tryAddMessage returns a boolean, not a Promise', async () => {
+      const res = agent.assertMessageReceived(({ headers, payload }) => {
+        assert.propertyVal(headers, 'host', `127.0.0.1:${agent.port}`)
+        assert.isArray(payload)
+        // Verify we got the expected spans from the test
+        assert.strictEqual(payload.length, 2)
+        assert.strictEqual(payload[0][0].name, 'azure.servicebus.create')
+        assert.strictEqual(payload[1][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(), 'tryAddMessage-test.mjs', agent.port, spawnEnv)
+
+      await res
+    }).timeout(20000)
   })
 })
 

packages/datadog-plugin-azure-service-bus/test/integration-test/server.mjs

@@ -54,14 +54,14 @@ await sender1.scheduleMessages(amqpMessages, scheduledEnqueueTimeUtc)
 
 // queue batching
 const batch1 = await sender1.createMessageBatch()
-await batch1.tryAddMessage(messages[0])
-await batch1.tryAddMessage(messages[1])
+batch1.tryAddMessage(messages[0])
+batch1.tryAddMessage(messages[1])
 await sender1.sendMessages(batch1)
 
 // topic batching
 const batch2 = await sender2.createMessageBatch()
-await batch2.tryAddMessage(messages[0])
-await batch2.tryAddMessage(messages[1])
+batch2.tryAddMessage(messages[0])
+batch2.tryAddMessage(messages[1])
 await sender2.sendMessages(batch2)
 
 await sender1.close()

packages/datadog-plugin-azure-service-bus/test/integration-test/tryAddMessage-test.mjs

@@ -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()