fix(compose): allow overriding the IScopes used by SentryTraced via LocalSentryScopes

CHANGELOG.md

@@ -2,6 +2,18 @@
 
 ## Unreleased
 
+### Features
+
+- Add `LocalSentryScopes` to `sentry-compose`, allowing `SentryTraced` to trace against a custom `IScopes` instance instead of always defaulting to `Sentry.getCurrentScopes()` ([#5838](https://github.com/getsentry/sentry-java/pull/5838))
+  - Example usage:
+  ```kotlin
+  val scopes = Scopes(options)
+  CompositionLocalProvider(LocalSentryScopes provides scopes) {
+    // this uses custom scopes now
+    SentryTraced(tag = "custom") { Box {} }
+  }
+  ```
+
 ### Improvements
 
 - Skip building Android manifest metadata debug log messages when debug logging is disabled, reducing allocations during SDK init ([#5790](https://github.com/getsentry/sentry-java/pull/5790))

sentry-compose/api/android/sentry-compose.api

@@ -12,6 +12,7 @@ public final class io/sentry/compose/SentryComposeHelperKt {
 
 public final class io/sentry/compose/SentryComposeTracingKt {
 	public static final fun SentryTraced (Ljava/lang/String;Landroidx/compose/ui/Modifier;ZLkotlin/jvm/functions/Function3;Landroidx/compose/runtime/Composer;II)V
+	public static final fun getLocalSentryScopes ()Landroidx/compose/runtime/ProvidableCompositionLocal;
 }
 
 public final class io/sentry/compose/SentryModifier {

sentry-compose/build.gradle.kts

@@ -68,9 +68,11 @@ kotlin {
         implementation(libs.androidx.test.rules)
         implementation(libs.androidx.test.runner)
         implementation(libs.kotlin.test.junit)
+        implementation(libs.androidx.compose.foundation)
         implementation(libs.mockito.inline)
         implementation(libs.mockito.kotlin)
         implementation(libs.roboelectric)
+        implementation(projects.sentryTestSupport)
       }
     }
   }

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

@@ -3,12 +3,15 @@ package io.sentry.compose
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.BoxScope
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.Immutable
-import androidx.compose.runtime.compositionLocalOf
+import androidx.compose.runtime.ProvidableCompositionLocal
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.staticCompositionLocalOf
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.drawWithContent
+import io.sentry.IScopes
 import io.sentry.ISpan
 import io.sentry.Sentry
 import io.sentry.SpanOptions
@@ -24,15 +27,24 @@ private const val OP_TRACE_ORIGIN = "auto.ui.jetpack_compose"
 
 @Immutable private class ImmutableHolder<T>(var item: T)
 
-private fun getRootSpan(): ISpan? {
+// Defaults to null rather than eagerly resolving Sentry.getCurrentScopes(): a CompositionLocal's
+// default value is computed at most once for the process lifetime, so a non-null default would
+// permanently cache whatever scopes were current on first read (e.g. a NoOp scopes if read before
+// Sentry.init()). SentryTraced instead falls back to Sentry.getCurrentScopes() on every call when
+// no value has been explicitly provided, so it always reflects the current scopes.
+public val LocalSentryScopes: ProvidableCompositionLocal<IScopes?> = staticCompositionLocalOf {
+  null
+}
+
+private fun getRootSpan(scopes: IScopes): ISpan? {
   var rootSpan: ISpan? = null
-  Sentry.configureScope { rootSpan = it.transaction }
+  scopes.configureScope { rootSpan = it.transaction }
   return rootSpan
 }
 
-private val localSentryCompositionParentSpan = compositionLocalOf {
+private fun createCompositionParentSpan(scopes: IScopes): ImmutableHolder<ISpan?> =
   ImmutableHolder(
-    getRootSpan()
+    getRootSpan(scopes)
       ?.startChild(
         OP_PARENT_COMPOSITION,
         "Jetpack Compose Initial Composition",
@@ -44,11 +56,10 @@ private val localSentryCompositionParentSpan = compositionLocalOf {
       )
       ?.apply { spanContext.origin = OP_TRACE_ORIGIN }
   )
-}
 
-private val localSentryRenderingParentSpan = compositionLocalOf {
+private fun createRenderingParentSpan(scopes: IScopes): ImmutableHolder<ISpan?> =
   ImmutableHolder(
-    getRootSpan()
+    getRootSpan(scopes)
       ?.startChild(
         OP_PARENT_RENDER,
         "Jetpack Compose Initial Render",
@@ -60,8 +71,49 @@ private val localSentryRenderingParentSpan = compositionLocalOf {
       )
       ?.apply { spanContext.origin = OP_TRACE_ORIGIN }
   )
+
+private class RootSpans {
+  // Entries are cleaned up deterministically via reference counting (see retain/release) rather
+  // than relying on GC to reclaim a weakly-keyed map: a value here (ImmutableHolder -> Span)
+  // strongly references the IScopes that created it (Span keeps its Scopes alive), so a
+  // WeakHashMap keyed on IScopes would never actually evict its own key, and a value wrapped in a
+  // WeakReference could be collected between recompositions even while a SentryTraced call under
+  // that scopes is still in composition, silently breaking root span sharing.
+  val compositionSpans = HashMap<IScopes, ImmutableHolder<ISpan?>>()
+  val renderingSpans = HashMap<IScopes, ImmutableHolder<ISpan?>>()
+  private val refCounts = HashMap<IScopes, Int>()
+
+  fun retain(scopes: IScopes) {
+    refCounts[scopes] = (refCounts[scopes] ?: 0) + 1
+  }
+
+  fun release(scopes: IScopes) {
+    val remaining = (refCounts[scopes] ?: 1) - 1
+    if (remaining <= 0) {
+      refCounts.remove(scopes)
+      compositionSpans.remove(scopes)
+      renderingSpans.remove(scopes)
+    } else {
+      refCounts[scopes] = remaining
+    }
+  }
 }
 
+private fun getOrCreateParentSpan(
+  map: MutableMap<IScopes, ImmutableHolder<ISpan?>>,
+  scopes: IScopes,
+  create: (IScopes) -> ImmutableHolder<ISpan?>,
+): ImmutableHolder<ISpan?> =
+  // Only cache the holder once it actually contains a span; a null result (no transaction bound
+  // to the scopes yet) is recomputed on the next call so a later transaction is still picked up.
+  map[scopes] ?: create(scopes).also { if (it.item != null) map[scopes] = it }
+
+// Cached once per Composition and shared by every SentryTraced call within it, mirroring the
+// old eagerly-computed `compositionLocalOf { ... }` default (which Compose resolves once and
+// reuses for every `.current` read that has no ancestor Provider, sibling or not). Keyed per
+// IScopes so distinct custom scopes each get their own root span instead of colliding.
+private val LocalRootSpans = staticCompositionLocalOf { RootSpans() }
+
 @ExperimentalComposeUiApi
 @Composable
 public fun SentryTraced(
@@ -70,8 +122,20 @@ public fun SentryTraced(
   enableUserInteractionTracing: Boolean = true,
   content: @Composable BoxScope.() -> Unit,
 ) {
-  val parentCompositionSpan = localSentryCompositionParentSpan.current
-  val parentRenderingSpan = localSentryRenderingParentSpan.current
+  val scopes = LocalSentryScopes.current ?: Sentry.getCurrentScopes()
+  val rootSpans = LocalRootSpans.current
+  // Retain synchronously during composition (not inside a DisposableEffect) so it always runs
+  // before any effect-phase work for this frame, including another SentryTraced call's dispose:
+  // Compose runs the dispose of an outgoing node before the effects of an incoming one in the
+  // same recomposition, so retaining from an effect could let the shared entry's refcount hit
+  // zero (and get evicted) between an old and a new SentryTraced call sharing the same scopes.
+  remember(rootSpans, scopes) { rootSpans.retain(scopes) }
+  DisposableEffect(rootSpans, scopes) { onDispose { rootSpans.release(scopes) } }
+  val parentCompositionSpan =
+    getOrCreateParentSpan(rootSpans.compositionSpans, scopes, ::createCompositionParentSpan)
+  val parentRenderingSpan =
+    getOrCreateParentSpan(rootSpans.renderingSpans, scopes, ::createRenderingParentSpan)
+
   val compositionSpan =
     parentCompositionSpan.item?.startChild(OP_COMPOSE, tag)?.apply {
       spanContext.origin = OP_TRACE_ORIGIN

sentry-compose/src/androidUnitTest/kotlin/io/sentry/compose/SentryTracedTest.kt

@@ -0,0 +1,150 @@
+package io.sentry.compose
+
+import android.app.Application
+import android.content.ComponentName
+import androidx.activity.ComponentActivity
+import androidx.compose.foundation.layout.Box
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.key
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.ExperimentalComposeUiApi
+import androidx.compose.ui.test.junit4.createAndroidComposeRule
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import io.sentry.IScopes
+import io.sentry.ITransaction
+import io.sentry.Sentry
+import io.sentry.SentryOptions
+import io.sentry.TransactionOptions
+import io.sentry.test.createTestScopes
+import kotlin.test.assertEquals
+import org.junit.After
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TestWatcher
+import org.junit.runner.Description
+import org.junit.runner.RunWith
+import org.robolectric.Shadows
+import org.robolectric.annotation.Config
+
+@OptIn(ExperimentalComposeUiApi::class)
+@RunWith(AndroidJUnit4::class)
+@Config(sdk = [30])
+class SentryTracedTest {
+  // workaround for robolectric tests with composeRule
+  // from https://github.com/robolectric/robolectric/pull/4736#issuecomment-1831034882
+  @get:Rule(order = 1)
+  val addActivityToRobolectricRule =
+    object : TestWatcher() {
+      override fun starting(description: Description?) {
+        super.starting(description)
+        val appContext: Application = ApplicationProvider.getApplicationContext()
+        Shadows.shadowOf(appContext.packageManager)
+          .addActivityIfNotPresent(
+            ComponentName(appContext.packageName, ComponentActivity::class.java.name)
+          )
+      }
+    }
+
+  @get:Rule(order = 2) val rule = createAndroidComposeRule<ComponentActivity>()
+
+  @After
+  fun tearDown() {
+    Sentry.close()
+  }
+
+  private fun newTracingScopes(): IScopes =
+    createTestScopes(
+      SentryOptions().apply {
+        dsn = "https://key@sentry.io/proj"
+        tracesSampleRate = 1.0
+      }
+    )
+
+  private fun IScopes.startBoundTransaction(name: String): ITransaction =
+    startTransaction(name, "test", TransactionOptions().apply { isBindToScope = true })
+
+  @Test
+  fun `SentryTraced creates its root span on the scopes provided via LocalSentryScopes`() {
+    val scopes = newTracingScopes()
+    val tx = scopes.startBoundTransaction("custom-scopes-tx")
+
+    rule.setContent {
+      CompositionLocalProvider(LocalSentryScopes provides scopes) {
+        SentryTraced(tag = "custom") { Box {} }
+      }
+    }
+    rule.waitForIdle()
+
+    assertEquals(1, tx.spans.count { it.operation == "ui.compose.composition" })
+    assertEquals(1, tx.spans.count { it.operation == "ui.compose" })
+  }
+
+  @Test
+  fun `sibling SentryTraced composables under the same scopes share one root span`() {
+    val scopes = newTracingScopes()
+    val tx = scopes.startBoundTransaction("custom-scopes-tx")
+
+    rule.setContent {
+      CompositionLocalProvider(LocalSentryScopes provides scopes) {
+        SentryTraced(tag = "first") { Box {} }
+        SentryTraced(tag = "second") { Box {} }
+      }
+    }
+    rule.waitForIdle()
+
+    assertEquals(1, tx.spans.count { it.operation == "ui.compose.composition" })
+    assertEquals(2, tx.spans.count { it.operation == "ui.compose" })
+  }
+
+  @Test
+  fun `SentryTraced composables under different scopes do not interfere with each other`() {
+    val scopesA = newTracingScopes()
+    val txA = scopesA.startBoundTransaction("scopes-a-tx")
+    val scopesB = newTracingScopes()
+    val txB = scopesB.startBoundTransaction("scopes-b-tx")
+
+    rule.setContent {
+      CompositionLocalProvider(LocalSentryScopes provides scopesA) {
+        SentryTraced(tag = "a") { Box {} }
+      }
+      CompositionLocalProvider(LocalSentryScopes provides scopesB) {
+        SentryTraced(tag = "b") { Box {} }
+      }
+    }
+    rule.waitForIdle()
+
+    assertEquals(1, txA.spans.count { it.operation == "ui.compose.composition" })
+    assertEquals(1, txA.spans.count { it.operation == "ui.compose" })
+    assertEquals(1, txB.spans.count { it.operation == "ui.compose.composition" })
+    assertEquals(1, txB.spans.count { it.operation == "ui.compose" })
+  }
+
+  @Test
+  fun `repeatedly replacing a traced composable under the same scopes keeps sharing the root span`() {
+    // Each swap disposes the outgoing keyed composable and mounts a new one under the same
+    // scopes within a single recomposition. Compose dispatches the outgoing composable's
+    // onDispose before the incoming one's DisposableEffect runs, so a second swap is needed to
+    // surface a root span cache that was cleared out from under a still-live retain.
+    val scopes = newTracingScopes()
+    val tx = scopes.startBoundTransaction("custom-scopes-tx")
+    var step by mutableStateOf(0)
+
+    rule.setContent {
+      CompositionLocalProvider(LocalSentryScopes provides scopes) {
+        key(step) { SentryTraced(tag = "step-$step") { Box {} } }
+      }
+    }
+    rule.waitForIdle()
+
+    step = 1
+    rule.waitForIdle()
+
+    step = 2
+    rule.waitForIdle()
+
+    assertEquals(1, tx.spans.count { it.operation == "ui.compose.composition" })
+  }
+}