Support creating custom spans and attributes in the startup trace (#1824

)

## Goal

Add the ability to add custom child spans and attributes to UI load traces. Also exposed a public method to get the SDK timestamp so the custom spans can use the same clock as the rest of the spans i the traces

## Testing
Added unit and integration tests

embrace-android-api/api/embrace-android-api.api

@@ -56,10 +56,15 @@ public abstract interface class io/embrace/android/embracesdk/internal/api/Instr
 	public abstract fun addAttributeToLoadTrace (Landroid/app/Activity;Ljava/lang/String;Ljava/lang/String;)V
 	public abstract fun addChildSpanToLoadTrace (Landroid/app/Activity;Ljava/lang/String;JJ)V
 	public abstract fun addChildSpanToLoadTrace (Landroid/app/Activity;Ljava/lang/String;JJLjava/util/Map;Ljava/util/List;Lio/embrace/android/embracesdk/spans/ErrorCode;)V
+	public abstract fun addStartupChildSpan (Ljava/lang/String;JJ)V
+	public abstract fun addStartupChildSpan (Ljava/lang/String;JJLjava/util/Map;Ljava/util/List;Lio/embrace/android/embracesdk/spans/ErrorCode;)V
+	public abstract fun addStartupTraceAttribute (Ljava/lang/String;Ljava/lang/String;)V
+	public abstract fun getSdkCurrentTimeMs ()J
 }
 
 public final class io/embrace/android/embracesdk/internal/api/InstrumentationApi$DefaultImpls {
 	public static fun addChildSpanToLoadTrace (Lio/embrace/android/embracesdk/internal/api/InstrumentationApi;Landroid/app/Activity;Ljava/lang/String;JJ)V
+	public static fun addStartupChildSpan (Lio/embrace/android/embracesdk/internal/api/InstrumentationApi;Ljava/lang/String;JJ)V
 }
 
 public abstract interface class io/embrace/android/embracesdk/internal/api/InternalWebViewApi {
@@ -101,6 +106,7 @@ public abstract interface class io/embrace/android/embracesdk/internal/api/SdkAp
 
 public final class io/embrace/android/embracesdk/internal/api/SdkApi$DefaultImpls {
 	public static fun addChildSpanToLoadTrace (Lio/embrace/android/embracesdk/internal/api/SdkApi;Landroid/app/Activity;Ljava/lang/String;JJ)V
+	public static fun addStartupChildSpan (Lio/embrace/android/embracesdk/internal/api/SdkApi;Ljava/lang/String;JJ)V
 	public static fun createSpan (Lio/embrace/android/embracesdk/internal/api/SdkApi;Ljava/lang/String;Lio/embrace/android/embracesdk/spans/AutoTerminationMode;)Lio/embrace/android/embracesdk/spans/EmbraceSpan;
 	public static fun recordCompletedSpan (Lio/embrace/android/embracesdk/internal/api/SdkApi;Ljava/lang/String;JJ)Z
 	public static fun recordCompletedSpan (Lio/embrace/android/embracesdk/internal/api/SdkApi;Ljava/lang/String;JJLio/embrace/android/embracesdk/spans/EmbraceSpan;)Z

embrace-android-api/src/main/kotlin/io/embrace/android/embracesdk/internal/api/InstrumentationApi.kt

@@ -17,6 +17,12 @@ public interface InstrumentationApi {
      */
     public fun activityLoaded(activity: Activity)
 
+    /**
+     * Return the timestamp of the SDK clock. This value should be used when creating spans with specific start and
+     * end times so that
+     */
+    public fun getSdkCurrentTimeMs(): Long
+
     /**
      * Add an attribute to the trace generated by the loading of the given [Activity]
      */
@@ -53,4 +59,38 @@ public interface InstrumentationApi {
         events: List<EmbraceSpanEvent>,
         errorCode: ErrorCode?,
     )
+
+    /**
+     * Add an attribute to the app startup trace
+     */
+    public fun addStartupTraceAttribute(key: String, value: String)
+
+    /**
+     * Add a successfully completed child span to the app startup trace
+     */
+    public fun addStartupChildSpan(
+        name: String,
+        startTimeMs: Long,
+        endTimeMs: Long,
+    ): Unit = addStartupChildSpan(
+        name = name,
+        startTimeMs = startTimeMs,
+        endTimeMs = endTimeMs,
+        attributes = emptyMap(),
+        events = emptyList(),
+        errorCode = null,
+    )
+
+    /**
+     * Add a completed child span to the app startup trace with the given  attributes and span events.
+     * Specify an [ErrorCode] if the span didn't complete successfully.
+     */
+    public fun addStartupChildSpan(
+        name: String,
+        startTimeMs: Long,
+        endTimeMs: Long,
+        attributes: Map<String, String>,
+        events: List<EmbraceSpanEvent>,
+        errorCode: ErrorCode?,
+    )
 }

embrace-android-features/src/main/kotlin/io/embrace/android/embracesdk/internal/capture/startup/AppStartupDataCollector.kt

@@ -59,4 +59,9 @@ interface AppStartupDataCollector {
      * Set an arbitrary time interval during startup that is of note
      */
     fun addTrackedInterval(name: String, startTimeMs: Long, endTimeMs: Long)
+
+    /**
+     * Add custom attribute to the root span of the trace logged for app startup
+     */
+    fun addAttribute(key: String, value: String)
 }

embrace-android-features/src/main/kotlin/io/embrace/android/embracesdk/internal/capture/startup/AppStartupTraceEmitter.kt

@@ -12,6 +12,7 @@ import io.embrace.android.embracesdk.internal.utils.VersionChecker
 import io.embrace.android.embracesdk.internal.worker.BackgroundWorker
 import io.embrace.android.embracesdk.spans.EmbraceSpan
 import io.opentelemetry.sdk.common.Clock
+import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ConcurrentLinkedQueue
 import java.util.concurrent.atomic.AtomicBoolean
 
@@ -46,6 +47,7 @@ internal class AppStartupTraceEmitter(
     private val processCreateRequestedMs: Long?
     private val processCreatedMs: Long?
     private val additionalTrackedIntervals = ConcurrentLinkedQueue<TrackedInterval>()
+    private val customAttributes: MutableMap<String, String> = ConcurrentHashMap()
 
     init {
         val timestampAtDeviceStart = nowMs() - clock.nanoTime().nanosToMillis()
@@ -159,6 +161,10 @@ internal class AppStartupTraceEmitter(
         )
     }
 
+    override fun addAttribute(key: String, value: String) {
+        customAttributes[key] = value
+    }
+
     /**
      * Called when app startup is considered complete, i.e. the data can be used and any additional updates can be ignored
      */
@@ -250,6 +256,7 @@ internal class AppStartupTraceEmitter(
                     parent = startupTrace,
                     startTimeMs = trackedInterval.startTimeMs,
                     endTimeMs = trackedInterval.endTimeMs,
+                    internal = false,
                 )
             }
         } while (additionalTrackedIntervals.isNotEmpty())
@@ -388,6 +395,7 @@ internal class AppStartupTraceEmitter(
     private fun nowMs(): Long = clock.now().nanosToMillis()
 
     private fun PersistableEmbraceSpan.addTraceMetadata() {
+        addCustomAttributes()
         processCreateDelay()?.let { delay ->
             addAttribute("process-create-delay-ms", delay.toString())
         }
@@ -417,6 +425,12 @@ internal class AppStartupTraceEmitter(
         }
     }
 
+    private fun PersistableEmbraceSpan.addCustomAttributes() {
+        customAttributes.forEach {
+            addAttribute(it.key, it.value)
+        }
+    }
+
     private data class TrackedInterval(
         val name: String,
         val startTimeMs: Long,

embrace-android-features/src/test/java/io/embrace/android/embracesdk/internal/capture/startup/AppStartupTraceEmitterTest.kt

@@ -223,23 +223,27 @@ internal class AppStartupTraceEmitterTest {
     private fun verifyColdStartWithRender(processCreateDelayMs: Long? = null) {
         clock.tick(100L)
         appStartupTraceEmitter.applicationInitStart()
+        val customSpanStartMs = clock.now()
         clock.tick(15L)
         val (sdkInitStart, sdkInitEnd) = startSdk()
+        appStartupTraceEmitter.addAttribute("custom-attribute", "true")
         appStartupTraceEmitter.applicationInitEnd()
         val applicationInitEnd = clock.now()
         clock.tick(50L)
+        appStartupTraceEmitter.addTrackedInterval("custom-span", customSpanStartMs, applicationInitEnd)
 
         val activityCreateEvents = createStartupActivity()
         val traceEnd = startupActivityRender().second
 
-        assertEquals(7, spanSink.completedSpans().size)
+        assertEquals(8, spanSink.completedSpans().size)
         val spanMap = spanSink.completedSpans().associateBy { it.name }
         val trace = checkNotNull(spanMap["emb-cold-time-to-initial-display"])
         val processInit = checkNotNull(spanMap["emb-process-init"])
         val embraceInit = checkNotNull(spanMap["emb-embrace-init"])
         val activityInitDelay = checkNotNull(spanMap["emb-activity-init-gap"])
         val activityCreate = checkNotNull(spanMap["emb-activity-create"])
         val firstRender = checkNotNull(spanMap["emb-first-frame-render"])
+        val customSpan = checkNotNull(spanMap["custom-span"])
 
         val startupActivityStart = checkNotNull(activityCreateEvents.create)
         val startupActivityEnd = checkNotNull((activityCreateEvents.finished))
@@ -252,12 +256,14 @@ internal class AppStartupTraceEmitterTest {
             expectedActivityPreCreatedMs = activityCreateEvents.preCreate,
             expectedActivityPostCreatedMs = activityCreateEvents.postCreate,
             expectedFirstActivityLifecycleEventMs = activityCreateEvents.firstEvent,
+            expectedCustomAttributes = mapOf("custom-attribute" to "true")
         )
         assertChildSpan(processInit, DEFAULT_FAKE_CURRENT_TIME, applicationInitEnd)
         assertChildSpan(embraceInit, sdkInitStart, sdkInitEnd)
         assertChildSpan(activityInitDelay, applicationInitEnd, startupActivityStart)
         assertChildSpan(activityCreate, startupActivityStart, startupActivityEnd)
         assertChildSpan(firstRender, startupActivityEnd, traceEnd)
+        assertChildSpan(customSpan, customSpanStartMs, applicationInitEnd)
         assertEquals(0, logger.internalErrorMessages.size)
     }
 
@@ -542,6 +548,7 @@ internal class AppStartupTraceEmitterTest {
         expectedActivityPreCreatedMs: Long? = null,
         expectedActivityPostCreatedMs: Long? = null,
         expectedFirstActivityLifecycleEventMs: Long? = null,
+        expectedCustomAttributes: Map<String, String> = emptyMap(),
     ) {
         val trace = input.toNewPayload()
         assertEquals(expectedStartTimeMs, trace.startTimeNanos?.nanosToMillis())
@@ -565,6 +572,10 @@ internal class AppStartupTraceEmitterTest {
         assertEquals("false", attrs.findAttributeValue("embrace-init-in-foreground"))
         assertEquals("main", attrs.findAttributeValue("embrace-init-thread-name"))
         assertEquals(1, dataCollectionCompletedCallbackInvokedCount)
+
+        expectedCustomAttributes.forEach { entry ->
+            assertEquals(entry.value, trace.attributes?.findAttributeValue(entry.key))
+        }
     }
 
     private fun assertChildSpan(span: EmbraceSpanData, expectedStartTimeNanos: Long, expectedEndTimeNanos: Long) {

embrace-android-sdk/api/embrace-android-sdk.api

@@ -8,6 +8,9 @@ public final class io/embrace/android/embracesdk/Embrace : io/embrace/android/em
 	public fun addLogRecordExporter (Lio/opentelemetry/sdk/logs/export/LogRecordExporter;)V
 	public fun addSessionProperty (Ljava/lang/String;Ljava/lang/String;Z)Z
 	public fun addSpanExporter (Lio/opentelemetry/sdk/trace/export/SpanExporter;)V
+	public fun addStartupChildSpan (Ljava/lang/String;JJ)V
+	public fun addStartupChildSpan (Ljava/lang/String;JJLjava/util/Map;Ljava/util/List;Lio/embrace/android/embracesdk/spans/ErrorCode;)V
+	public fun addStartupTraceAttribute (Ljava/lang/String;Ljava/lang/String;)V
 	public fun addUserPersona (Ljava/lang/String;)V
 	public fun clearAllUserPersonas ()V
 	public fun clearUserEmail ()V
@@ -26,6 +29,7 @@ public final class io/embrace/android/embracesdk/Embrace : io/embrace/android/em
 	public static final fun getInstance ()Lio/embrace/android/embracesdk/Embrace;
 	public fun getLastRunEndState ()Lio/embrace/android/embracesdk/LastRunEndState;
 	public fun getOpenTelemetry ()Lio/opentelemetry/api/OpenTelemetry;
+	public fun getSdkCurrentTimeMs ()J
 	public fun getSpan (Ljava/lang/String;)Lio/embrace/android/embracesdk/spans/EmbraceSpan;
 	public fun isStarted ()Z
 	public fun logCustomStacktrace ([Ljava/lang/StackTraceElement;)V

embrace-android-sdk/src/integrationTest/kotlin/io/embrace/android/embracesdk/testcases/AppStartupTraceTest.kt

@@ -0,0 +1,58 @@
+package io.embrace.android.embracesdk.testcases
+
+import android.os.Build
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import io.embrace.android.embracesdk.assertions.findSpansOfType
+import io.embrace.android.embracesdk.fakes.config.FakeEnabledFeatureConfig
+import io.embrace.android.embracesdk.fakes.config.FakeInstrumentedConfig
+import io.embrace.android.embracesdk.internal.arch.schema.EmbType
+import io.embrace.android.embracesdk.internal.spans.findAttributeValue
+import io.embrace.android.embracesdk.testframework.IntegrationTestRule
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.annotation.Config
+
+@Config(sdk = [Build.VERSION_CODES.LOLLIPOP])
+@RunWith(AndroidJUnit4::class)
+internal class AppStartupTraceTest {
+    @Rule
+    @JvmField
+    val testRule: IntegrationTestRule = IntegrationTestRule()
+
+    @Test
+    fun `startup spans recorded in foreground session when background activity is enabled`() {
+        testRule.runTest(
+            instrumentedConfig = FakeInstrumentedConfig(
+                enabledFeatures = FakeEnabledFeatureConfig(
+                    bgActivityCapture = true
+                )
+            ),
+            testCaseAction = {
+                val customStartTimeMs = clock.now()
+                val customEndTimeMs = clock.tick(100L)
+                embrace.addStartupChildSpan("custom-span", customStartTimeMs, customEndTimeMs)
+                embrace.addStartupTraceAttribute("custom-attribute", "yes")
+                simulateOpeningActivities(
+                    addStartupActivity = false,
+                    startInBackground = true
+                )
+            },
+            assertAction = {
+                with(getSingleSessionEnvelope()) {
+                    val spans = findSpansOfType(EmbType.Performance.Default).associateBy { it.name }
+                    assertTrue(spans.isNotEmpty())
+                    with(checkNotNull(spans["emb-cold-time-to-initial-display"])) {
+                        assertEquals("yes", attributes?.findAttributeValue("custom-attribute"))
+                    }
+                    assertTrue(spans.containsKey("emb-embrace-init"))
+                    assertTrue(spans.containsKey("custom-span"))
+                    assertTrue(spans.containsKey("emb-activity-create"))
+                    assertTrue(spans.containsKey("emb-activity-resume"))
+                }
+            }
+        )
+    }
+}

embrace-android-sdk/src/main/java/io/embrace/android/embracesdk/Embrace.kt

@@ -428,6 +428,8 @@ public class Embrace private constructor(
         impl.activityLoaded(activity)
     }
 
+    override fun getSdkCurrentTimeMs(): Long = impl.getSdkCurrentTimeMs()
+
     override fun addAttributeToLoadTrace(activity: Activity, key: String, value: String) {
         impl.addAttributeToLoadTrace(activity, key, value)
     }
@@ -451,4 +453,26 @@ public class Embrace private constructor(
             errorCode = errorCode
         )
     }
+
+    override fun addStartupTraceAttribute(key: String, value: String) {
+        impl.addStartupTraceAttribute(key, value)
+    }
+
+    override fun addStartupChildSpan(
+        name: String,
+        startTimeMs: Long,
+        endTimeMs: Long,
+        attributes: Map<String, String>,
+        events: List<EmbraceSpanEvent>,
+        errorCode: ErrorCode?,
+    ) {
+        impl.addStartupChildSpan(
+            name = name,
+            startTimeMs = startTimeMs,
+            endTimeMs = endTimeMs,
+            attributes = attributes,
+            events = events,
+            errorCode = errorCode
+        )
+    }
 }

embrace-android-sdk/src/main/java/io/embrace/android/embracesdk/internal/api/delegate/InstrumentationApiDelegate.kt

@@ -18,13 +18,18 @@ internal class InstrumentationApiDelegate(
     private val uiLoadTraceEmitter by embraceImplInject(sdkCallChecker) {
         bootstrapper.dataCaptureServiceModule.uiLoadDataListener
     }
+    private val appStartupDataCollector by embraceImplInject(sdkCallChecker) {
+        bootstrapper.dataCaptureServiceModule.appStartupDataCollector
+    }
 
     override fun activityLoaded(activity: Activity) {
         if (sdkCallChecker.check("activity_fully_loaded")) {
             uiLoadTraceEmitter?.complete(traceInstanceId(activity), clock.now())
         }
     }
 
+    override fun getSdkCurrentTimeMs(): Long = clock.now()
+
     override fun addAttributeToLoadTrace(activity: Activity, key: String, value: String) {
         if (sdkCallChecker.check("add_attribute_to_load_trace")) {
             uiLoadTraceEmitter?.addAttribute(traceInstanceId(activity), key, value)
@@ -52,4 +57,27 @@ internal class InstrumentationApiDelegate(
             )
         }
     }
+
+    override fun addStartupTraceAttribute(key: String, value: String) {
+        if (sdkCallChecker.check("add_attribute_to_app_startup_trace")) {
+            appStartupDataCollector?.addAttribute(key, value)
+        }
+    }
+
+    override fun addStartupChildSpan(
+        name: String,
+        startTimeMs: Long,
+        endTimeMs: Long,
+        attributes: Map<String, String>,
+        events: List<EmbraceSpanEvent>,
+        errorCode: ErrorCode?,
+    ) {
+        if (sdkCallChecker.check("add_child_span_to_app_startup_trace")) {
+            appStartupDataCollector?.addTrackedInterval(
+                name = name,
+                startTimeMs = startTimeMs,
+                endTimeMs = endTimeMs
+            )
+        }
+    }
 }