Measure time from composition start to the first draw in Jetpack Compose

UntilFirstDrawTracer.kt

import android.os.Handler
import android.os.Looper
import android.view.ViewTreeObserver
import android.view.ViewTreeObserver.OnDrawListener
import androidx.compose.runtime.Composable
import androidx.compose.runtime.DisposableEffect
import androidx.compose.runtime.SideEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import androidx.compose.ui.platform.LocalView
import kotlin.time.TimeSource

/**
 * Tracks and measures the time from the start of the composition to the first onDraw event.
 * This is useful for performance tracking, as it calculates the elapsed time
 * it takes for a Compose screen to reach its first visible frame.
 *
 * Example Usage:
 *
 * @Composable
 * fun TopicScreen() {
 *   UntilFirstDrawTracer { time ->
 *     Tracer.trace("TopicScreen", time)
 *   }
 *
 *   // UI content for the TopicScreen goes here
 * }
 *
 * @param onMeasured - Callback that receives the elapsed time in milliseconds
 *                     from composition start to the first onDraw event.
 */
@Composable
fun UntilFirstDrawTracer(onMeasured: (Long) -> Unit) {
    val startTime = remember { TimeSource.Monotonic.markNow() }

    val view = LocalView.current
    val viewTreeObserver = view.viewTreeObserver

    DisposableEffect(viewTreeObserver) {
        val listener = OnFirstDrawListener(
            viewTreeObserver = viewTreeObserver,
            onFinish = {
                onMeasured(startTime.elapsedNow().inWholeMilliseconds)
            },
        )

        viewTreeObserver.addOnDrawListener(listener)
        onDispose {
            viewTreeObserver.removeOnDrawListener(listener)
        }
    }
}

private class OnFirstDrawListener(
    private val viewTreeObserver: ViewTreeObserver,
    private val onFinish: () -> Unit,
) : OnDrawListener {
    private var firstOnDrawHappened: Boolean = false

    override fun onDraw() {
        if (!firstOnDrawHappened) {
            firstOnDrawHappened = true
            onFinish()
            Handler(Looper.getMainLooper()).post {
                dispose()
            }
        }
    }

    fun dispose() {
        if (viewTreeObserver.isAlive) {
            viewTreeObserver.removeOnDrawListener(this)
        }
    }
}