Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit 625faabe authored by Peter Kalauskas's avatar Peter Kalauskas Committed by Automerger Merge Worker
Browse files

tracinglib: new helper for naming flow scopes am: 32ce2d93

parents 1118acc6 32ce2d93
Loading
Loading
Loading
Loading
+13 −14
Original line number Diff line number Diff line
# Coroutine Tracing

This library contains utilities for tracing coroutines. Coroutines cannot normally be traced using
the `android.os.Trace` APIs because it will often lead to malformed trace sections. This is because
each `Trace.beginSection` must have a matching `Trace.endSection` on the same thread before the
scope is finished, so if they are used around a suspend point, the trace section will remain open
while other unrelated work executes. It could even remain open indefinitely if the coroutine is
canceled.

To address this, we introduce a function `traceCoroutine("name") {}` that can be used for tracing
sections of coroutine code. When invoked, a trace section with the given name will start
This library contains utilities for tracing coroutines. Coroutines cannot be traced using the
`android.os.Trace` APIs because suspension points will lead to malformed trace sections. This is
because each `Trace.beginSection` must have a matching `Trace.endSection`; if a coroutine suspends
before `Trace.endSection` is called, the trace section will remain open while other unrelated work
executes.

To address this, we introduce a function `traceCoroutine("name") { ... }` that can be used for
tracing sections of coroutine code. When invoked, a trace section with the given name will start
immediately, and its name will also be written to an object in the current `CoroutineContext` used
for coroutine-local storage. When the coroutine suspends, all trace sections will end immediately.
When resumed, the coroutine will read the names of the previous sections from coroutine-local
@@ -61,19 +60,19 @@ val coldFlow = flow {
  emit(1)
  emit(2)
  emit(3)
}.withTraceName("my-flow")
}

