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 @file:OptIn(ExperimentalTypeInference::class)
17 
18 package com.android.app.tracing
19 
20 import android.os.Trace
21 import com.android.app.tracing.TraceUtils.traceAsync
22 import java.util.concurrent.atomic.AtomicInteger
23 import kotlin.experimental.ExperimentalTypeInference
24 import kotlinx.coroutines.channels.ProducerScope
25 import kotlinx.coroutines.channels.awaitClose
26 import kotlinx.coroutines.flow.Flow
27 import kotlinx.coroutines.flow.callbackFlow
28 import kotlinx.coroutines.flow.conflate
29 import kotlinx.coroutines.flow.onEach
30 
31 /** Utilities to trace Flows */
32 public object FlowTracing {
33 
34     private const val TAG = "FlowTracing"
35     private const val DEFAULT_ASYNC_TRACK_NAME = TAG
36     private val counter = AtomicInteger(0)
37 
38     /** Logs each flow element to a trace. */
traceEachnull39     public inline fun <T> Flow<T>.traceEach(
40         flowName: String,
41         logcat: Boolean = false,
42         traceEmissionCount: Boolean = false,
43         crossinline valueToString: (T) -> String = { it.toString() },
44     ): Flow<T> {
45         val stateLogger = TraceStateLogger(flowName, logcat = logcat)
46         val baseFlow = if (traceEmissionCount) traceEmissionCount(flowName) else this
<lambda>null47         return baseFlow.onEach { stateLogger.log(valueToString(it)) }
48     }
49 
50     /** Records value of a given numeric flow as a counter track in traces. */
traceAsCounternull51     public fun <T : Number> Flow<T>.traceAsCounter(
52         counterName: String,
53         traceEmissionCount: Boolean = false,
54         valueToInt: (T) -> Int = { it.toInt() },
55     ): Flow<T> {
56         val baseFlow = if (traceEmissionCount) traceEmissionCount(counterName) else this
<lambda>null57         return baseFlow.onEach {
58             if (Trace.isEnabled()) {
59                 Trace.traceCounter(Trace.TRACE_TAG_APP, counterName, valueToInt(it))
60             }
61         }
62     }
63 
64     /** Adds a counter track to monitor emissions from a specific flow.] */
traceEmissionCountnull65     public fun <T> Flow<T>.traceEmissionCount(
66         flowName: String,
67         uniqueSuffix: Boolean = false,
68     ): Flow<T> {
69         val trackName by lazy {
70             "$flowName#emissionCount" + if (uniqueSuffix) "\$${counter.addAndGet(1)}" else ""
71         }
72         var count = 0
73         return onEach {
74             count += 1
75             Trace.traceCounter(Trace.TRACE_TAG_APP, trackName, count)
76         }
77     }
78 
79     /**
80      * Adds a counter track to monitor emissions from a specific flow.
81      *
82      * [flowName] is lazy: it would be computed only if tracing is enabled and only the first time.
83      */
traceEmissionCountnull84     public fun <T> Flow<T>.traceEmissionCount(
85         flowName: () -> String,
86         uniqueSuffix: Boolean = false,
87     ): Flow<T> {
88         val trackName by lazy {
89             "${flowName()}#emissionCount" + if (uniqueSuffix) "\$${counter.addAndGet(1)}" else ""
90         }
91         var count = 0
92         return onEach {
93             count += 1
94             if (Trace.isEnabled()) {
95                 Trace.traceCounter(Trace.TRACE_TAG_APP, trackName, count)
96             }
97         }
98     }
99 
100     /**
101      * Makes [awaitClose] output Perfetto traces.
102      *
103      * There will be 2 traces:
104      * - One in the thread this is being executed on
105      * - One in a track having [DEFAULT_ASYNC_TRACK_NAME] name.
106      *
107      * This allows to easily have visibility into what's happening in awaitClose.
108      */
<lambda>null109     public suspend fun ProducerScope<*>.tracedAwaitClose(name: String, block: () -> Unit = {}) {
<lambda>null110         awaitClose {
111             val traceName = { "$name#TracedAwaitClose" }
112             traceAsync(DEFAULT_ASYNC_TRACK_NAME, traceName) { traceSection(traceName) { block() } }
113         }
114     }
115 
116     /**
117      * Traced version of [callbackFlow].
118      *
119      * Adds tracing in 2 ways:
120      * - An async slice will appear in the [DEFAULT_ASYNC_TRACK_NAME] named track.
121      * - A counter will be increased at every emission
122      *
123      * Should be used with [tracedAwaitClose] (when needed).
124      */
tracedConflatedCallbackFlownull125     public fun <T> tracedConflatedCallbackFlow(
126         name: String,
127         @BuilderInference block: suspend ProducerScope<T>.() -> Unit,
128     ): Flow<T> {
129         return callbackFlow {
130                 traceAsync(DEFAULT_ASYNC_TRACK_NAME, { "$name#CallbackFlowBlock" }) {
131                     block(this@callbackFlow)
132                 }
133             }
134             .conflate()
135             .traceEmissionCount(name, uniqueSuffix = true)
136     }
137 }
138