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

UntilLastDrawTracer.kt

package ru.drinkit.benchmarkutils.compose

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.remember
import androidx.compose.ui.platform.LocalView
import kotlinx.datetime.Clock

private const val TIMEOUT = 1000L

/**
 * Tracks and measures the time from the start of the composition to the last onDraw event,
 * providing a single [onMeasured] callback with the total elapsed time.
 *
 * This is useful for capturing the time it takes for the screen to settle fully
 * after all initial rendering and layout operations have completed.
 *
 * Example Usage:
 *
 * @Composable
 * fun MyScreen() {
 *   UntilLastDrawTracer { time ->
 *     Tracer.trace("MyScreen", time)
 *   }
 *
 *   // UI content for MyScreen goes here
 * }
 *
 * @param onMeasured - Callback that receives the elapsed time in milliseconds
 *                     from the start of composition to the last onDraw event.
 */
@Composable
fun UntilLastDrawTracer(onMeasured: (Long) -> Unit) {
    val startTime = remember { System.currentTimeMillis() }

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

    DisposableEffect(viewTreeObserver) {
        val listener = OnLastDrawListener(
            viewTreeObserver = viewTreeObserver,
            startTime = startTime,
            onMeasured = onMeasured
        )

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

private class OnLastDrawListener(
    private val viewTreeObserver: ViewTreeObserver,
    private val startTime: Long,
    private val onMeasured: (Long) -> Unit,
) : OnDrawListener {

    val handler = Handler(Looper.getMainLooper())
    private var lastTime: Long = 0L

    override fun onDraw() {
        val checkTime = Clock.System.now().toEpochMilliseconds()
        lastTime = checkTime
        handler.postDelayed(
            {
                if (checkTime == lastTime) {
                    dispose()
                    onMeasured(lastTime - startTime)
                }
            },
            TIMEOUT,
        )

    }

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