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

Commit ed3da0e4 authored by Treehugger Robot's avatar Treehugger Robot Committed by Android (Google) Code Review
Browse files

Merge "tracinglib: delegate traceAs calls" into main

parents a471aaff 44bc7c80
Loading
Loading
Loading
Loading
+16 −0
Original line number Diff line number Diff line
@@ -125,6 +125,22 @@ public inline fun <T> traceSection(tag: () -> String, block: () -> T): T {
    }
}

/**
 * Like [com.android.app.tracing.traceSection], but uses `crossinline` so we don't accidentally
 * introduce non-local returns. This is less convenient to use, but it ensures we will not
 * accidentally pass a suspending function to this method.
 */
@OptIn(ExperimentalContracts::class)
internal inline fun <T> traceBlocking(sectionName: String, crossinline block: () -> T): T {
    contract { callsInPlace(block, InvocationKind.EXACTLY_ONCE) }
    Trace.traceBegin(Trace.TRACE_TAG_APP, sectionName)
    return try {
        block()
    } finally {
        Trace.traceEnd(Trace.TRACE_TAG_APP)
    }
}

@OptIn(ExperimentalContracts::class)
public object TraceUtils {
    public const val TAG: String = "TraceUtils"
+70 −55
Original line number Diff line number Diff line
@@ -19,7 +19,7 @@ package com.android.app.tracing.coroutines.flow
import com.android.app.tracing.coroutines.CoroutineTraceName
import com.android.app.tracing.coroutines.traceCoroutine
import com.android.app.tracing.coroutines.traceName
import com.android.app.tracing.traceSection
import com.android.app.tracing.traceBlocking
import kotlin.experimental.ExperimentalTypeInference
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -65,10 +65,12 @@ internal inline fun <T, R> Flow<T>.unsafeTransform(
): Flow<R> = unsafeFlow { collect { value -> transform(value) } }

@OptIn(ExperimentalForInheritanceCoroutinesApi::class)
private open class TracedSharedFlow<T>(private val name: String, private val flow: SharedFlow<T>) :
    SharedFlow<T> by flow {
private open class TracedSharedFlow<out T>(
    private val name: String,
    private val flow: SharedFlow<T>,
) : SharedFlow<T> {
    override val replayCache: List<T>
        get() = traceSection("replayCache:$name") { flow.replayCache }
        get() = traceBlocking("replayCache:$name") { flow.replayCache }

    override suspend fun collect(collector: FlowCollector<T>): Nothing {
        traceCoroutine("collect:$name") {
@@ -78,38 +80,33 @@ private open class TracedSharedFlow<T>(private val name: String, private val flo
}

@OptIn(ExperimentalForInheritanceCoroutinesApi::class)
private open class TracedStateFlow<T>(private val name: String, private val flow: StateFlow<T>) :
    StateFlow<T> by flow, TracedSharedFlow<T>(name, flow) {
private open class TracedStateFlow<out T>(
    private val name: String,
    private val flow: StateFlow<T>,
) : StateFlow<T>, TracedSharedFlow<T>(name, flow) {
    override val value: T
        get() = traceCoroutine("get:$name") { flow.value }

    override val replayCache: List<T>
        get() = super.replayCache

    override suspend fun collect(collector: FlowCollector<T>): Nothing {
        super.collect(collector)
    }
        get() = traceBlocking("get:$name") { flow.value }
}

@OptIn(ExperimentalForInheritanceCoroutinesApi::class)
private class TracedMutableSharedFlow<T>(
private open class TracedMutableSharedFlow<T>(
    private val name: String,
    private val flow: MutableSharedFlow<T>,
) : MutableSharedFlow<T> by flow, TracedSharedFlow<T>(name, flow) {
) : MutableSharedFlow<T>, TracedSharedFlow<T>(name, flow) {
    override val subscriptionCount: StateFlow<Int>
        get() = traceBlocking("subscriptionCount:$name") { flow.subscriptionCount }

    override val replayCache: List<T>
        get() = super.replayCache
    @ExperimentalCoroutinesApi
    override fun resetReplayCache() {
        traceBlocking("resetReplayCache:$name") { flow.resetReplayCache() }
    }

    override suspend fun emit(value: T) {
        traceCoroutine("emit:$name") { flow.emit(value) }
    }

    override fun tryEmit(value: T): Boolean {
        return traceCoroutine("tryEmit:$name") { flow.tryEmit(value) }
    }

    override suspend fun collect(collector: FlowCollector<T>): Nothing {
        super.collect(collector)
        return traceBlocking("tryEmit:$name") { flow.tryEmit(value) }
    }
}

@@ -117,31 +114,15 @@ private class TracedMutableSharedFlow<T>(
private class TracedMutableStateFlow<T>(
    private val name: String,
    private val flow: MutableStateFlow<T>,
) : MutableStateFlow<T> by flow, TracedStateFlow<T>(name, flow) {

) : MutableStateFlow<T>, TracedMutableSharedFlow<T>(name, flow) {
    override var value: T
        get() = super.value
        get() = traceBlocking("get:$name") { flow.value }
        set(newValue) {
            traceSection("updateState:$name") { flow.value = newValue }
        }

    override val replayCache: List<T>
        get() = super.replayCache

    override suspend fun emit(value: T) {
        traceCoroutine("emit:$name") { flow.emit(value) }
    }

    override suspend fun collect(collector: FlowCollector<T>): Nothing {
        traceCoroutine("collect:$name") {
            flow.collect { traceCoroutine("emit:$name") { collector.emit(it) } }
        }
            traceBlocking("updateState:$name") { flow.value = newValue }
        }

    override fun compareAndSet(expect: T, update: T): Boolean {
        traceSection("updateState:$name") {
            return flow.compareAndSet(expect, update)
        }
        return traceBlocking("compareAndSet:$name") { flow.compareAndSet(expect, update) }
    }
}

@@ -172,27 +153,61 @@ public fun <T> Flow<T>.flowName(name: String): Flow<T> = traceAs(name)

public fun <T> Flow<T>.traceAs(name: String): Flow<T> {
    return if (com.android.systemui.Flags.coroutineTracing()) {
        return when (this) {
            is SharedFlow -> traceAs(name)
            else ->
                unsafeFlow {
                    traceCoroutine("collect:$name") {
                        collect { value -> traceCoroutine("emit:$name") { emit(value) } }
                    }
                }
        }
    } else {
        this
    }
}

public fun <T> SharedFlow<T>.traceAs(name: String): SharedFlow<T> =
    if (com.android.systemui.Flags.coroutineTracing()) TracedSharedFlow(name, this) else this
public fun <T> SharedFlow<T>.traceAs(name: String): SharedFlow<T> {
    return if (com.android.systemui.Flags.coroutineTracing()) {
        when (this) {
            is MutableSharedFlow -> traceAs(name)
            is StateFlow -> traceAs(name)
            else -> TracedSharedFlow(name, this)
        }
    } else {
        this
    }
}

public fun <T> StateFlow<T>.traceAs(name: String): StateFlow<T> =
    if (com.android.systemui.Flags.coroutineTracing()) TracedStateFlow(name, this) else this
public fun <T> StateFlow<T>.traceAs(name: String): StateFlow<T> {
    return if (com.android.systemui.Flags.coroutineTracing()) {
        when (this) {
            is MutableStateFlow -> traceAs(name)
            else -> TracedStateFlow(name, this)
        }
    } else {
        this
    }
}

public fun <T> MutableSharedFlow<T>.traceAs(name: String): MutableSharedFlow<T> =
    if (com.android.systemui.Flags.coroutineTracing()) TracedMutableSharedFlow(name, this) else this
public fun <T> MutableSharedFlow<T>.traceAs(name: String): MutableSharedFlow<T> {
    return if (com.android.systemui.Flags.coroutineTracing()) {
        when (this) {
            is MutableStateFlow -> traceAs(name)
            else -> TracedMutableSharedFlow(name, this)
        }
    } else {
        this
    }
}

public fun <T> MutableStateFlow<T>.traceAs(name: String): MutableStateFlow<T> =
    if (com.android.systemui.Flags.coroutineTracing()) TracedMutableStateFlow(name, this) else this
public fun <T> MutableStateFlow<T>.traceAs(name: String): MutableStateFlow<T> {
    return if (com.android.systemui.Flags.coroutineTracing()) {
        TracedMutableStateFlow(name, this)
    } else {
        this
    }
}

public fun <T> Flow<T>.onEachTraced(name: String, action: suspend (T) -> Unit): Flow<T> {
    return onEach { value -> traceCoroutine(name) { action(value) } }