Suggestion to improve BraveTraceContextProvider#getTraceContextLater

[be-hase]

Feature Request

Is your feature request related to a problem? Please describe

Currently, there is an little issue with TracingContext propagation when using Kotlin Coroutines.

In coroutines, there is a class called ThreadContextElement. This allows for Context propagation with a ThreadLocal-like feel. If we implement the following class, we can realize the propagation of TraceContext in Coroutines.

class TracingThreadContextElement(
    private val traceContext: TraceContext? = currentTraceContext()?.get(),
) : ThreadContextElement<CurrentTraceContext.Scope?> {
    override val key: CoroutineContext.Key<*>
        get() = KEY

    override fun updateThreadContext(context: CoroutineContext): CurrentTraceContext.Scope? {
        return currentTraceContext()?.maybeScope(traceContext)
    }

    override fun restoreThreadContext(context: CoroutineContext, oldState: CurrentTraceContext.Scope?) {
        oldState?.close()
    }

    companion object {
        val KEY = object : CoroutineContext.Key<TracingThreadContextElement> {}

        fun currentTraceContext(): CurrentTraceContext? {
            return Tracing.current()?.currentTraceContext()
        }
    }
}

withContext(TracingThreadContextElement(...)) {

    // traceContext is propagated in this coroutine scope

}

Coroutines implementation of lettuce is reactor-based. Therefore, a ReactorContext must be set up to propagate the TracingContext. https://github.com/lettuce-io/lettuce-core/blob/7be77cb8298b456cb68754f71c5eb3fed412c0f9/src/main/java/io/lettuce/core/tracing/BraveTracing.java#L39-L46

Coroutine User needs to be aware of Reactor, which is undesirable.

Describe the solution you'd like

How about fallback to getTraceContext() when there was no ReactorContext in this part? https://github.com/lettuce-io/lettuce-core/blob/7be77cb8298b456cb68754f71c5eb3fed412c0f9/src/main/java/io/lettuce/core/tracing/BraveTracing.java#L484-L495

@Override
public Mono<TraceContext> getTraceContextLater() {

    return Mono.deferContextual(Mono::justOrEmpty)
               .filter(it -> it.hasKey(Span.class) || it.hasKey(brave.propagation.TraceContext.class)).map(it -> {

                if (it.hasKey(Span.class)) {
                    return new BraveTraceContext(it.get(Span.class).context());
                }

                return new BraveTraceContext(it.get(brave.propagation.TraceContext.class));
            })
            .switchIfEmpty(Mono.justOrEmpty(getTraceContext())); // HERE !
}

This way, the TracingContext can be propagated if the aforementioned TracingThreadContextElement is used.

More to the point, even when using the reactive interface in Spring MVC, it is possible to propagate a TracingContext that is tied to the request thread.

@GetMapping("/reactive")
fun reactive(): String {
    return reactive.set("hoge", "hoge").block()!!
}

Describe alternatives you've considered

The problem can also be solved by lettuce user writing code like the following:

class CustomBraveTracing(delegate: io.lettuce.core.tracing.Tracing) :
    io.lettuce.core.tracing.Tracing by delegate {
    override fun initialTraceContextProvider(): TraceContextProvider {
        return CustomBraveTraceContextProvider.INSTANCE
    }
}

enum class CustomBraveTraceContextProvider : TraceContextProvider {
    INSTANCE;

    override fun getTraceContext(): TraceContext? {
        val tracer = Tracing.currentTracer()
        if (tracer != null) {
            val span = tracer.currentSpan()
            if (span != null) {
                return BraveTraceContext.create(span.context())
            }
        }
        return null
    }

    override fun getTraceContextLater(): Mono<TraceContext> {
        return Mono.justOrEmpty(traceContext)
    }
}

Teachability, Documentation, Adoption, Migration Strategy

NONE

[mp911de]

The tracing context propagation code is written in a defensive way to avoid accidental propagation from the thread-local context. It would make sense to explore ways to fall back to the trace context if coroutine types are on the classpath.

[be-hase]

The tracing context propagation code is written in a defensive way to avoid accidental propagation from the thread-local context

Yes, there is certainly a probability of accidental propagation. It's gonna be hard to make a decision.

It would make sense to explore ways to fall back to the trace context if coroutine types are on the classpath.

The idea also looks good. However, depending on the application, there are cases where reactor and coroutines are used together. (For example, Kotlin and Java mixed application)

Therefore, ideally it seems better to use thread-local only for Coroutines Interface.