fix(tracing): make peer service finalization idempotent

packages/datadog-plugin-mongodb-core/test/mongodb.spec.js

@@ -107,7 +107,8 @@ describe('Plugin', () => {
             'mongodb-core',
             (done) => collection.insertOne({ a: 1 }, {}, done),
             'test',
-            'peer.service'
+            'peer.service',
+            { component: 'mongodb' }
           )
 
           // The bulkWrite span is opened with `db.name` set, so it flows through the same
@@ -118,7 +119,11 @@ describe('Plugin', () => {
             () => collection.bulkWrite([{ insertOne: { document: { a: 1 } } }]),
             'test',
             'peer.service',
-            { desc: 'with bulkWrite' }
+            {
+              component: 'mongodb',
+              desc: 'with bulkWrite',
+              resource: () => `bulkWrite test.${collectionName}`,
+            }
           )
 
           it('should do automatic instrumentation', done => {

packages/datadog-plugin-mongoose/test/index.spec.js

@@ -73,7 +73,8 @@ describe('Plugin', () => {
           return new PeerCat({ name: 'PeerCat' }).save()
         },
         () => dbName,
-        'peer.service')
+        'peer.service',
+        { component: 'mongodb' })
 
       it('should propagate context with write operations', () => {
         const Cat = mongoose.model('Cat1', { name: String })

packages/dd-trace/src/plugins/outbound.js

@@ -125,7 +125,10 @@ class OutboundPlugin extends TracingPlugin {
    */
   tagPeerService (span) {
     if (this._tracerConfig.spanComputePeerService) {
-      const peerData = this.getPeerService(span.context().getTags())
+      const tags = span.context().getTags()
+      if (tags[PEER_SERVICE_SOURCE_KEY] !== undefined) return
+
+      const peerData = this.getPeerService(tags)
       if (peerData !== undefined) {
         span.addTags(this.getPeerServiceRemap(peerData))
       }

packages/dd-trace/test/plugins/outbound.spec.js

@@ -50,6 +50,44 @@ describe('OuboundPlugin', () => {
       sinon.assert.notCalled(getRemapStub)
     })
 
+    it('should not recompute peer service after its source was set', () => {
+      computePeerServiceStub.value({ spanComputePeerService: true })
+      const tags = {
+        'peer.service': 'mypeerservice',
+        '_dd.peer.service.source': 'out.host',
+      }
+      instance.tagPeerService({ context: () => ({ getTags: () => tags }) })
+
+      sinon.assert.notCalled(getPeerServiceStub)
+      sinon.assert.notCalled(getRemapStub)
+    })
+
+    it('should not recompute peer service when only its source was set', () => {
+      computePeerServiceStub.value({ spanComputePeerService: true })
+      const tags = {
+        '_dd.peer.service.source': 'out.host',
+      }
+      instance.tagPeerService({ context: () => ({ getTags: () => tags }) })
+
+      sinon.assert.notCalled(getPeerServiceStub)
+      sinon.assert.notCalled(getRemapStub)
+    })
+
+    it('should compute peer service when only its value was set', () => {
+      computePeerServiceStub.value({ spanComputePeerService: true })
+      getPeerServiceStub.returns({
+        'peer.service': 'mypeerservice',
+        '_dd.peer.service.source': 'peer.service',
+      })
+      const tags = {
+        'peer.service': 'mypeerservice',
+      }
+      instance.tagPeerService({ context: () => ({ getTags: () => tags }), addTags: () => {} })
+
+      sinon.assert.called(getPeerServiceStub)
+      sinon.assert.called(getRemapStub)
+    })
+
     it('should do nothing when disabled', () => {
       computePeerServiceStub.value({ spanComputePeerService: false })
       instance.tagPeerService({ context: () => ({ _tags: {}, getTags () { return this._tags } }), addTags: () => {} })

packages/dd-trace/test/setup/mocha.js

@@ -182,6 +182,14 @@ function withNamingSchema (
   })
 }
 
+/**
+ * @param {() => import('../../src/proxy')} tracer
+ * @param {string} pluginName
+ * @param {((callback: (error?: Error) => void) => unknown) | (() => Promise<unknown>)} spanGenerationFn
+ * @param {string | (() => string)} service
+ * @param {string} serviceSource
+ * @param {{ component?: string, desc?: string, resource?: string | (() => string) }} [opts]
+ */
 function withPeerService (tracer, pluginName, spanGenerationFn, service, serviceSource, opts = {}) {
   describe('peer service computation' + (opts.desc ? ` ${opts.desc}` : ''), function () {
     this.timeout(10000)
@@ -199,32 +207,72 @@ function withPeerService (tracer, pluginName, spanGenerationFn, service, service
     })
 
     it('should compute peer service', async () => {
-      const useCallback = spanGenerationFn.length === 1
-      const spanGenerationPromise = useCallback
-        ? new Promise(/** @type {() => void} */ (resolve, reject) => {
-          const result = spanGenerationFn((err) => err ? reject(err) : resolve())
-          // Some callback based methods are a mixture of callback and promise,
-          // depending on the module version. Await the promises as well.
-          if (util.types.isPromise(result)) {
-            result.then?.(resolve, reject)
+      const currentTracer = global._ddtrace
+      const parentSpan = currentTracer.startSpan('peer-service.test')
+      const traceId = BigInt(parentSpan.context().toTraceId())
+      const component = opts.component ?? pluginName
+      let traceAssertion
+
+      /**
+       * @param {Array<Array<{ trace_id: bigint, resource: string, meta: Record<string, string> }>>} traces
+       */
+      function assertPeerServiceSpan (traces) {
+        const expectedService = typeof service === 'function' ? service() : service
+        const expectedResource = typeof opts.resource === 'function' ? opts.resource() : opts.resource
+
+        for (const trace of traces) {
+          for (const span of trace) {
+            if (
+              span.trace_id === traceId &&
+              span.meta.component === component &&
+              (expectedResource === undefined || span.resource === expectedResource) &&
+              span.meta['peer.service'] === expectedService &&
+              span.meta['_dd.peer.service.source'] === serviceSource
+            ) {
+              return
+            }
           }
-        })
-        : spanGenerationFn()
+        }
 
-      assert.strictEqual(
-        typeof spanGenerationPromise?.then, 'function',
-        'spanGenerationFn should return a promise in case no callback is defined. Received: ' +
-          util.inspect(spanGenerationPromise, { depth: 1 }),
-      )
+        assert.fail(
+          `No ${component} span in trace ${traceId} had peer.service=${expectedService} and ` +
+          `_dd.peer.service.source=${serviceSource}.\n\nCandidate Traces:\n${util.inspect(traces, { depth: null })}`
+        )
+      }
 
-      await Promise.all([
-        getAgent().assertSomeTraces(traces => {
-          const span = traces[0][0]
-          assert.strictEqual(span.meta['peer.service'], typeof service === 'function' ? service() : service)
-          assert.strictEqual(span.meta['_dd.peer.service.source'], serviceSource)
-        }),
-        spanGenerationPromise,
-      ])
+      try {
+        traceAssertion = getAgent().assertSomeTraces(assertPeerServiceSpan, { timeoutMs: 9000 })
+        const useCallback = spanGenerationFn.length === 1
+        const spanGenerationPromise = currentTracer.scope().activate(parentSpan, () => {
+          return useCallback
+            ? new Promise(/** @type {() => void} */ (resolve, reject) => {
+              const result = spanGenerationFn((error) => error ? reject(error) : resolve())
+              // Some callback based methods are a mixture of callback and promise,
+              // depending on the module version. Await the promises as well.
+              if (util.types.isPromise(result)) {
+                result.then?.(resolve, reject)
+              }
+            })
+            : spanGenerationFn()
+        })
+
+        assert.strictEqual(
+          typeof spanGenerationPromise?.then, 'function',
+          'spanGenerationFn should return a promise in case no callback is defined. Received: ' +
+            util.inspect(spanGenerationPromise, { depth: 1 }),
+        )
+
+        await Promise.all([
+          traceAssertion,
+          (async () => {
+            await spanGenerationPromise
+            parentSpan.finish()
+          })(),
+        ])
+      } finally {
+        parentSpan.finish()
+        traceAssertion?.cancel()
+      }
     })
   })
 }