1 /*
2 * Copyright (C) 2023 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 package com.android.app.tracing.coroutines
18
19 import com.android.systemui.Flags
20 import kotlin.contracts.ExperimentalContracts
21 import kotlin.contracts.InvocationKind
22 import kotlin.contracts.contract
23 import kotlin.coroutines.CoroutineContext
24 import kotlin.coroutines.EmptyCoroutineContext
25 import kotlinx.coroutines.CoroutineScope
26 import kotlinx.coroutines.CoroutineStart
27 import kotlinx.coroutines.Deferred
28 import kotlinx.coroutines.Job
29 import kotlinx.coroutines.async
30 import kotlinx.coroutines.coroutineScope
31 import kotlinx.coroutines.launch
32 import kotlinx.coroutines.runBlocking
33 import kotlinx.coroutines.withContext
34
35 @OptIn(ExperimentalContracts::class)
coroutineScopeTracednull36 public suspend inline fun <R> coroutineScopeTraced(
37 traceName: String,
38 crossinline block: suspend CoroutineScope.() -> R,
39 ): R {
40 contract { callsInPlace(block, InvocationKind.EXACTLY_ONCE) }
41 return coroutineScope {
42 traceCoroutine(traceName) {
43 return@coroutineScope block()
44 }
45 }
46 }
47
48 /**
49 * Convenience function for calling [CoroutineScope.launch] with [traceCoroutine] to enable tracing.
50 *
51 * @see traceCoroutine
52 */
launchTracednull53 public inline fun CoroutineScope.launchTraced(
54 crossinline spanName: () -> String,
55 context: CoroutineContext = EmptyCoroutineContext,
56 start: CoroutineStart = CoroutineStart.DEFAULT,
57 noinline block: suspend CoroutineScope.() -> Unit,
58 ): Job {
59 return launch(nameCoroutine(spanName) + context, start, block)
60 }
61
62 /**
63 * Convenience function for calling [CoroutineScope.launch] with [traceCoroutine] to enable tracing.
64 *
65 * @see traceCoroutine
66 */
launchTracednull67 public fun CoroutineScope.launchTraced(
68 spanName: String? = null,
69 context: CoroutineContext = EmptyCoroutineContext,
70 start: CoroutineStart = CoroutineStart.DEFAULT,
71 block: suspend CoroutineScope.() -> Unit,
72 ): Job = launchTraced({ spanName ?: block::class.simpleName ?: "launch" }, context, start, block)
73
74 /**
75 * Convenience function for calling [CoroutineScope.async] with [traceCoroutine] enable tracing
76 *
77 * @see traceCoroutine
78 */
asyncTracednull79 public inline fun <T> CoroutineScope.asyncTraced(
80 spanName: () -> String,
81 context: CoroutineContext = EmptyCoroutineContext,
82 start: CoroutineStart = CoroutineStart.DEFAULT,
83 noinline block: suspend CoroutineScope.() -> T,
84 ): Deferred<T> = async(nameCoroutine(spanName) + context, start, block)
85
86 /**
87 * Convenience function for calling [CoroutineScope.async] with [traceCoroutine] enable tracing.
88 *
89 * @see traceCoroutine
90 */
91 public fun <T> CoroutineScope.asyncTraced(
92 spanName: String? = null,
93 context: CoroutineContext = EmptyCoroutineContext,
94 start: CoroutineStart = CoroutineStart.DEFAULT,
95 block: suspend CoroutineScope.() -> T,
96 ): Deferred<T> =
97 asyncTraced({ spanName ?: block::class.simpleName ?: "async" }, context, start, block)
98
99 /**
100 * Convenience function for calling [runBlocking] with [traceCoroutine] to enable tracing.
101 *
102 * @see traceCoroutine
103 */
runBlockingTracednull104 public inline fun <T> runBlockingTraced(
105 spanName: () -> String,
106 context: CoroutineContext,
107 noinline block: suspend CoroutineScope.() -> T,
108 ): T = runBlocking(nameCoroutine(spanName) + context, block)
109
110 /**
111 * Convenience function for calling [runBlocking] with [traceCoroutine] to enable tracing.
112 *
113 * @see traceCoroutine
114 */
115 public fun <T> runBlockingTraced(
116 spanName: String? = null,
117 context: CoroutineContext,
118 block: suspend CoroutineScope.() -> T,
119 ): T = runBlockingTraced({ spanName ?: block::class.simpleName ?: "runBlocking" }, context, block)
120
121 /**
122 * Convenience function for calling [withContext] with [traceCoroutine] to enable tracing.
123 *
124 * @see traceCoroutine
125 */
withContextTracednull126 public suspend fun <T> withContextTraced(
127 spanName: String? = null,
128 context: CoroutineContext,
129 block: suspend CoroutineScope.() -> T,
130 ): T = withContextTraced({ spanName ?: block::class.simpleName ?: "withContext" }, context, block)
131
132 /**
133 * Convenience function for calling [withContext] with [traceCoroutine] to enable tracing.
134 *
135 * @see traceCoroutine
136 */
withContextTracednull137 public suspend inline fun <T> withContextTraced(
138 spanName: () -> String,
139 context: CoroutineContext,
140 noinline block: suspend CoroutineScope.() -> T,
141 ): T = withContext(nameCoroutine(spanName) + context, block)
142
143 /**
144 * Traces a section of work of a `suspend` [block]. The trace sections will appear on the thread
145 * that is currently executing the [block] of work. If the [block] is suspended, all trace sections
146 * added using this API will end until the [block] is resumed, which could happen either on this
147 * thread or on another thread. If a child coroutine is started, it will *NOT* inherit the trace
148 * sections of its parent; however, it will include metadata in the trace section pointing to the
149 * parent.
150 *
151 * The current [CoroutineContext] must have a [TraceContextElement] for this API to work. Otherwise,
152 * the trace sections will be dropped.
153 *
154 * For example, in the following trace, Thread #1 starts a coroutine, suspends, and continues the
155 * coroutine on Thread #2. Next, Thread #2 start a child coroutine in an unconfined manner. Then,
156 * still on Thread #2, the original coroutine suspends, the child resumes, and the child suspends.
157 * Then, the original coroutine resumes on Thread#1.
158 *
159 * ```
160 * -----------------------------------------------------------------------------------------------|
161 * Thread #1 | [== Slice A ==] [==== Slice A ====]
162 * | [== B ==] [=== B ===]
163 * -----------------------------------------------------------------------------------------------|
164 * Thread #2 | [==== Slice A ====] [=== C ====]
165 * | [======= B =======]
166 * | [=== C ====]
167 * -----------------------------------------------------------------------------------------------|
168 * ```
169 *
170 * @param spanName The name of the code section to appear in the trace
171 * @see traceCoroutine
172 */
173 @OptIn(ExperimentalContracts::class)
174 public inline fun <T> traceCoroutine(spanName: () -> String, block: () -> T): T {
175 contract {
176 callsInPlace(spanName, InvocationKind.AT_MOST_ONCE)
177 callsInPlace(block, InvocationKind.EXACTLY_ONCE)
178 }
179
180 // For coroutine tracing to work, trace spans must be added and removed even when
181 // tracing is not active (i.e. when TRACE_TAG_APP is disabled). Otherwise, when the
182 // coroutine resumes when tracing is active, we won't know its name.
183 val traceData = if (Flags.coroutineTracing()) traceThreadLocal.get() else null
184 traceData?.beginSpan(spanName())
185 try {
186 return block()
187 } finally {
188 traceData?.endSpan()
189 }
190 }
191
192 /** @see traceCoroutine */
traceCoroutinenull193 public inline fun <T> traceCoroutine(spanName: String, block: () -> T): T =
194 traceCoroutine({ spanName }, block)
195