Merge "New tracinglib library for coroutine tracing" into main

// Copyright (C) 2023 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 {

    default_applicable_licenses: ["Android-Apache-2.0"],

}

java_library {

    name: "tracinglib",

    static_libs: [

        "kotlinx_coroutines_android",

    ],

    srcs: [

        "src/**/*.kt"

    ],

    kotlincflags: ["-Xjvm-default=all"],

    platform_apis: true,

}

/*

 * Copyright (C) 2023 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 com.android.app.tracing

import com.android.app.tracing.TraceUtils.Companion.instant

import com.android.app.tracing.TraceUtils.Companion.traceCoroutine

import kotlin.coroutines.CoroutineContext

import kotlinx.coroutines.CopyableThreadContextElement

import kotlinx.coroutines.CoroutineDispatcher

import kotlinx.coroutines.DelicateCoroutinesApi

import kotlinx.coroutines.ExperimentalCoroutinesApi

/**

 * Used for safely persisting [TraceData] state when coroutines are suspended and resumed.

 *

 * This is internal machinery for [traceCoroutine]. It cannot be made `internal` or `private`

 * because [traceCoroutine] is a Public-API inline function.

 *

 * @see traceCoroutine

 */

@OptIn(DelicateCoroutinesApi::class)

@ExperimentalCoroutinesApi

class TraceContextElement(private val traceData: TraceData = TraceData()) :

    CopyableThreadContextElement<TraceData?> {

    companion object Key : CoroutineContext.Key<TraceContextElement>

    override val key: CoroutineContext.Key<TraceContextElement> = Key

    @OptIn(ExperimentalStdlibApi::class)

    override fun updateThreadContext(context: CoroutineContext): TraceData? {

        val oldState = threadLocalTrace.get()

        oldState?.endAllOnThread()

        threadLocalTrace.set(traceData)

        instant("resuming ${context[CoroutineDispatcher]}")

        traceData.beginAllOnThread()

        return oldState

    }

    @OptIn(ExperimentalStdlibApi::class)

    override fun restoreThreadContext(context: CoroutineContext, oldState: TraceData?) {

        instant("suspending ${context[CoroutineDispatcher]}")

        traceData.endAllOnThread()

        threadLocalTrace.set(oldState)

        oldState?.beginAllOnThread()

    }

    override fun copyForChild(): CopyableThreadContextElement<TraceData?> {

        return TraceContextElement(traceData.copy())

    }

    override fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext {

        return TraceContextElement(traceData.copy())

    }

}

/*

 * Copyright (C) 2023 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 com.android.app.tracing

import android.os.Build

import android.util.Log

import com.android.app.tracing.TraceUtils.Companion.beginSlice

import com.android.app.tracing.TraceUtils.Companion.endSlice

import com.android.app.tracing.TraceUtils.Companion.traceCoroutine

import kotlin.random.Random

/**

 * Used for giving each thread a unique [TraceData] for thread-local storage. `null` by default.

 * [threadLocalTrace] can only be used when it is paired with a [TraceContextElement].

 *

 * This ThreadLocal will be `null` if either 1) we aren't in a coroutine, or 2) the coroutine we are

 * in does not have a [TraceContextElement].

 *

 * This is internal machinery for [traceCoroutine]. It cannot be made `internal` or `private`

 * because [traceCoroutine] is a Public-API inline function.

 *

 * @see traceCoroutine

 */

val threadLocalTrace = ThreadLocal<TraceData?>()

/**

 * Used for storing trace sections so that they can be added and removed from the currently running

 * thread when the coroutine is suspended and resumed.

 *

 * This is internal machinery for [traceCoroutine]. It cannot be made `internal` or `private`

 * because [traceCoroutine] is a Public-API inline function.

 *

 * @see traceCoroutine

 */

class TraceData {

    private var slices = mutableListOf<TraceSection>()

    /** Adds current trace slices back to the current thread. Called when coroutine is resumed. */

    fun beginAllOnThread() {

        slices.forEach { beginSlice(it.name) }

    }

    /**

     * Removes all current trace slices from the current thread. Called when coroutine is suspended.

     */

