fix(compose): Stop SentryTraced from reusing stale parent spans

Fix the stale parent problem that causes spans from all SentryTraced instances for the entire app process to be dropped once the initial active transaction finishes.

Prior to this commit, SentryTraced used process-wide composition locals to bind the transaction active whenever the first SentryTraced for an app process entered the composition, and to reuse it for all SentryTraced instances thereafter. That meant all SentryTraced spans for the entire app would be dropped for the lifetime of the app process once the initial transaction finished. Oof.

This commit sets things right by having each SentryTraced composable request the current active transaction and update the generation of spans accordingly.

Because we're no longer relying on a single transaction + parent span pair, SentryTraced now needs to manage the creation of possibly multiple parent span pairs, as the owning transaction updates. That logic lives in the new ParentSpans class.

Author: 0xadam-brown

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+### Fixes
+
+`SentryTraced` now checks for its owning transaction dynamically rather than once per app process, which led to `SentryTraced` spans being dropped once the original transaction finished ([#6057](https://github.com/getsentry/sentry-java/pull/6057))
+
 ## 8.55.0
 
 ### Features

sentry-compose/build.gradle.kts

@@ -62,11 +62,14 @@ kotlin {
     }
     getByName("androidUnitTest") {
       dependencies {
+        implementation(libs.androidx.compose.foundation)
+        implementation(libs.androidx.compose.foundation.layout)
         implementation(libs.androidx.compose.ui.test.junit4)
         implementation(libs.androidx.navigation.compose)
         implementation(libs.androidx.test.ext.junit)
         implementation(libs.androidx.test.rules)
         implementation(libs.androidx.test.runner)
+        implementation(libs.google.truth)
         implementation(libs.kotlin.test.junit)
         implementation(libs.mockito.inline)
         implementation(libs.mockito.kotlin)

sentry-compose/src/androidMain/kotlin/io/sentry/compose/SentryComposeTracing.kt

@@ -4,74 +4,47 @@ import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.BoxScope
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.SideEffect
-import androidx.compose.runtime.compositionLocalOf
 import androidx.compose.runtime.remember
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.drawWithContent
 import io.sentry.ISpan
+import io.sentry.NoOpSpan
 import io.sentry.Sentry
 import io.sentry.SentryDate
 import io.sentry.SpanOptions
 import io.sentry.compose.SentryModifier.sentryTag
+import java.lang.ref.WeakReference
+import java.util.WeakHashMap
 
-private const val OP_PARENT_COMPOSITION = "ui.compose.composition"
-private const val OP_COMPOSE = "ui.compose"
+private const val DESCRIPTION_COMPOSITION_PARENT = "Jetpack Compose Initial Composition"
+private const val OP_COMPOSITION_PARENT = "ui.compose.composition"
+private const val OP_COMPOSITION_CHILD = "ui.compose"
 
-private const val OP_PARENT_RENDER = "ui.compose.rendering"
-private const val OP_RENDER = "ui.render"
+private const val DESCRIPTION_RENDER_PARENT = "Jetpack Compose Initial Render"
+private const val OP_RENDER_PARENT = "ui.compose.rendering"
+private const val OP_RENDER_CHILD = "ui.render"
 
 private const val OP_TRACE_ORIGIN = "auto.ui.jetpack_compose"
 
-private val localSentryCompositionParentSpan = compositionLocalOf {
-  getRootSpan()
-    // Create a single parent span to own composition spans emitted by all SentryTraced composables
-    // during the root's lifetime.
-    ?.startChild(
-      OP_PARENT_COMPOSITION,
-      "Jetpack Compose Initial Composition",
-      SpanOptions().apply {
-        isTrimStart = true
-        isTrimEnd = true
-        isIdle = true
-      },
-    )
-    ?.apply { spanContext.origin = OP_TRACE_ORIGIN }
-}
-
-private val localSentryRenderingParentSpan = compositionLocalOf {
-  getRootSpan()
-    // Create a single parent span to own render spans emitted by all SentryTraced composables
-    // during the root's lifetime.
-    ?.startChild(
-      OP_PARENT_RENDER,
-      "Jetpack Compose Initial Render",
-      SpanOptions().apply {
-        isTrimStart = true
-        isTrimEnd = true
-        isIdle = true
-      },
-    )
-    ?.apply { spanContext.origin = OP_TRACE_ORIGIN }
-}
-
-/**
- * A substitute for Compose's `MutableState` that doesn't register itself with the snapshot system,
- * so mutating [value] never triggers recomposition.
- */
-private class MutableRef<T>(var value: T)
-
 /**
- * Creates a single span for tracking the time required to compose the wrapped [content], and a span
- * for its initial draw.
+ * Creates a span for the initial composition of the wrapped [content], and a span for its initial
+ * rendering.
  *
  * Spans are approximate and include work performed by any composables [content] invokes. Abandoned
  * recompositions are ignored.
  *
- * Spans live under a set of parents shared by all `SentryTraced` composables. Every `SentryTraced`
- * contributes at most one `ui.compose` child and one `ui.render` child per parent lifetime:
+ * **Span organization**
+ *
+ * All spans produced are rooted under an owner span defined by the environment `SentryTraced` runs
+ * in. `SentryTraced` composables with the same owner share two common parent spans
+ * (`ui.compose.composition` and `ui.compose.rendering`). Each `SentryTraced` in the group emits at
+ * most one `ui.compose` span to the composition parent and one `ui.render` span to the render
+ * parent.
+ *
+ * The end result looks something like this:
  * ```
- * Root span
+ * Owner span
  * │
  * ├─ ui.compose.composition  "Jetpack Compose Initial Composition"
  * │   ├─ ui.compose   "product_info"
@@ -82,9 +55,8 @@ private class MutableRef<T>(var value: T)
  *     └─ ui.render    "add_to_cart_button"
  * ```
  *
- * Here `ui.compose.composition` and `ui.compose.rendering` are the shared parents. A `SentryTraced`
- * generates the "product_info" spans, and a separate `SentryTraced` generates the
- * "add_to_cart_button" spans.
+ * (Here, there were only two `SentryTraced` composables in the group. One emitted "product_info"
+ * spans, the other emitted "add_to_cart_button" spans.)
  */
 @ExperimentalComposeUiApi
 @Composable
@@ -95,31 +67,31 @@ public fun SentryTraced(
   content: @Composable BoxScope.() -> Unit,
 ) {
   val baseModifier = if (enableUserInteractionTracing) modifier.sentryTag(tag) else modifier
+  val scopes = Sentry.getCurrentScopes()
+  val ownerSpan = scopes.transaction ?: NoOpSpan.getInstance()
 
-  val parentCompositionSpan = localSentryCompositionParentSpan.current
-  val parentRenderingSpan = localSentryRenderingParentSpan.current
-
-  val alreadyComposed = remember(parentCompositionSpan) { MutableRef(false) }
-  val alreadyRendered = remember(parentRenderingSpan) { MutableRef(false) }
-  val dateProvider = Sentry.getCurrentScopes().options.dateProvider
+  val alreadyComposed = remember(ownerSpan) { MutableRef(false) }
+  val alreadyRendered = remember(ownerSpan) { MutableRef(false) }
+  val shouldRecordSpans = !ownerSpan.dropsChildSpans
 
-  // Only record spans if we have a parent for them.
+  val dateProvider = scopes.options.dateProvider
   val compositionStart =
-    if (!alreadyComposed.value) parentCompositionSpan?.let { dateProvider.now() } else null
+    if (shouldRecordSpans && !alreadyComposed.value) dateProvider.now() else null
 
   Box(
     modifier =
       baseModifier.drawWithContent {
-        if (alreadyRendered.value || parentRenderingSpan == null) {
+        if (!shouldRecordSpans || alreadyRendered.value) {
           drawContent()
-        } else {
-          val renderStart = dateProvider.now()
-          drawContent()
-          val renderEnd = dateProvider.now()
-
-          alreadyRendered.value = true
-          recordRenderSpan(parentRenderingSpan, tag, renderStart, renderEnd)
+          return@drawWithContent
         }
+
+        val renderStart = dateProvider.now()
+        drawContent()
+        val renderEnd = dateProvider.now()
+
+        alreadyRendered.value = true
+        recordRenderSpan(ownerSpan, tag, renderStart, renderEnd)
       },
     propagateMinConstraints = true,
   ) {
@@ -130,44 +102,137 @@ public fun SentryTraced(
     val compositionEnd = dateProvider.now()
 
     SideEffect {
-      recordCompositionSpan(
-        parentSpan = parentCompositionSpan,
-        tag = tag,
-        startTimestamp = compositionStart,
-        endTimestamp = compositionEnd,
-      )
-
       alreadyComposed.value = true
+      recordCompositionSpan(ownerSpan, tag, compositionStart, compositionEnd)
     }
   }
 }
 
-private fun getRootSpan(): ISpan? {
-  var rootSpan: ISpan? = null
-  Sentry.configureScope { rootSpan = it.transaction }
-  return rootSpan
-}
-
 private fun recordCompositionSpan(
-  parentSpan: ISpan?,
+  ownerSpan: ISpan,
   tag: String,
   startTimestamp: SentryDate,
   endTimestamp: SentryDate,
 ) {
-  parentSpan?.startChild(OP_COMPOSE, tag, startTimestamp)?.apply {
+  val parentSpan = ParentSpans.getOrCreateCompositionSpan(ownerSpan, startTimestamp) ?: return
+
+  parentSpan.startChild(OP_COMPOSITION_CHILD, tag, startTimestamp).apply {
     spanContext.origin = OP_TRACE_ORIGIN
     finish(null, endTimestamp)
   }
 }
 
 private fun recordRenderSpan(
-  parentSpan: ISpan?,
+  ownerSpan: ISpan,
   tag: String,
   startTimestamp: SentryDate,
   endTimestamp: SentryDate,
 ) {
-  parentSpan?.startChild(OP_RENDER, tag, startTimestamp)?.apply {
+  val parentSpan = ParentSpans.getOrCreateRenderSpan(ownerSpan, startTimestamp) ?: return
+
+  parentSpan.startChild(OP_RENDER_CHILD, tag, startTimestamp).apply {
     spanContext.origin = OP_TRACE_ORIGIN
     finish(null, endTimestamp)
   }
 }
+
+/**
+ * Returns true if spans parented under the receiver will be dropped (and therefore aren't worth
+ * creating in the first place).
+ */
+private val ISpan.dropsChildSpans: Boolean
+  // NoOp spans return false for isFinished, so we check for no-op status directly.
+  get() = this.isFinished || this.isNoOp
+
+/**
+ * Weak holder of [OP_COMPOSITION_PARENT] and [OP_RENDER_PARENT] spans generated by [SentryTraced]
+ * composables.
+ *
+ * Instances are cached in a weak map keyed by owner span, so multiple [SentryTraced] composables
+ * with the same owner share one composition parent and one render parent.
+ *
+ * **Not threadsafe:** Access must be confined to Compose UI-thread callbacks.
+ */
+private class ParentSpans {
+
+  private var compositionParentSpan: WeakReference<ISpan>? = null
+  private var renderParentSpan: WeakReference<ISpan>? = null
+
+  companion object {
+
+    private val ownerSpanToParentSpans = WeakHashMap<ISpan, ParentSpans>()
+
+    fun getOrCreateCompositionSpan(ownerSpan: ISpan, startTimestamp: SentryDate): ISpan? =
+      getFor(ownerSpan).getOrCreateCompositionSpan(ownerSpan, startTimestamp)
+
+    fun getOrCreateRenderSpan(ownerSpan: ISpan, startTimestamp: SentryDate): ISpan? =
+      getFor(ownerSpan).getOrCreateRenderSpan(ownerSpan, startTimestamp)
+
+    private fun getFor(ownerSpan: ISpan): ParentSpans =
+      ownerSpanToParentSpans.getOrPut(ownerSpan) { ParentSpans() }
+  }
+
+  private fun getOrCreateCompositionSpan(ownerSpan: ISpan, startTimestamp: SentryDate): ISpan? =
+    getOrCreate(
+      ownerSpan = ownerSpan,
+      startTimestamp = startTimestamp,
+      cached = compositionParentSpan,
+      operation = OP_COMPOSITION_PARENT,
+      description = DESCRIPTION_COMPOSITION_PARENT,
+    ) {
+      compositionParentSpan = it
+    }
+
+  private fun getOrCreateRenderSpan(ownerSpan: ISpan, startTimestamp: SentryDate): ISpan? =
+    getOrCreate(
+      ownerSpan = ownerSpan,
+      startTimestamp = startTimestamp,
+      cached = renderParentSpan,
+      operation = OP_RENDER_PARENT,
+      description = DESCRIPTION_RENDER_PARENT,
+    ) {
+      renderParentSpan = it
+    }
+
+  private fun getOrCreate(
+    ownerSpan: ISpan,
+    startTimestamp: SentryDate,
+    cached: WeakReference<ISpan>?,
+    operation: String,
+    description: String,
+    setCached: (WeakReference<ISpan>) -> Unit,
+  ): ISpan? {
+    cached
+      ?.get()
+      ?.takeUnless { it.dropsChildSpans }
+      ?.let {
+        return it
+      }
+
+    val parentSpan =
+      ownerSpan.startChild(
+        operation,
+        description,
+        SpanOptions().apply {
+          setStartTimestamp(startTimestamp)
+          isTrimStart = true
+          isTrimEnd = true
+          isIdle = true
+        },
+      )
+
+    if (parentSpan.dropsChildSpans) {
+      return null
+    }
+
+    parentSpan.spanContext.origin = OP_TRACE_ORIGIN
+    setCached(WeakReference(parentSpan))
+    return parentSpan
+  }
+}
+
+/**
+ * A substitute for Compose's `MutableState` that doesn't register itself with the snapshot system,
+ * so mutating [value] won't trigger recomposition.
+ */
+private class MutableRef<T>(var value: T)