From 4ee30f7e3d264d428f653688ee4ab1dffae4e484 Mon Sep 17 00:00:00 2001
From: Thomas Watson <w@tson.dk>
Date: Fri, 11 Oct 2024 13:11:18 +0200
Subject: [PATCH] Add support for exit spans in Code Origin for Spans

---
 packages/dd-trace/src/plugins/outbound.js     |  9 +++
 .../dd-trace/test/plugins/outbound.spec.js    | 71 +++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/packages/dd-trace/src/plugins/outbound.js b/packages/dd-trace/src/plugins/outbound.js
index f0a9509269e..728fe269b84 100644
--- a/packages/dd-trace/src/plugins/outbound.js
+++ b/packages/dd-trace/src/plugins/outbound.js
@@ -7,6 +7,7 @@ const {
   PEER_SERVICE_REMAP_KEY
 } = require('../constants')
 const TracingPlugin = require('./tracing')
+const { exitTag } = require('../../../datadog-plugin-code-origin/src/lib/tags')
 
 const COMMON_PEER_SVC_SOURCE_TAGS = [
   'net.peer.name',
@@ -25,6 +26,14 @@ class OutboundPlugin extends TracingPlugin {
     })
   }
 
+  startSpan (...args) {
+    const span = super.startSpan(...args)
+    if (this._tracerConfig.codeOriginForSpansEnabled) {
+      span.addTags(exitTag(this.startSpan))
+    }
+    return span
+  }
+
   getPeerService (tags) {
     /**
      * Compute `peer.service` and associated metadata from available tags, based
diff --git a/packages/dd-trace/test/plugins/outbound.spec.js b/packages/dd-trace/test/plugins/outbound.spec.js
index 5709c789575..e49e35aad17 100644
--- a/packages/dd-trace/test/plugins/outbound.spec.js
+++ b/packages/dd-trace/test/plugins/outbound.spec.js
@@ -157,4 +157,75 @@ describe('OuboundPlugin', () => {
       })
     })
   })
+
+  describe('code origin tags', () => {
+    let instance = null
+
+    beforeEach(() => {
+      const tracerStub = {
+        _tracer: {
+          startSpan: sinon.stub().returns({
+            addTags: sinon.spy()
+          })
+        }
+      }
+      instance = new OutboundPlugin(tracerStub)
+    })
+
+    it('should not add exit tags to span if codeOriginForSpansEnabled is false', () => {
+      sinon.stub(instance, '_tracerConfig').value({ codeOriginForSpansEnabled: false })
+      const span = instance.startSpan('test')
+      expect(span.addTags).to.not.have.been.called
+    })
+
+    it('should add exit tags to span if codeOriginForSpansEnabled is true', () => {
+      sinon.stub(instance, '_tracerConfig').value({ codeOriginForSpansEnabled: true })
+
+      const lineNumber = getNextLineNumber()
+      const span = instance.startSpan('test')
+
+      expect(span.addTags).to.have.been.calledOnce
+      const args = span.addTags.args[0]
+      expect(args).to.have.property('length', 1)
+      const tags = parseTags(args[0])
+
+      expect(tags).to.nested.include({ '_dd.code_origin.type': 'exit' })
+      expect(tags._dd.code_origin).to.have.property('frames').to.be.an('array').with.length.above(0)
+
+      for (const frame of tags._dd.code_origin.frames) {
+        expect(frame).to.have.property('file', __filename)
+        expect(frame).to.have.property('line').to.match(/^\d+$/)
+        expect(frame).to.have.property('column').to.match(/^\d+$/)
+        expect(frame).to.have.property('type').to.a('string')
+      }
+
+      const topFrame = tags._dd.code_origin.frames[0]
+      expect(topFrame).to.have.property('line', lineNumber)
+    })
+  })
 })
+
+function getNextLineNumber () {
+  return String(Number(new Error().stack.split('\n')[2].match(/:(\d+):/)[1]) + 1)
+}
+
+function parseTags (tags) {
+  const parsedTags = {}
+  for (const [tag, value] of Object.entries(tags)) {
+    const keys = tag.split('.')
+    let current = parsedTags
+    let depth = 0
+    for (const key of keys) {
+      if (!current[key]) {
+        if (depth === keys.length - 1) {
+          current[key] = value
+          break
+        }
+        current[key] = keys[depth + 1]?.match(/^\d+$/) ? [] : {}
+      }
+      current = current[key]
+      depth++
+    }
+  }
+  return parsedTags
+}