TraceTest.kt

package com.example

import kotlin.annotation.AnnotationRetention.RUNTIME
import kotlin.annotation.AnnotationTarget.FUNCTION

/**
 * @see TracingRunListener
 */
@Retention(RUNTIME)
@Target(FUNCTION)
annotation class TraceTest

TracingRunListener.kt

package com.example

import android.os.Debug
import android.os.SystemClock
import org.junit.runner.Description
import org.junit.runner.notification.RunListener
import timber.log.Timber
import java.io.File

/**
 * Turns on sampling profiling on Firebase UI test methods annotated with [TraceTest]
 * [TracingRunListener] is enabled for all Firebase UI tests and selectively turns on profiling by
 * looking up the annotations on the test methods.
 *
 * The resulting traces are stored on the device in /sdcard/traces and can be downloaded by Firebase to
 * gcloud storage with `--directories-to-pull "/sdcard/traces/"`
 *
 * Once you've downloaded the trace file, open the Profiler in Android Studio
 * (Cmd + Shift + A and type Profiler). Click the + sign, select "Load from file..." and open up
 * the trace file. Now go read this: https://developer.android.com/studio/profile/cpu-profiler.html
 *
 * First you'll want to select a thread in the Threads view. You're most likely interested in
 * "main" and "Instr: runner name". Once you've selected a thread, Call Chart will give you a sense
 * of the timeline / where CPU time is spent over time. Flame Chart is not time based but groups
 * together identical stack traces, so it helps you identify calls that happens a lot.
 *
 * When looking at the Bottom Up tab, it's useful to try sorting by Total but also Self (which would
 * e.g. surface thread waiting calls).
 */
internal class TracingRunListener : RunListener() {

  @Volatile
  var tracingTest = false

  private val Description.traceTest
    get() = getAnnotation(TraceTest::class.java) != null

  private val Description.traceFile
    get() = "$FOLDER$className.$methodName-${SystemClock.uptimeMillis()}.trace"

  override fun testRunStarted(description: Description?) {
    val defaultHandler = Thread.getDefaultUncaughtExceptionHandler()
    Thread.setDefaultUncaughtExceptionHandler { thread, throwable ->
      stopMethodTracing("Stopped profiler trace because process crashed")
      defaultHandler.uncaughtException(thread, throwable)
    }
  }

  override fun testStarted(description: Description) {
    if (description.traceTest) {
      tracingTest = true
      val folder = File(FOLDER)
      folder.mkdirs()
      val traceFile = description.traceFile
      Timber.d("Starting a profiler trace, stored at %s", traceFile)
      Debug.startMethodTracingSampling(traceFile, BUFFER_SIZE, INTERVAL_MICROS)
    }
  }

  override fun testFinished(description: Description) {
    stopMethodTracing("Stopped profiler trace because test finished")
  }

  private fun stopMethodTracing(logMessage: String) {
    if (tracingTest) {
      tracingTest = false
      Timber.d(logMessage)
      Debug.stopMethodTracing()
    }
  }

  companion object {
    private const val FOLDER = "/sdcard/traces/"
    /**
     * A buffer of that size is immediately allocated in memory. The final files is smaller than that.
     * https://android.googlesource.com/platform/art/+/master/runtime/trace.cc#556
     */
    private const val BUFFER_SIZE = 50 * 1024 * 1024
    private const val INTERVAL_MICROS = 1000
  }
}