From 591c40ed874c6b05b16a9bafcdf39560a7eea114 Mon Sep 17 00:00:00 2001
From: mhlidd <matthew.li@datadoghq.com>
Date: Fri, 8 Nov 2024 16:27:43 -0500
Subject: [PATCH] ddtrace/tracer: add support for span links on Span Context
 (#2973)

---
 ddtrace/ddtrace.go             | 12 ++++++++++++
 ddtrace/tracer/context_test.go | 35 ++++++++++++++++++++++++++++++++++
 ddtrace/tracer/spancontext.go  | 12 ++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/ddtrace/ddtrace.go b/ddtrace/ddtrace.go
index e311b5ff25..768146e7ac 100644
--- a/ddtrace/ddtrace.go
+++ b/ddtrace/ddtrace.go
@@ -33,6 +33,18 @@ type SpanContextW3C interface {
 	TraceID128Bytes() [16]byte
 }
 
+// SpanContextWithLinks represents a SpanContext with additional methods for
+// access to the SpanLinks on the span context, if present.
+type SpanContextWithLinks interface {
+	SpanContext
+
+	// SpanLinks returns the span links on the SpanContext.
+	SpanLinks() []SpanLink
+
+	// Setlinks takes in a slice of SpanLinks and sets them to the SpanContext.
+	SetLinks(links []SpanLink)
+}
+
 // Tracer specifies an implementation of the Datadog tracer which allows starting
 // and propagating spans. The official implementation if exposed as functions
 // within the "tracer" package.
diff --git a/ddtrace/tracer/context_test.go b/ddtrace/tracer/context_test.go
index 2ed39e9d2a..8a5fd617f0 100644
--- a/ddtrace/tracer/context_test.go
+++ b/ddtrace/tracer/context_test.go
@@ -80,6 +80,41 @@ func TestStartSpanFromContext(t *testing.T) {
 	assert.Equal("/", got.Resource)
 }
 
+func TestStartSpanWithSpanLinks(t *testing.T) {
+	_, _, _, stop := startTestTracer(t)
+	defer stop()
+	spanLink := ddtrace.SpanLink{TraceID: 789, TraceIDHigh: 0, SpanID: 789, Attributes: map[string]string{"reason": "terminated_context", "context_headers": "datadog"}, Flags: 0}
+	var ctx ddtrace.SpanContext
+	ctx = &spanContext{spanID: 789, traceID: traceIDFrom64Bits(789), spanLinks: []ddtrace.SpanLink{spanLink}}
+
+	t.Run("spanContext with spanLinks satisfies SpanContextWithLinks interface", func(t *testing.T) {
+		ctxLink, ok := ctx.(ddtrace.SpanContextWithLinks)
+		assert.True(t, ok)
+		assert.Equal(t, len(ctxLink.SpanLinks()), 1)
+		assert.Equal(t, ctxLink.SpanLinks()[0], spanLink)
+	})
+
+	t.Run("create span from spancontext with links", func(t *testing.T) {
+		var s ddtrace.Span
+		s, _ = StartSpanFromContext(
+			context.Background(),
+			"http.request",
+			WithSpanLinks([]ddtrace.SpanLink{spanLink}),
+			ChildOf(ctx),
+		)
+		//checking that a span links are added to a child span that is created where span links are passed as an StartSpanOption
+		sp, ok := s.(*span)
+		if !ok {
+			assert.Fail(t, "couldn't cast to span")
+		}
+
+		assert.Equal(t, 1, len(sp.SpanLinks))
+		assert.Equal(t, spanLink, sp.SpanLinks[0])
+
+		assert.Equal(t, 0, len(sp.context.spanLinks)) // ensure that the span links are not added to the parent context
+	})
+}
+
 func TestStartSpanFromContextRace(t *testing.T) {
 	_, _, _, stop := startTestTracer(t)
 	defer stop()
diff --git a/ddtrace/tracer/spancontext.go b/ddtrace/tracer/spancontext.go
index 29da1b80b1..44419241db 100644
--- a/ddtrace/tracer/spancontext.go
+++ b/ddtrace/tracer/spancontext.go
@@ -111,6 +111,8 @@ type spanContext struct {
 	baggage    map[string]string
 	hasBaggage uint32 // atomic int for quick checking presence of baggage. 0 indicates no baggage, otherwise baggage exists.
 	origin     string // e.g. "synthetics"
+
+	spanLinks []ddtrace.SpanLink // links to related spans in separate|external|disconnected traces
 }
 
 // newSpanContext creates a new SpanContext to serve as context for the given
@@ -166,6 +168,8 @@ func (c *spanContext) SpanID() uint64 { return c.spanID }
 // TraceID implements ddtrace.SpanContext.
 func (c *spanContext) TraceID() uint64 { return c.traceID.Lower() }
 
+func (c *spanContext) TraceIDUpper() uint64 { return c.traceID.Upper() }
+
 // TraceID128 implements ddtrace.SpanContextW3C.
 func (c *spanContext) TraceID128() string {
 	if c == nil {
@@ -179,6 +183,14 @@ func (c *spanContext) TraceID128Bytes() [16]byte {
 	return c.traceID
 }
 
+func (c *spanContext) SpanLinks() []ddtrace.SpanLink {
+	return c.spanLinks
+}
+
+func (c *spanContext) SetLinks(links []ddtrace.SpanLink) {
+	c.spanLinks = links
+}
+
 // ForeachBaggageItem implements ddtrace.SpanContext.
 func (c *spanContext) ForeachBaggageItem(handler func(k, v string) bool) {
 	if atomic.LoadUint32(&c.hasBaggage) == 0 {