    fun endAllOnThread() {

        for (i in 0..slices.size) {

            endSlice()

        }

    }

    /**

     * Creates a new trace section with a unique ID and adds it to the current trace data. The slice

     * will also be added to the current thread immediately. This slice will not propagate to parent

     * coroutines, or to child coroutines that have already started. The unique ID is used to verify

     * that the [endSpan] is corresponds to a [beginSpan].

     */

    fun beginSpan(name: String): Int {

        val newSlice = TraceSection(name, Random.nextInt(FIRST_VALID_SPAN, Int.MAX_VALUE))

        slices.add(newSlice)

        beginSlice(name)

        return newSlice.id

    }

    /**

     * Used by [TraceContextElement] when launching a child coroutine so that the child coroutine's

     * state is isolated from the parent.

     */

    fun copy(): TraceData {

        return TraceData().also { it.slices.addAll(slices) }

    }

    /**

     * Ends the trace section and validates it corresponds with an earlier call to [beginSpan]. The

     * trace slice will immediately be removed from the current thread. This information will not

     * propagate to parent coroutines, or to child coroutines that have already started.

     */

    fun endSpan(id: Int) {

        val v = slices.removeLast()

        if (v.id != id) {

            if (STRICT_MODE) {

                throw IllegalArgumentException(errorMsg)

            } else if (Log.isLoggable(TAG, Log.VERBOSE)) {

                Log.v(TAG, errorMsg)

            }

        }

        endSlice()

    }

    companion object {

        private const val TAG = "TraceData"

        const val INVALID_SPAN = -1

        const val FIRST_VALID_SPAN = 1

        /**

         * If true, throw an exception instead of printing a warning when trace sections beginnings

         * and ends are mismatched.

         */

        private val STRICT_MODE = Build.IS_ENG

        private const val errorMsg =

            "Mismatched trace section. This likely means you are accessing the trace local " +

                "storage (threadLocalTrace) without a corresponding CopyableThreadContextElement." +

                " This could happen if you are using a global dispatcher like Dispatchers.IO." +

                " To fix this, use one of the coroutine contexts provided by the dagger scope " +

                "(e.g. \"@Main CoroutineContext\")."

    }

}

/*

 * Copyright (C) 2023 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 com.android.app.tracing

import com.android.app.tracing.TraceUtils.Companion.traceCoroutine

/**

 * Represents a section of code executing in a coroutine. This can be split up into multiple slices

 * on different threads as the coroutine is suspended and resumed.

 *

 * This is internal machinery for [traceCoroutine]. It cannot be made `internal` or `private`

 * because [traceCoroutine] is a Public-API inline function.

 *

 * @param name the name of the slice to appear on the current thread's track.

 * @param id used for matching the beginning and end of trace sections and validating correctness

 * @see traceCoroutine

 */

data class TraceSection(

    val name: String,

    val id: Int,

)

/*

 * Copyright (C) 2023 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 com.android.app.tracing

import android.os.Trace

/**

 * Utility class used to log state changes easily in a track with a custom name.

 *

 * Example of usage:

 * ```kotlin

 * class MyClass {

 *    val screenStateLogger = TraceStateLogger("Screen state")

 *

 *    fun onTurnedOn() { screenStateLogger.log("on") }

 *    fun onTurnedOff() { screenStateLogger.log("off") }

 * }

 * ```

 *

 * This creates a new slice in a perfetto trace only if the state is different than the previous

 * one.

 */

class TraceStateLogger(

    private val trackName: String,

    private val logOnlyIfDifferent: Boolean = true,

    private val instantEvent: Boolean = true

) {

    private var previousValue: String? = null

    /** If needed, logs the value to a track with name [trackName]. */

    fun log(newValue: String) {

        if (instantEvent) {

            Trace.instantForTrack(Trace.TRACE_TAG_APP, trackName, newValue)

        }

        if (logOnlyIfDifferent && previousValue == newValue) return

        Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, trackName, 0)

        Trace.asyncTraceForTrackBegin(Trace.TRACE_TAG_APP, trackName, newValue, 0)

        previousValue = newValue

    }

}