Fix a race condition when tracing Kotlin coroutines.

* When using a `ThreadContextElement` to track suspensions and resumes, its important that we make
 `updateThreadContext` and `restoreThreadContext` safe, when calls to them are concurrent and interleaved.

* For e.g. this is something that can happen when creating a child coroutine that in turn uses a multi threaded dispatcher.

```
coroutine.updateThreadContext() // Thread #1
... coroutine body ...
// suspension + immediate dispatch happen here
coroutine.updateThreadContext() // Thread #2, coroutine is already resumed
// ... coroutine body after suspension point on Thread #2 ...
coroutine.restoreThreadContext() // Thread #1, is invoked late because Thread #1 is slow
coroutine.restoreThreadContext() // Thread #2, may happen in parallel with the previous restore
```

* Previously when `updateThreadContext(...)` would get called on `Thread #2`, we won't emit an additional `begin` slice, for the new track because we are in `STATE_BEGIN` on the old thread.
* When ending a slice, we could end on `Thread #2` and not on `Thread #1` which ends up corrupting traces.

Test: This is something very difficult to to test, but I added additional integration tests, that involve jumping threads.
Example trace: https://ui.perfetto.dev/#!/?s=575b99d9a235654f9de85aadaaa0ce63f31e4770

Change-Id: I18a49c3f6b65fa41ab5f9ad855377ac779aa4326

Author: tikurahul

tracing/tracing-wire/src/desktopTest/kotlin/androidx/tracing/wire/TracingDemoTest.kt

@@ -120,6 +120,20 @@ class TracingDemoTest {
         }
     }
 
+    @Test
+    internal fun testContextSwitching(): Unit = runBlocking {
+        var index = 0
+        driver.use {
+            tracer.traceCoroutine(category = "category", name = "switch") {
+                repeat(100) {
+                    index += 1
+                    val dispatcher = if (index % 2 == 0) Dispatchers.Default else Dispatchers.IO
+                    withContext(dispatcher) { delay(100) }
+                }
+            }
+        }
+    }
+
     @Test
     // The amount of time spent sleeping does not affect the outcome of the test.
     @Suppress("BanThreadSleep")

tracing/tracing/src/commonMain/kotlin/androidx/tracing/ContextElements.kt

@@ -16,6 +16,8 @@
 
 package androidx.tracing
 
+import androidx.collection.MutableLongIntMap
+import androidx.collection.mutableLongIntMapOf
 import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 
@@ -40,14 +42,30 @@ internal constructor(
      */
     public open val flowIds: List<Long>,
 ) : AbstractCoroutineContextElement(key = KEY), PropagationToken, AutoCloseable {
-    // Always starts in a `begin` state.
-    @JvmField internal var started: Int = STATE_BEGIN
+
+    // Calls to update and restore will happen concurrently on multiple threads, if we happen
+    // to jump threads between suspends and resumes. Therefore, it's important for us to track
+    // begin and end per thread id. For more information refer to the `Reentrancy and thread-safety`
+    // section in the documentation for `CopyableThreadContextElement`.
+    @JvmField internal val started: MutableLongIntMap = mutableLongIntMapOf()
+
+    init {
+        synchronized(this) {
+            // Always starts in a `begin` state.
+            started[currentThreadId()] = STATE_BEGIN
+        }
+    }
 
     @Suppress("NOTHING_TO_INLINE")
     internal inline fun synchronizedCompareAndSet(expected: Int, newValue: Int): Boolean {
         return synchronized(this) {
-            if (started == expected) {
-                started = newValue
+            // Treat the absence of an entry here as `STATE_END` given we have not emitted a
+            // trace packet yet on this Thread. This is true for every child coroutine, but
+            // **not** for the coroutine that kicked things off (handled by the init block).
+            val id = currentThreadId()
+            val current = started.getOrDefault(key = id, defaultValue = STATE_END)
+            if (current == expected) {
+                started[id] = newValue
                 true
             } else {
                 false

tracing/tracing/src/commonMain/kotlin/androidx/tracing/Threads.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2026 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.tracing
+
+internal expect inline fun currentThreadId(): Long

tracing/tracing/src/jvmAndAndroidMain/kotlin/androidx/tracing/Threads.jvmAndAndroid.kt

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2026 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.tracing
+
+@Suppress("NOTHING_TO_INLINE", "DEPRECATION")
+internal actual inline fun currentThreadId(): Long {
+    return Thread.currentThread().id
+}