Better Stack Traces with Coroutines

gistfile1.txt

import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.joetr.coroutines.CoroutineScopeExtensions.launchWithBreadcrumbs
import com.joetr.coroutines.CoroutinesBreadcrumbExceptionHandler.createBreadcrumbsExceptionHandler
import com.joetr.coroutines.FlowExtensions.stateInWithBreadcrumbs
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*
import java.util.*
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext

class MainViewModel : ViewModel() {

    val state = flow<String> {
        throw IllegalArgumentException("oh no")
    }.stateInWithBreadcrumbs(
        scope = viewModelScope,
        started = SharingStarted.Eagerly,
        initialValue = "",
    )

    fun testException() {
        viewModelScope.launchWithBreadcrumbs {
            throwException1()
        }
    }

    private suspend fun throwException1() {
        delay(300)
        throwException2()
    }

    private suspend fun throwException2() {
        delay(300)
        throwException3()
    }

    private suspend fun throwException3() {
        delay(300)
        throw IllegalArgumentException("oh no")
    }
}

internal object CoroutinesBreadcrumbExceptionHandler {

    private fun createBreadcrumbsException() =
        BreadcrumbException()
            .apply {
                val queue: Queue<StackTraceElement> = LinkedList(stackTrace.asList())

                val iterator = queue.iterator()
                while (iterator.hasNext()) {
                    val next = iterator.next()

                    // we need to remove all of the 'helper' classes from the stack trace
                    // as that really just adds noise and isn't overly helpful
                    if (next.className == CoroutinesBreadcrumbExceptionHandler::class.qualifiedName ||
                        next.className == FlowExtensions::class.qualifiedName ||
                        next.className == CoroutineScopeExtensions::class.qualifiedName
                    ) {
                        iterator.remove()
                    } else {
                        break
                    }
                }
                stackTrace = queue.toTypedArray()
            }

    private fun Throwable.addBreadcrumbs(
        breadcrumbsException: BreadcrumbException,
    ): Throwable {
        var lastCause = cause
        while (lastCause?.cause != null) {
            lastCause = lastCause.cause
        }

        return lastCause?.initCause(breadcrumbsException) ?: initCause(breadcrumbsException)
    }

    internal fun createBreadcrumbsExceptionHandler(): CoroutineExceptionHandler {
        val breadcrumbsException = createBreadcrumbsException()

        return CoroutineExceptionHandler { _, throwable ->
            val currentThread = Thread.currentThread()
            run {
                currentThread.uncaughtExceptionHandler ?: Thread.getDefaultUncaughtExceptionHandler()
            }.uncaughtException(
                currentThread,
                throwable.addBreadcrumbs(breadcrumbsException),
            )
        }
    }
}

internal class BreadcrumbException : RuntimeException("More details about the original stack trace:")

object CoroutineScopeExtensions {
    fun CoroutineScope.launchWithBreadcrumbs(
        context: CoroutineContext = EmptyCoroutineContext,
        start: CoroutineStart = CoroutineStart.DEFAULT,
        block: suspend CoroutineScope.() -> Unit,
    ): Job {
        return this.launch(
            context = context + createBreadcrumbsExceptionHandler(),
            start = start,
            block = block,
        )
    }
}

object FlowExtensions {
    fun <T> Flow<T>.stateInWithBreadcrumbs(
        scope: CoroutineScope,
        started: SharingStarted,
        initialValue: T,
    ): StateFlow<T> {
        return stateIn(
            scope = scope + createBreadcrumbsExceptionHandler(),
            started = started,
            initialValue = initialValue,
        )
    }
}