data class SpanContext(val span: Span) : AbstractCoroutineContextElement(SpanContext) {
    companion object Key : CoroutineContext.Key<SpanContext>
}

fun CoroutineContext.currentSpan(): Span {
    return get(SpanContext)?.span ?: throw IllegalStateException("No active Span in context")
}

suspend fun <T> withTrace(
    name: String,
    spanName: String,
    block: suspend CoroutineScope.() -> T
): T {
    val trace = startTrace(name, spanName)
    return withSpan(trace, block)
}

suspend fun <T> withSpan(name: String, block: suspend CoroutineScope.() -> T): T {
    val span = currentCoroutineContext().currentSpan()
    val childSpan = span.startChild(name)
    return withSpan(childSpan, block)
}

suspend fun <T> withSpan(span: Span, block: suspend CoroutineScope.() -> T): T {
    return try {
        withContext(SpanContext(span), block)
    } catch (e: Throwable) {
        if (e is CancellationException) {
            span.setStatus("cancelled")
        } else {
            span.setThrowable(e)
            span.setStatus("error")
        }
        throw e
    } finally {
        span.finish()
    }
}

suspend fun getData(): String {
  return withSpan("getData") {
    delay(1000)
    "foo"
  }
}