Add support for SpanLinks (#2053)

bidetofevil · web-flow · commit 9825312e61c6 · 2025-03-21T07:38:22.000-07:00
## Goal

Add basic support for SpanLinks. A couple of things still have to be finalized:

- Payload shape
- Limits

Thought I'd get this in review first and made the approach changes after

## Testing
Unit and integration tests added
diff --git a/embrace-android-api/api/embrace-android-api.api b/embrace-android-api/api/embrace-android-api.api
@@ -263,6 +263,10 @@ public abstract interface class io/embrace/android/embracesdk/spans/EmbraceSpan
 	public abstract fun addEvent (Ljava/lang/String;)Z
 	public abstract fun addEvent (Ljava/lang/String;Ljava/lang/Long;)Z
 	public abstract fun addEvent (Ljava/lang/String;Ljava/lang/Long;Ljava/util/Map;)Z
+	public abstract fun addLink (Lio/embrace/android/embracesdk/spans/EmbraceSpan;)Z
+	public abstract fun addLink (Lio/embrace/android/embracesdk/spans/EmbraceSpan;Ljava/util/Map;)Z
+	public abstract fun addLink (Lio/opentelemetry/api/trace/SpanContext;)Z
+	public abstract fun addLink (Lio/opentelemetry/api/trace/SpanContext;Ljava/util/Map;)Z
 	public abstract fun getAutoTerminationMode ()Lio/embrace/android/embracesdk/spans/AutoTerminationMode;
 	public abstract fun getParent ()Lio/embrace/android/embracesdk/spans/EmbraceSpan;
 	public abstract fun getSpanContext ()Lio/opentelemetry/api/trace/SpanContext;
@@ -283,6 +287,9 @@ public abstract interface class io/embrace/android/embracesdk/spans/EmbraceSpan
 public final class io/embrace/android/embracesdk/spans/EmbraceSpan$DefaultImpls {
 	public static fun addEvent (Lio/embrace/android/embracesdk/spans/EmbraceSpan;Ljava/lang/String;)Z
 	public static fun addEvent (Lio/embrace/android/embracesdk/spans/EmbraceSpan;Ljava/lang/String;Ljava/lang/Long;)Z
+	public static fun addLink (Lio/embrace/android/embracesdk/spans/EmbraceSpan;Lio/embrace/android/embracesdk/spans/EmbraceSpan;)Z
+	public static fun addLink (Lio/embrace/android/embracesdk/spans/EmbraceSpan;Lio/embrace/android/embracesdk/spans/EmbraceSpan;Ljava/util/Map;)Z
+	public static fun addLink (Lio/embrace/android/embracesdk/spans/EmbraceSpan;Lio/opentelemetry/api/trace/SpanContext;)Z
 	public static fun recordException (Lio/embrace/android/embracesdk/spans/EmbraceSpan;Ljava/lang/Throwable;)Z
 	public static fun start (Lio/embrace/android/embracesdk/spans/EmbraceSpan;)Z
 	public static fun stop (Lio/embrace/android/embracesdk/spans/EmbraceSpan;)Z
diff --git a/embrace-android-api/src/main/kotlin/io/embrace/android/embracesdk/spans/EmbraceSpan.kt b/embrace-android-api/src/main/kotlin/io/embrace/android/embracesdk/spans/EmbraceSpan.kt
@@ -131,4 +131,33 @@ public interface EmbraceSpan {
      * Update the name of the span. Returns false if the update was not successful, like when it has already been stopped
      */
     public fun updateName(newName: String): Boolean
+
+    /**
+     * Add a link to the given [EmbraceSpan]
+     */
+    public fun addLink(linkedSpan: EmbraceSpan): Boolean = addLink(linkedSpan = linkedSpan, attributes = null)
+
+    /**
+     * Add a link to the span with the given [SpanContext]
+     */
+    public fun addLink(
+        linkedSpanContext: SpanContext
+    ): Boolean = addLink(linkedSpanContext = linkedSpanContext, attributes = null)
+
+    /**
+     * Add a link to the given [EmbraceSpan] with the given attributes
+     */
+    public fun addLink(linkedSpan: EmbraceSpan, attributes: Map<String, String>?): Boolean {
+        val spanContext = linkedSpan.spanContext
+        return if (spanContext != null) {
+            addLink(linkedSpanContext = spanContext, attributes = attributes)
+        } else {
+            false
+        }
+    }
+
+    /**
+     * Add a link to the span with the given [SpanContext] with the given attributes
+     */
+    public fun addLink(linkedSpanContext: SpanContext, attributes: Map<String, String>?): Boolean
 }
diff --git a/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/gating/SpanSanitizer.kt b/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/gating/SpanSanitizer.kt
@@ -34,7 +34,8 @@ internal class SpanSanitizer(
             sessionSpan.endTimeNanos,
             sessionSpan.status,
             sanitizedEvents,
-            sessionSpan.attributes
+            sessionSpan.attributes,
+            sessionSpan.links,
         )
         sanitizedSpans.add(sanitizedSessionSpan)
         return sanitizedSpans
diff --git a/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/opentelemetry/EmbSpan.kt b/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/opentelemetry/EmbSpan.kt
@@ -38,6 +38,11 @@ class EmbSpan(
         return this
     }
 
+    override fun addLink(spanContext: SpanContext, attributes: Attributes): Span {
+        embraceSpan.addLink(spanContext, attributes.toStringMap())
+        return this
+    }
+
     override fun setStatus(statusCode: StatusCode, description: String): Span {
         if (isRecording) {
             embraceSpan.setStatus(statusCode, description)
diff --git a/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/payload/SpanMapper.kt b/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/payload/SpanMapper.kt
@@ -5,6 +5,7 @@ import io.embrace.android.embracesdk.internal.arch.schema.EmbType
 import io.embrace.android.embracesdk.internal.arch.schema.ErrorCodeAttribute
 import io.embrace.android.embracesdk.internal.clock.millisToNanos
 import io.embrace.android.embracesdk.internal.clock.nanosToMillis
+import io.embrace.android.embracesdk.internal.spans.EmbraceLinkData
 import io.embrace.android.embracesdk.internal.spans.EmbraceSpanData
 import io.embrace.android.embracesdk.internal.spans.hasFixedAttribute
 import io.embrace.android.embracesdk.internal.spans.setFixedAttribute
@@ -22,7 +23,8 @@ fun EmbraceSpanData.toNewPayload(): Span = Span(
     endTimeNanos = endTimeNanos,
     status = status.toStatus(),
     events = events.map(EmbraceSpanEvent::toNewPayload),
-    attributes = attributes.toNewPayload()
+    attributes = attributes.toNewPayload(),
+    links = links,
 )
 
 fun EmbraceSpanEvent.toNewPayload(): SpanEvent = SpanEvent(
@@ -43,6 +45,8 @@ fun Map<String, String>.toNewPayload(): List<Attribute> =
 fun List<Attribute>.toOldPayload(): Map<String, String> =
     associate { Pair(it.key ?: "", it.data ?: "") }.filterKeys { it.isNotBlank() }
 
+fun EmbraceLinkData.toPayload() = Link(spanContext.spanId, attributes?.toNewPayload())
+
 fun Span.toOldPayload(): EmbraceSpanData {
     return EmbraceSpanData(
         traceId = traceId ?: "",
@@ -58,7 +62,8 @@ fun Span.toOldPayload(): EmbraceSpanData {
             else -> StatusCode.UNSET
         },
         events = events?.mapNotNull { it.toOldPayload() } ?: emptyList(),
-        attributes = attributes?.toOldPayload() ?: emptyMap()
+        attributes = attributes?.toOldPayload() ?: emptyMap(),
+        links = links ?: emptyList()
     )
 }
 
diff --git a/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/EmbraceLinkData.kt b/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/EmbraceLinkData.kt
@@ -0,0 +1,8 @@
+package io.embrace.android.embracesdk.internal.spans
+
+import io.opentelemetry.api.trace.SpanContext
+
+data class EmbraceLinkData(
+    val spanContext: SpanContext,
+    val attributes: Map<String, String>?
+)
diff --git a/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/EmbraceSpanData.kt b/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/EmbraceSpanData.kt
@@ -1,6 +1,7 @@
 package io.embrace.android.embracesdk.internal.spans
 
 import io.embrace.android.embracesdk.internal.clock.nanosToMillis
+import io.embrace.android.embracesdk.internal.payload.Link
 import io.embrace.android.embracesdk.spans.EmbraceSpanEvent
 import io.opentelemetry.api.trace.StatusCode
 import io.opentelemetry.sdk.trace.data.EventData
@@ -26,6 +27,8 @@ data class EmbraceSpanData(
     val events: List<EmbraceSpanEvent> = emptyList(),
 
     val attributes: Map<String, String> = emptyMap(),
+
+    val links: List<Link> = emptyList(),
 ) {
 
     companion object {
diff --git a/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/EmbraceSpanImpl.kt b/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/EmbraceSpanImpl.kt
@@ -16,6 +16,7 @@ import io.embrace.android.embracesdk.internal.config.instrumented.schema.OtelLim
 import io.embrace.android.embracesdk.internal.payload.Attribute
 import io.embrace.android.embracesdk.internal.payload.Span
 import io.embrace.android.embracesdk.internal.payload.toNewPayload
+import io.embrace.android.embracesdk.internal.payload.toPayload
 import io.embrace.android.embracesdk.internal.utils.truncatedStacktraceText
 import io.embrace.android.embracesdk.spans.AutoTerminationMode
 import io.embrace.android.embracesdk.spans.EmbraceSpan
@@ -64,6 +65,7 @@ internal class EmbraceSpanImpl(
     private val customAttributes = ConcurrentHashMap<String, String>().apply {
         putAll(spanBuilder.getCustomAttributes())
     }
+    private val systemLinks = ConcurrentLinkedQueue<EmbraceLinkData>()
 
     // size for ConcurrentLinkedQueues is not a constant operation, so it could be subject to race conditions
     // do the bookkeeping separately so we don't have to worry about this
@@ -127,29 +129,9 @@ internal class EmbraceSpanImpl(
             }
 
             startedSpan.get()?.let { spanToStop ->
-                systemAttributes.forEach { systemAttribute ->
-                    spanToStop.setAttribute(systemAttribute.key, systemAttribute.value)
-                }
-                customAttributes.redactIfSensitive().forEach { attribute ->
-                    spanToStop.setAttribute(attribute.key, attribute.value)
-                }
-
-                val redactedCustomEvents = customEvents.map { it.copy(attributes = it.attributes.redactIfSensitive()) }
-
-                (systemEvents + redactedCustomEvents).forEach { event ->
-                    val eventAttributes = if (event.attributes.isNotEmpty()) {
-                        Attributes.builder().fromMap(event.attributes, spanBuilder.internal).build()
-                    } else {
-                        Attributes.empty()
-                    }
-
-                    spanToStop.addEvent(
-                        event.name,
-                        eventAttributes,
-                        event.timestampNanos,
-                        TimeUnit.NANOSECONDS
-                    )
-                }
+                populateAttributes(spanToStop)
+                populateEvents(spanToStop)
+                populateLinks(spanToStop)
 
                 if (errorCode != null) {
                     setStatus(StatusCode.ERROR)
@@ -171,6 +153,17 @@ internal class EmbraceSpanImpl(
         return successful
     }
 
+    private fun populateLinks(spanToStop: io.opentelemetry.api.trace.Span) {
+        systemLinks.forEach {
+            val linkAttributes = if (it.attributes != null) {
+                Attributes.builder().fromMap(attributes = it.attributes, false).build()
+            } else {
+                Attributes.empty()
+            }
+            spanToStop.addLink(it.spanContext, linkAttributes)
+        }
+    }
+
     override fun addEvent(name: String, timestampMs: Long?, attributes: Map<String, String>?): Boolean =
         recordEvent(customEvents, customEventCount, limits.getMaxCustomEventCount()) {
             EmbraceSpanEvent.create(
@@ -268,6 +261,20 @@ internal class EmbraceSpanImpl(
         return false
     }
 
+    override fun addLink(linkedSpanContext: SpanContext, attributes: Map<String, String>?): Boolean {
+        if (systemLinks.size < limits.getMaxTotalLinkCount()) {
+            synchronized(systemLinks) {
+                if (systemLinks.size < limits.getMaxTotalLinkCount()) {
+                    systemLinks.add(EmbraceLinkData(linkedSpanContext, attributes))
+                    spanRepository.notifySpanUpdate()
+                    return true
+                }
+            }
+        }
+
+        return false
+    }
+
     override fun asNewContext(): Context? = startedSpan.get()?.run { spanBuilder.parentContext.with(this) }
 
     override fun snapshot(): Span? {
@@ -282,7 +289,8 @@ internal class EmbraceSpanImpl(
                 endTimeNanos = spanEndTimeMs?.millisToNanos(),
                 status = status,
                 events = systemEvents.map(EmbraceSpanEvent::toNewPayload) + redactedCustomEvents.map(EmbraceSpanEvent::toNewPayload),
-                attributes = getAttributesPayload()
+                attributes = getAttributesPayload(),
+                links = systemLinks.toList().map { it.toPayload() }
             )
         } else {
             null
@@ -348,4 +356,32 @@ internal class EmbraceSpanImpl(
             }
         }
     }
+
+    private fun populateAttributes(spanToStop: io.opentelemetry.api.trace.Span) {
+        systemAttributes.forEach { systemAttribute ->
+            spanToStop.setAttribute(systemAttribute.key, systemAttribute.value)
+        }
+        customAttributes.redactIfSensitive().forEach { attribute ->
+            spanToStop.setAttribute(attribute.key, attribute.value)
+        }
+    }
+
+    private fun populateEvents(spanToStop: io.opentelemetry.api.trace.Span) {
+        val redactedCustomEvents = customEvents.map { it.copy(attributes = it.attributes.redactIfSensitive()) }
+
+        (systemEvents + redactedCustomEvents).forEach { event ->
+            val eventAttributes = if (event.attributes.isNotEmpty()) {
+                Attributes.builder().fromMap(event.attributes, spanBuilder.internal).build()
+            } else {
+                Attributes.empty()
+            }
+
+            spanToStop.addEvent(
+                event.name,
+                eventAttributes,
+                event.timestampNanos,
+                TimeUnit.NANOSECONDS
+            )
+        }
+    }
 }
diff --git a/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/SpanDataExt.kt b/embrace-android-core/src/main/kotlin/io/embrace/android/embracesdk/internal/spans/SpanDataExt.kt
@@ -1,5 +1,7 @@
 package io.embrace.android.embracesdk.internal.spans
 
+import io.embrace.android.embracesdk.internal.payload.Link
+import io.embrace.android.embracesdk.internal.payload.toNewPayload
 import io.embrace.android.embracesdk.internal.spans.EmbraceSpanData.Companion.fromEventData
 import io.opentelemetry.sdk.trace.data.SpanData
 
@@ -13,4 +15,5 @@ fun SpanData.toEmbraceSpanData(): EmbraceSpanData = EmbraceSpanData(
     status = status.statusCode,
     events = fromEventData(eventDataList = events),
     attributes = attributes.toStringMap(),
+    links = links.map { Link(it.spanContext.spanId, it.attributes.toNewPayload()) }
 )
diff --git a/embrace-android-core/src/test/java/io/embrace/android/embracesdk/internal/opentelemetry/EmbSpanTest.kt b/embrace-android-core/src/test/java/io/embrace/android/embracesdk/internal/opentelemetry/EmbSpanTest.kt
@@ -176,4 +176,18 @@ internal class EmbSpanTest {
 
         assertEquals(attributesCount + 10, fakeEmbraceSpan.attributes.size)
     }
+
+    @Test
+    fun `add span link`() {
+        with(embSpan) {
+            val linkedSpanContext = checkNotNull(FakePersistableEmbraceSpan.started().spanContext)
+            addLink(linkedSpanContext, Attributes.builder().put("boolean", true).build())
+            with(fakeEmbraceSpan.links.single()) {
+                assertEquals(linkedSpanContext.spanId, spanId)
+                val attribute = checkNotNull(attributes?.single())
+                assertEquals("boolean", attribute.key)
+                assertEquals("true", attribute.data)
+            }
+        }
+    }
 }
diff --git a/embrace-android-payload/src/main/kotlin/io/embrace/android/embracesdk/internal/config/instrumented/schema/OtelLimitsConfig.kt b/embrace-android-payload/src/main/kotlin/io/embrace/android/embracesdk/internal/config/instrumented/schema/OtelLimitsConfig.kt
@@ -13,6 +13,7 @@ interface OtelLimitsConfig {
     fun getMaxTotalEventCount(): Int = 11000
     fun getMaxCustomAttributeCount(): Int = 50
     fun getMaxTotalAttributeCount(): Int = 300
+    fun getMaxTotalLinkCount(): Int = 100
     fun getMaxInternalAttributeKeyLength(): Int = 1000
     fun getMaxInternalAttributeValueLength(): Int = 2000
     fun getMaxCustomAttributeKeyLength(): Int = 50
diff --git a/embrace-android-payload/src/main/kotlin/io/embrace/android/embracesdk/internal/payload/Link.kt b/embrace-android-payload/src/main/kotlin/io/embrace/android/embracesdk/internal/payload/Link.kt
@@ -0,0 +1,15 @@
+package io.embrace.android.embracesdk.internal.payload
+
+import com.squareup.moshi.Json
+import com.squareup.moshi.JsonClass
+
+@JsonClass(generateAdapter = true)
+data class Link(
+
+    /* The ID of the span that this link is to */
+    @Json(name = "span_id")
+    val spanId: String? = null,
+
+    @Json(name = "attributes")
+    val attributes: List<Attribute>? = null,
+)
diff --git a/embrace-android-payload/src/main/kotlin/io/embrace/android/embracesdk/internal/payload/Span.kt b/embrace-android-payload/src/main/kotlin/io/embrace/android/embracesdk/internal/payload/Span.kt
@@ -53,6 +53,9 @@ data class Span(
 
     @Json(name = "attributes")
     val attributes: List<Attribute>? = null,
+
+    @Json(name = "links")
+    val links: List<Link>? = null,
 ) {
 
     /**
diff --git a/embrace-android-sdk/src/integrationTest/kotlin/io/embrace/android/embracesdk/testcases/TracingApiTest.kt b/embrace-android-sdk/src/integrationTest/kotlin/io/embrace/android/embracesdk/testcases/TracingApiTest.kt
@@ -3,6 +3,7 @@ package io.embrace.android.embracesdk.testcases
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import io.embrace.android.embracesdk.arch.assertIsTypePerformance
 import io.embrace.android.embracesdk.assertions.assertEmbraceSpanData
+import io.embrace.android.embracesdk.assertions.findSpanByName
 import io.embrace.android.embracesdk.concurrency.SingleThreadTestScheduledExecutor
 import io.embrace.android.embracesdk.fakes.FakeSpanExporter
 import io.embrace.android.embracesdk.fakes.config.FakeEnabledFeatureConfig
@@ -329,6 +330,43 @@ internal class TracingApiTest {
         )
     }
 
+    @Test
+    fun `span links`() {
+        var linkedSpanId: String? = null
+        testRule.runTest(
+            testCaseAction = {
+                recordSession {
+                    clock.tick(100L)
+                    val op = checkNotNull(embrace.startSpan("my-op"))
+                    clock.tick(100L)
+                    val op2 = checkNotNull(embrace.startSpan("my-op-2"))
+                    clock.tick(100L)
+                    op2.addLink(op, mapOf("test" to "value"))
+                    op2.stop()
+                    op.stop()
+                    linkedSpanId = op.spanId
+                }
+            },
+            assertAction = {
+                with(getSingleSessionEnvelope()) {
+                    val op = findSpanByName(name = "my-op")
+                    val op2 = findSpanByName(name = "my-op-2")
+                    assertNotNull(op)
+                    assertNotNull(op2)
+                    assertEquals(op.spanId, op2.links?.single()?.spanId)
+                }
+            },
+            otelExportAssertion = {
+                val spanWithLink = awaitSpans(1) { it.links.size == 1 }.single()
+                with(spanWithLink.links.single()) {
+                    assertEquals(linkedSpanId, spanContext.spanId)
+                    assertEquals("value", attributes.toNewPayload().single().data)
+                }
+            }
+        )
+    }
+
+
     private fun EmbracePayloadAssertionInterface.getSdkInitSpanFromBackgroundActivity(): List<Span> {
         val lastSentBackgroundActivity = getSingleSessionEnvelope(ApplicationState.BACKGROUND)
         val spans = checkNotNull(lastSentBackgroundActivity.data.spans)
diff --git a/embrace-test-fakes/src/main/kotlin/io/embrace/android/embracesdk/fakes/FakePersistableEmbraceSpan.kt b/embrace-test-fakes/src/main/kotlin/io/embrace/android/embracesdk/fakes/FakePersistableEmbraceSpan.kt
