Tracking Performance in Kotlin Suspending Functions with Firebase Performance Monitoring

Suspended cups — Photo by Chris Scott on Unsplash

Using Firebase Performance Monitoring to track the performance of a plain-ol’ function (Java or Kotlin) is drop-dead simple using the @AddTrace annotation:

@AddTrace("some-trace-name")
fun myBoringFunction() {
    println("nothing to see here!")
}

Wait a few moments, and you’ll see some-trace-name in the Firebase Console!

But what if you want to track the performance of a suspend function? Unfortunately, this is not as straightforward.

Attempt #1: Suspending Disbelief

Let’s say you have some uncomplicated suspending function, like so:

suspend fun myAsyncFunction() : Int {
    val result = 1 + 1
    delay(100L)   
    println("something to see here: $result")
    return result
}

A naive attempt to track this function’s performance would, as you might expect, to use @AddTrace:

@AddTrace("some-other-trace-name")
suspend fun myAsyncFunction() : Int { ... }

Just eyeballing it, it should work, right? Even if you compile, the firebase-perf Gradle plugin won’t complain, and the Firebase Performance Monitoring library won’t crash your app. So maybe everything will Just Work™️? 🤞🏽

Plot Twist: it won’t work.

To understand why it doesn’t work, you’ll need to understand two things:

1. How the Kotlin compiler produces a suspending state machine
2. How the firebase-perf Gradle plugin instruments the @AddTrace annotation into an actual Trace invocation

For #1, Manuel Vivo has a great article going into the nuts and bolts of what the suspend keyword actually means, and if you’re not familiar with the process, I recommend checking it out.
For our purposes, we’ll simplify, and say that the Kotlin compiler will rewrite your suspend method signature, appending a Continuation<T> parameter, which allows it to build the state machine.

For #2, I don’t work at Google, so I’m not entirely sure how the @AddTrace instrumentation works under the hood. But I’ve got some guesses!

Let’s take our previous example and “compile” it. We’ll make 3 changes:

1. Simplify the “guts” of the coroutine state machine
2. Add some “fake” Perf traces
3. “Decompile” everything into Kotlin source code, because nobody wants to read Kotlin JVM bytecode!

fun myAsyncFun(cont : Continuation) {
    val perf = FirebasePerformance.getInstance()
    val trace = perf.newTrace("some-other-trace-name")
    trace.start()
    val result = 1 + 1
    
    //begin: coroutine state machine
    when(cont.label){
        0 -> {
            //will call myAsyncFun with cont.label=1
            DelayKt.delay(100, cont) 
        }
        1 -> {
            println("something to see here: $result")
            //hands control+result back to caller
            cont.resume(result, ...) 
        }
    }
    //end: coroutine state machine
    
    trace.end()
}

The coroutine’s state machine is muddying the waters — it’s no wonder Firebase’s automated instrumentation can’t figure out where to put the trace.end()! What’s a performance-minded developer to do?

Attempt #2: If you Try, Finally you’ll succeed

The heart of the issue is that the Firebase performance instrumentation can’t query suspending Kotlin bytecode for the method entry and exit.
But, as developers, we know when a suspending function will enter and exit; in fact, there’s even a Java/Kotlin paradigm to codifies this: try/finally

try/finally is exactly what we need: no matter what happens inside a method body, finally will always be executed last (or, finally!) — this is a guarantee from the JVM, and you can see this reflected in any compiled Java/Kotlin bytecode. Read more here

Armed with this knowledge, and a little bit of Kotlin magic✨, we can write a little wrapper that can be used anywhere:

inline fun <E> trace(name : String, block: (Trace) -> E): E {
    val trace = startTrace(name) //creates & starts a new Trace
    return try {
        block(trace)
    } finally {
        trace.stop()
    }
}

Breaking it down:

• the inline keyword means we won’t interfere with any coroutine state machine; the entire suspending invocation will occur in the try {} block
• startTrace just invokes FirebasePerformance, names the trace, and starts it.
• we accept a block: (Trace) -> E, so that the instrumented code can append information to the Trace, if needed
• in the finally block, we stop the trace, and return value produced by block
• any exceptions are propagated to the calling code, since there’s no catch defined

Sweet; simple; and, most importantly, it works!

Let’s rewrite our first suspending function using our new toy:

suspend fun myAsyncFunction() : Int {
    val result = 1 + 1
    delay(100L)   
    println("something to see here: $result")
    return result
}

becomes:

suspend fun myAsyncFunction() : Int {
    return trace("some-other-trace") {
        val result = 1 + 1
        delay(100L)   
        println("something to see here: $result")
        return@trace result
    }
}

This is our solution! If you want to achieve this same result without nesting your method into a lambda, you can push it into a private function:

suspend fun myAsyncFunction() : Int {
    return trace("some-other-trace") { _myAsyncFunction() }
}
private suspend fun _myAsyncFunction() : Int {
    val result = 1 + 1
    delay(100L)   
    println("something to see here: $result")
    return result
}

The Firebase Performance Monitoring Console will now show the total time myAsyncFunction runs, including the time that any actual suspension takes 🎉

Happy coding!

Next Steps

I’ve filed a feature request with Firebase to augment @AddTrace to support suspending functions. Let’s see what happens, maybe this blog post will be obsolete in 6 months 👀

Thanks to Akshay Chordiya and Doug Stevenson for reviewing the early drafts! 🙏🏾

Note: this story was originally published on my personal website: https://jvmname.dev/posts/2020/03/tracking-performance-in-kotlin-suspending-functions/

Sources | Extra Reading

The suspend modifier — Under the hood

Java Code To Byte Code - Part Two