coldFlow.collect {
coldFlow.collect("F") {
  println(it)
  delay(10)
  yield()
}
```

Would be traced as follows:

```
Thread #1 |  [=== my-flow:collect ===]    [=== my-flow:collect ===]    [=== my-flow:collect ===]
          |    [== my-flow:emit ==]         [== my-flow:emit ==]         [== my-flow:emit ==]
Thread #1 |  [====== collect:F ======]    [==== collect:F =====]    [====== collect:F ======]
          |    [== collect:F:emit ==]     [== collect:F:emit ==]    [== collect:F:emit ==]
```

# Building and Running
+40 −46
Original line number Diff line number Diff line
@@ -32,17 +32,15 @@ import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withContext

const val DEFAULT_TRACK_NAME = "Coroutines"

@OptIn(ExperimentalContracts::class)
suspend inline fun <R> coroutineScope(
suspend inline fun <R> coroutineScopeTraced(
    traceName: String,
    crossinline block: suspend CoroutineScope.() -> R,
): R {
    contract { callsInPlace(block, InvocationKind.EXACTLY_ONCE) }
    return traceCoroutine(traceName) {
        return@traceCoroutine coroutineScope wrappedCoroutineScope@{
            return@wrappedCoroutineScope block()
    return coroutineScope {
        traceCoroutine(traceName) {
            return@coroutineScope block()
        }
    }
}
@@ -52,31 +50,33 @@ suspend inline fun <R> coroutineScope(
 *
 * @see traceCoroutine
 */
inline fun CoroutineScope.launch(
inline fun CoroutineScope.launchTraced(
    crossinline spanName: () -> String,
    context: CoroutineContext = EmptyCoroutineContext,
    start: CoroutineStart = CoroutineStart.DEFAULT,
    noinline block: suspend CoroutineScope.() -> Unit,
): Job = launch(nameCoroutine(spanName) + context, start, block)
): Job {
    return launch(nameCoroutine(spanName) + context, start, block)
}

/**
 * Convenience function for calling [CoroutineScope.launch] with [traceCoroutine] to enable tracing.
 *
 * @see traceCoroutine
 */
fun CoroutineScope.launch(
    spanName: String,
fun CoroutineScope.launchTraced(
    spanName: String? = null,
    context: CoroutineContext = EmptyCoroutineContext,
    start: CoroutineStart = CoroutineStart.DEFAULT,
    block: suspend CoroutineScope.() -> Unit,
): Job = launch(nameCoroutine(spanName) + context, start, block)
): Job = launchTraced({ spanName ?: block::class.simpleName ?: "launch" }, context, start, block)

/**
 * Convenience function for calling [CoroutineScope.async] with [traceCoroutine] enable tracing
 *
 * @see traceCoroutine
 */
inline fun <T> CoroutineScope.async(
inline fun <T> CoroutineScope.asyncTraced(
    spanName: () -> String,
    context: CoroutineContext = EmptyCoroutineContext,
    start: CoroutineStart = CoroutineStart.DEFAULT,
@@ -88,19 +88,20 @@ inline fun <T> CoroutineScope.async(
 *
 * @see traceCoroutine
 */
fun <T> CoroutineScope.async(
    spanName: String,
fun <T> CoroutineScope.asyncTraced(
    spanName: String? = null,
    context: CoroutineContext = EmptyCoroutineContext,
    start: CoroutineStart = CoroutineStart.DEFAULT,
    block: suspend CoroutineScope.() -> T,
): Deferred<T> = async(nameCoroutine(spanName) + context, start, block)
): Deferred<T> =
    asyncTraced({ spanName ?: block::class.simpleName ?: "async" }, context, start, block)

/**
 * Convenience function for calling [runBlocking] with [traceCoroutine] to enable tracing.
 *
 * @see traceCoroutine
 */
inline fun <T> runBlocking(
inline fun <T> runBlockingTraced(
    spanName: () -> String,
    context: CoroutineContext,
    noinline block: suspend CoroutineScope.() -> T,
@@ -111,29 +112,29 @@ inline fun <T> runBlocking(
 *
 * @see traceCoroutine
 */
fun <T> runBlocking(
    spanName: String,
fun <T> runBlockingTraced(
    spanName: String? = null,
    context: CoroutineContext,
    block: suspend CoroutineScope.() -> T,
): T = runBlocking(nameCoroutine(spanName) + context, block)
): T = runBlockingTraced({ spanName ?: block::class.simpleName ?: "runBlocking" }, context, block)

/**
 * Convenience function for calling [withContext] with [traceCoroutine] to enable tracing.
 *
 * @see traceCoroutine
 */
suspend fun <T> withContext(
    spanName: String,
suspend fun <T> withContextTraced(
    spanName: String? = null,
    context: CoroutineContext,
    block: suspend CoroutineScope.() -> T,
): T = withContext(nameCoroutine(spanName) + context, block)
): T = withContextTraced({ spanName ?: block::class.simpleName ?: "withContext" }, context, block)

/**
 * Convenience function for calling [withContext] with [traceCoroutine] to enable tracing.
 *
 * @see traceCoroutine
 */
suspend inline fun <T> withContext(
suspend inline fun <T> withContextTraced(
    spanName: () -> String,
    context: CoroutineContext,
    noinline block: suspend CoroutineScope.() -> T,
@@ -143,37 +144,30 @@ suspend inline fun <T> withContext(
 * Traces a section of work of a `suspend` [block]. The trace sections will appear on the thread
 * that is currently executing the [block] of work. If the [block] is suspended, all trace sections
 * added using this API will end until the [block] is resumed, which could happen either on this
 * thread or on another thread. If a child coroutine is started, it will inherit the trace sections
 * of its parent. The child will continue to print these trace sections whether or not the parent
 * coroutine is still running them.
 * thread or on another thread. If a child coroutine is started, it will *NOT* inherit the trace
 * sections of its parent; however, it will include metadata in the trace section pointing to the
 * parent.
 *
 * The current [CoroutineContext] must have a [TraceContextElement] for this API to work. Otherwise,
 * the trace sections will be dropped.
 *
 * For example, in the following trace, Thread #1 ran some work, suspended, then continued working
 * on Thread #2. Meanwhile, Thread #2 created a new child coroutine which inherited its trace
 * sections. Then, the original coroutine resumed on Thread #1 before ending. Meanwhile Thread #3 is
 * still printing trace sections from its parent because they were copied when it was created. There
 * is no way for the parent to communicate to the child that it marked these slices as completed.
 * While this might seem counterintuitive, it allows us to pinpoint the origin of the child
 * coroutine's work.
 * For example, in the following trace, Thread #1 starts a coroutine, suspends, and continues the
 * coroutine on Thread #2. Next, Thread #2 start a child coroutine in an unconfined manner. Then,
 * still on Thread #2, the original coroutine suspends, the child resumes, and the child suspends.
 * Then, the original coroutine resumes on Thread#1.
 *
 * ```
 * Thread #1 | [==== Slice A ====]                        [==== Slice A ====]
 *           |       [==== B ====]                        [=== B ===]
 * --------------------------------------------------------------------------------------
 * Thread #2 |                    [====== Slice A ======]
 *           |                    [========= B =========]
 *           |                        [===== C ======]
 * --------------------------------------------------------------------------------------
 * Thread #3 |                            [== Slice A ==]                [== Slice A ==]
 *           |                            [===== B =====]                [===== B =====]
 *           |                            [===== C =====]                [===== C =====]
 *           |                                                               [=== D ===]
 * -----------------------------------------------------------------------------------------------|
 * Thread #1 | [== Slice A ==]                                            [==== Slice A ====]
 *           |       [== B ==]                                            [=== B ===]
 * -----------------------------------------------------------------------------------------------|
 * Thread #2 |                     [==== Slice A ====]    [=== C ====]
 *           |                     [======= B =======]
 *           |                         [=== C ====]
 * -----------------------------------------------------------------------------------------------|
 * ```
 *
 * @param name The name of the code section to appear in the trace
 * @see endSlice
 * @param spanName The name of the code section to appear in the trace
 * @see traceCoroutine
 */
@OptIn(ExperimentalContracts::class)
+48 −92
Original line number Diff line number Diff line
@@ -28,11 +28,8 @@ import java.util.stream.Stream
import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract
import kotlin.coroutines.AbstractCoroutineContextKey
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.coroutines.getPolymorphicElement
import kotlin.coroutines.minusPolymorphicKey
import kotlinx.coroutines.CopyableThreadContextElement
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.DelicateCoroutinesApi
@@ -92,19 +89,30 @@ val traceThreadLocal = TraceDataThreadLocal()
 *   names.
 * @param strictMode whether to add additional checks to coroutine tracing machinery. These checks
 *   are expensive and should only be used for testing.
 * @param shouldIgnoreClassName lambda that takes binary class name (as returned from
 *   [StackFrame.getClassName] and returns true if it should be ignored (e.g. search for relevant
 *   class name should continue) or false otherwise
 */
fun createCoroutineTracingContext(
    name: String = "UnnamedScope",
    walkStackForDefaultNames: Boolean = false,
    includeParentNames: Boolean = false,
    strictMode: Boolean = false,
    shouldIgnoreClassName: (String) -> Boolean = { false },
): CoroutineContext =
    if (Flags.coroutineTracing()) {
        TraceContextElement(
            scopeName = name,
            name = name,
            // Minor perf optimization: no need to create TraceData() for root scopes since all
            // launches require creation of child via [copyForChild] or [mergeForChild].
            contextTraceData = null,
            inheritedTracePrefix = "",
            coroutineDepth = 0,
            parentId = -1,
            walkStackForDefaultNames = walkStackForDefaultNames,
            includeParentNames = includeParentNames,
            strictMode = strictMode,
            shouldIgnoreClassName = shouldIgnoreClassName,
        )
    } else {
        EmptyCoroutineContext
@@ -141,29 +149,21 @@ inline fun nameCoroutine(name: () -> String): CoroutineContext {
 *
 * @property name the name of the current coroutine
 */
/**
 * A coroutine context element that can be used for naming the child coroutine under construction.
 *
 * @property name the name to be used for the child under construction
 * @see nameCoroutine
 */
@VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
open class BaseTraceElement(val name: String) : CoroutineContext.Element {
    companion object Key : CoroutineContext.Key<BaseTraceElement>
open class CoroutineTraceName(val name: String) : CoroutineContext.Element {
    companion object Key : CoroutineContext.Key<CoroutineTraceName>

    override val key: CoroutineContext.Key<*>
        get() = Key

    protected val currentId = ThreadLocalRandom.current().nextInt(1, Int.MAX_VALUE)

    @OptIn(ExperimentalStdlibApi::class)
    override fun <E : CoroutineContext.Element> get(key: CoroutineContext.Key<E>): E? {
        val rv = getPolymorphicElement(key)
        debug { "#get($key)=$rv" }
        return rv
    }

    @OptIn(ExperimentalStdlibApi::class)
    override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext {
        val rv = minusPolymorphicKey(key)
        debug { "#minusKey($key)=$rv" }
        return rv
    }

    @Deprecated(
        message =
            """
@@ -175,7 +175,7 @@ open class BaseTraceElement(val name: String) : CoroutineContext.Element {
        """,
        level = DeprecationLevel.ERROR,
    )
    operator fun plus(other: BaseTraceElement): BaseTraceElement {
    operator fun plus(other: CoroutineTraceName): CoroutineTraceName {
        debug { "#plus(${other.currentId})" }
        return other
    }
@@ -187,26 +187,6 @@ open class BaseTraceElement(val name: String) : CoroutineContext.Element {
    }
}

/**
 * A coroutine context element that can be used for naming the child coroutine under construction.
 *
 * @property name the name to be used for the child under construction
 * @see nameCoroutine
 */
@VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
class CoroutineTraceName(name: String) : BaseTraceElement(name) {
    @OptIn(ExperimentalStdlibApi::class)
    companion object Key :
        AbstractCoroutineContextKey<BaseTraceElement, CoroutineTraceName>(
            BaseTraceElement,
            { it as? CoroutineTraceName },
        )

    init {
        debug { "#init: name=$name" }
    }
}

/**
 * Used for tracking parent-child relationship of coroutines and persisting [TraceData] when
 * coroutines are suspended and resumed.
@@ -231,14 +211,17 @@ class CoroutineTraceName(name: String) : BaseTraceElement(name) {
 *   the child. This should only be used for testing because it can result in extremely long trace
 *   names.
 * @param strictMode whether to add additional checks to coroutine machinery. These checks are
 *   expensive and should only be used for testing. * @see createCoroutineTracingContext
 *   expensive and should only be used for testing.
 * @param shouldIgnoreClassName lambda that takes binary class name (as returned from
 *   [StackFrame.getClassName] and returns true if it should be ignored (e.g. search for relevant
 *   class name should continue) or false otherwise
 * @see createCoroutineTracingContext
 * @see nameCoroutine
 * @see traceCoroutine
 */
@OptIn(DelicateCoroutinesApi::class, ExperimentalCoroutinesApi::class)
@VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
class TraceContextElement
private constructor(
class TraceContextElement(
    name: String,
    @get:VisibleForTesting(otherwise = VisibleForTesting.PACKAGE_PRIVATE)
    val contextTraceData: TraceData?,
@@ -248,37 +231,8 @@ private constructor(
    private val walkStackForDefaultNames: Boolean,
    private val includeParentNames: Boolean,
    private val strictMode: Boolean,
) : CopyableThreadContextElement<TraceData?>, BaseTraceElement(name) {

    @OptIn(ExperimentalStdlibApi::class)
    companion object Key :
        AbstractCoroutineContextKey<BaseTraceElement, TraceContextElement>(
            BaseTraceElement,
            { it as? TraceContextElement },
        ) {
        override fun toString(): String {
            return "TraceContextElement.Key"
        }
    }

    @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
    constructor(
        scopeName: String,
        walkStackForDefaultNames: Boolean = false,
        includeParentNames: Boolean = false,
        strictMode: Boolean = false,
    ) : this(
        name = scopeName,
        // Minor perf optimization: no need to create TraceData() for root scopes since all launches
        // require creation of child via [copyForChild] or [mergeForChild].
        contextTraceData = null,
        inheritedTracePrefix = "",
        coroutineDepth = 0,
        parentId = -1,
        walkStackForDefaultNames = walkStackForDefaultNames,
        includeParentNames = includeParentNames,
        strictMode = strictMode,
    )
    private val shouldIgnoreClassName: (String) -> Boolean,
) : CopyableThreadContextElement<TraceData?>, CoroutineTraceName(name) {

    private var childCoroutineCount = AtomicInteger(0)

@@ -371,7 +325,7 @@ private constructor(
    override fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext {
        debug { "#mergeForChild" }
        if (DEBUG) {
            overwritingElement[TraceContextElement]?.let {
            (overwritingElement as? TraceContextElement)?.let {
                Log.e(
                    TAG,
                    "${this::class.java.simpleName}@$currentId#mergeForChild(@${it.currentId}): " +
@@ -380,15 +334,13 @@ private constructor(
                )
            }
        }
        val nameForChild =
            overwritingElement[CoroutineTraceName]?.name
                ?: overwritingElement[TraceContextElement]?.name
                ?: ""
        val nameForChild = (overwritingElement as CoroutineTraceName).name
        return createChildContext(nameForChild)
    }

    private fun createChildContext(
        name: String = if (walkStackForDefaultNames) walkStackForClassName() else ""
        name: String =
            if (walkStackForDefaultNames) walkStackForClassName(shouldIgnoreClassName) else ""
    ): TraceContextElement {
        debug { "#createChildContext: \"$name\" has new child with name \"${name}\"" }
        val childCount = childCoroutineCount.incrementAndGet()
@@ -402,26 +354,31 @@ private constructor(
            walkStackForDefaultNames = walkStackForDefaultNames,
            includeParentNames = includeParentNames,
            strictMode = strictMode,
            shouldIgnoreClassName = shouldIgnoreClassName,
        )
    }
}

/** Get a name for the trace section include the name of the call site. */
private fun walkStackForClassName(): String {
/**
 * Get a name for the trace section include the name of the call site.
 *
 * @param additionalDropPredicate additional checks for whether class should be ignored
 */
private fun walkStackForClassName(
    additionalDropPredicate: (String) -> Boolean = { false }
): String {
    Trace.traceBegin(Trace.TRACE_TAG_APP, "walkStackForClassName")
    try {
        var frame = ""
        StackWalker.getInstance().walk { s: Stream<StackFrame> ->
            s.dropWhile { f: StackFrame ->
                    with(f.className) {
                        startsWith("kotlin") ||
                            startsWith("android") ||
                            startsWith("com.android.internal.") ||
                            startsWith("com.android.app.tracing.")
                    }
                    val className = f.className
                    className.startsWith("kotlin") ||
                        className.startsWith("com.android.app.tracing.") ||
                        additionalDropPredicate(className)
                }
                .findFirst()
                .ifPresent { frame = it.className.substringAfterLast(".") }
                .ifPresent { frame = it.className.substringAfterLast(".") + "." + it.methodName }
        }
        return frame
    } catch (e: Exception) {
@@ -435,6 +392,5 @@ private fun walkStackForClassName(): String {
private const val UNEXPECTED_TRACE_DATA_ERROR_MESSAGE =
    "Overwriting context element with non-empty trace data. There should only be one " +
        "TraceContextElement per coroutine, and it should be installed in the root scope. "

internal val TAG = "CoroutineTracing"
internal const val TAG = "CoroutineTracing"
internal const val DEBUG = false
+146 −84

File changed.

Preview size limit exceeded, changes collapsed.

+12 −3
Original line number Diff line number Diff line
@@ -52,6 +52,8 @@ import kotlinx.coroutines.android.asCoroutineDispatcher

@Qualifier @MustBeDocumented @Retention(RUNTIME) annotation class FixedThreadC

@Qualifier @MustBeDocumented @Retention(RUNTIME) annotation class FixedThreadD

@Module
class ConcurrencyModule {

@@ -80,21 +82,28 @@ class ConcurrencyModule {
    @Singleton
    @FixedThreadA
    fun provideDispatcherA(): CoroutineDispatcher {
        return startThreadWithLooper("Thread A").threadHandler.asCoroutineDispatcher()
        return startThreadWithLooper("Thread:A").threadHandler.asCoroutineDispatcher()
    }

    @Provides
    @Singleton
    @FixedThreadB
    fun provideDispatcherB(): CoroutineDispatcher {
        return startThreadWithLooper("Thread B").threadHandler.asCoroutineDispatcher()
        return startThreadWithLooper("Thread:B").threadHandler.asCoroutineDispatcher()
    }

    @Provides
    @Singleton
    @FixedThreadC
    fun provideDispatcherC(): CoroutineDispatcher {
        return startThreadWithLooper("Thread C").threadHandler.asCoroutineDispatcher()
        return startThreadWithLooper("Thread:C").threadHandler.asCoroutineDispatcher()
    }

    @Provides
    @Singleton
    @FixedThreadD
    fun provideDispatcherD(): CoroutineDispatcher {
        return startThreadWithLooper("Thread:D").threadHandler.asCoroutineDispatcher()
    }
}

Loading