Measure time between onCreate call and the first onDraw call

UntilFirstDrawTracer.kt

import android.app.Activity
import android.os.Handler
import android.os.Looper
import android.os.SystemClock
import android.view.ViewTreeObserver
import android.view.ViewTreeObserver.OnDrawListener
import androidx.fragment.app.Fragment
import androidx.lifecycle.DefaultLifecycleObserver
import androidx.lifecycle.LifecycleOwner

/**
 * Trace method that gives you [onStart] and [onFinish] callbacks 
 * representing onCreate and first onDraw call
 * 
 * Example Usage (in Activity or Fragment):
 *
 * class MainActivity : AppCompatActivity() {
 *     override fun onCreate(savedInstanceState: Bundle?) {
 *         super.onCreate(savedInstanceState)
 *
 *         traceUntilFirstDraw(
 *             onStart = { Tracer.start("MyScreen") },
 *             onFinish = { Tracer.end("MyScreen") }
 *         )
 *     }
 * }
 */
fun LifecycleOwner.traceUntilFirstDraw(
  onStart: () -> Unit,
  onFinish: () -> Unit,
) {
  UntilFirstDrawTracer(
      lifecycleOwner = this,
      onStart = onStart,
      onFinish = onFinish,
  )
      .setup()
}

class UntilFirstDrawTracer(
  private val lifecycleOwner: LifecycleOwner,
  private val onStart: () -> Unit,
  private val onFinish: () -> Unit,
) {
  inner class FirstDrawObserver : DefaultLifecycleObserver {
    private var startTime: Long = 0L
    private var onFirstDrawListener: OnFirstDrawListener? = null

    override fun onCreate(owner: LifecycleOwner) {
      startTime = SystemClock.elapsedRealtime()
      onStart()
    }

    override fun onStart(owner: LifecycleOwner) {
      val viewTreeObserver: ViewTreeObserver? = when (owner) {
        is Fragment -> owner.view?.viewTreeObserver
        is Activity -> owner.window.decorView.viewTreeObserver
        else -> null
      }

      viewTreeObserver?.let { nonNullViewTreeObserver ->
        onFirstDrawListener = OnFirstDrawListener(
            viewTreeObserver = nonNullViewTreeObserver,
            startTime = startTime,
            onFinish = onFinish,
        )
        nonNullViewTreeObserver.addOnDrawListener(onFirstDrawListener)
      }
    }

    override fun onStop(owner: LifecycleOwner) {
      onFirstDrawListener?.dispose()
    }
  }

  fun setup() {
    lifecycleOwner.lifecycle.addObserver(FirstDrawObserver())
  }
}

/**
 * Trace method that calculates time between onCreate and the first onDraw 
 * and return one [onMeasured] callback with calculated time
 * 
 * Example Usage (in Activity or Fragment):
 * 
 * class MyFragment: Fragment() {
 *   override fun onCreate(savedInstanceState: Bundle?) {
 *     traceUntilFirstDraw { time -> Tracer.trace("MyScreen", time) }
 *
 *     super.onCreate(savedInstanceState)
 *   }
 * }
 */
fun LifecycleOwner.traceUntilFirstDraw(
  onMeasured: (Long) -> Unit,
) {
  UntilFirstDrawAutoTracer(
      lifecycleOwner = this,
      onMeasured = onMeasured,
  )
      .setup()
}

class UntilFirstDrawAutoTracer(
  private val lifecycleOwner: LifecycleOwner,
  private val onMeasured: (Long) -> Unit,
) {

  inner class FirstDrawObserver : DefaultLifecycleObserver {
    private var startTime: Long = 0L
    private var onFirstDrawListener: OnFirstDrawListener? = null

    override fun onCreate(owner: LifecycleOwner) {
      startTime = SystemClock.elapsedRealtime()
    }

    override fun onStart(owner: LifecycleOwner) {
      val viewTreeObserver: ViewTreeObserver? = when (owner) {
        is Fragment -> owner.view?.viewTreeObserver
        is Activity -> owner.window.decorView.viewTreeObserver
        else -> null
      }

      viewTreeObserver?.let { nonNullViewTreeObserver ->
        onFirstDrawListener = OnFirstDrawListener(
            viewTreeObserver = nonNullViewTreeObserver,
            startTime = startTime,
            onMeasured = onMeasured,
        )
        nonNullViewTreeObserver.addOnDrawListener(onFirstDrawListener)
      }
    }

    override fun onStop(owner: LifecycleOwner) {
      onFirstDrawListener?.dispose()
    }
  }

  fun setup() {
    lifecycleOwner.lifecycle.addObserver(FirstDrawObserver())
  }
}

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

  override fun onDraw() {
    if (!firstOnDrawHappened) {
      firstOnDrawHappened = true

      val finishTime = SystemClock.elapsedRealtime()
      onMeasured?.invoke(finishTime - startTime)

      onFinish?.invoke()

      Handler(Looper.getMainLooper()).post {
        dispose()
      }
    }
  }

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