Better support for working with Observation from Kotlin

[ilya40umov]

Right now (as far as I can tell), there is no extra support for using Observation API from Kotlin code.

For example, to create a new Observation in normal, blocking code, one first needs to invent a convenience function like this:

private fun <T> ObservationRegistry.observeNotNull(name: String, block: () -> T): T =
    Observation.start(name, this).observe(Supplier { block() })!!

And when it comes to wrapping a suspending block of code within Observation it gets even more complicated:

private suspend fun <T> ObservationRegistry.observe(name: String, block: suspend () -> T): T {
    val observation = Observation.createNotStarted(name, this)
    observation.start()
    return try {
        mono(observationRegistry.asContextElement() + Dispatchers.Unconfined) {
            block()
        }.contextWrite { context ->
            // XXX this seems to be the safest option of making the new observation current
            context.put(ObservationThreadLocalAccessor.KEY, observation)
        }.awaitSingle()
    } catch (error: Throwable) {
        throw error
    } finally {
        // XXX or maybe it's safer to stop this in `doOnTerminate` of the created mono?
        observation.stop()
    }
}

At the same time, when working with Otel APIs directly, there is an option to construct a span/context and put it into scope via Otel's ContextExtensions.kt

[jonatan-ivanov]

I think this issue should be in Micrometer instead where the Observation API is and some other bits of Kotlin support: https://github.com/micrometer-metrics/micrometer/blob/main/micrometer-core/src/main/kotlin/io/micrometer/core/instrument/kotlin/

[shakuzen]

Thanks for opening the issue. Is this something you would be willing to make a pull request for, @ilya40umov? I don't regularly use Kotlin so it will probably be easier for folks who do to suggest the changes.

[ilya40umov]

@shakuzen I'll definitely try looking into this (not sure sure if I'll have time for this next week though).

[ilya40umov]

So far I was able to improve a bit upon the second extension function I posted originally, so now it looks as follows:

suspend fun <T> ObservationRegistry.observe(name: String, block: suspend () -> T): T {
    val observation = Observation.start(name, this)
    val coroutineContext = currentCoroutineContext().minusKey(Job.Key)
    return mono(coroutineContext + observationRegistry.asContextElement()) {
        block()
    }.contextWrite { context ->
        context.put(ObservationThreadLocalAccessor.KEY, observation)
    }.doOnTerminate {
        observation.stop()
    }.awaitSingle()
}

This is basically making sure that it preserves most of the original coroutine context, similar to how TransactionalOperatorExtensions.kt#L48 is doing it.

Then it can be used from Kotlin code like this:

observationRegistry.observe(name = "constructAnswer") {
  delay(10L)
  "42"
}

I am a bit conflicted about this being an extension function of ObservationRegistry though.
On the one hand, it makes it easy to use it. On the other hand, it's hiding away a bunch of methods available for Observation construction otherwise (e.g. Observation.start has a bunch of overloads).
Some of this problematic could be solved by adding some optional arguments to the method and making them null by default.

Or this whole method could be added as an extension of already existing Observation.
But then using it would require referring to observationRegistry twice.
E.g.

Observation.start("constructAnswer", observationRegistry).observe(observationRegistry) {
  delay(10L)
  "42"
}

This is because the method also needs a reference to observationRegistry, which I can't get from an already constructed Observation.

[ilya40umov]

Hm. Actually, I may have over-complicated what is really needed.
Edit: Looks like this indeed can be implemented without relying on Reactor Context, but it's a bit tricky. (I first thought it was very straightforward, until I tried using an implementation relying purely on openScope() from a separate thread and adding Dispatchers.Unconfined to the context.).

The following implementation should do the trick without tapping into mono / contextWrite from Reactor:

private suspend fun <T> Observation.observeAndWait(
    observationRegistry: ObservationRegistry, 
    block: suspend () -> T
): T {
    start()
    return withContext(
        openScope().use { observationRegistry.asContextElement() }
    ) {
        try {
            block()
        } catch (error: Throwable) {
            error(error)
            throw error
        } finally {
            stop()
        }
    }
}

As an optimization, it may be possible to avoid adding observationRegistry.asContextElement() to the coroutine context when it has already been added (e.g. by a CoWebFilter or similar means), for which purpose the coroutine context could be checked for containing KotlinObservationContextElement.KEY first. I didn't try it yet, as KotlinObservationContextElement is not public.

[ilya40umov]

As for the "normal", i.e. non-suspending API, we could extend Observation like follows:

fun <T : Any> Observation.observeNotNull(block: () -> T): T = observeBlock(block)
    
fun <T : Any?> Observation.observeNullable(block: () -> T): T = observeBlock(block)

private fun <T> Observation.observeBlock(block: () -> T): T {
    start()
    return try {
        openScope().use {
            block()
        }
    } catch (error: Throwable) {
        error(error)
        throw error
    } finally {
        stop()
    }
}

I'm only including two methods, i.e. observeNotNull and observeNullable because I don't like the name of observeBlock and I am not sure how else to call it, so that it would have a name different from observe.

[ilya40umov]

@shakuzen I will think about it for a few more days and then will raise a PR with the changes I mentioned above and hopefully some tests.

[ilya40umov]

Just for reference, here are a couple of examples of how to both versions of the APIs (suspending and non-suspending) look in Kotlin code:

• Observation.observeNotNull
• Observation.observeAndWait

[ilya40umov]

I've opened a PR where I ended up adding 2 extension functions to ObservationRegistry. I will try to test those (especially the coroutine version) a bit more, but I would happy to hear any feedback on this change overall.

[ilya40umov]

I've added a second PR which is adding the same 2 extensions on Observation instead of ObservationRegistry. Let's see which one gets through the review.

[ArvindEnvoy]

@ilya40umov is it possible to map the CoroutineContext to an @Observed annotation as well? I'm not fully sure how the AOP libraries work and facilitate those annotations, but ideally we could use @Observed on suspend functions in Kotlin, and the Observation would manage the content of the coroutine like it does in your above example.

[ilya40umov]

@ArvindEnvoy I've created a separate issue to track the topic of @Observed annotation. I think it's definitely a good idea to make the associated aspect handle suspending methods as well, but let's see what the maintainers will say.

[ilya40umov]

@jonatan-ivanov @shakuzen I totally forgot about this issue / associated PRs, as I ended up implementing a solution based on #4823 in a shared library at my current job. Ideally, it would be still nice to decide which of the two PRs you guys would rather keep: #4772 or #4823 .

[marcingrzejszczak]

Hey, I've totally missed the PRs, sorry. What's the difference between the two?

[ilya40umov]

So, one is adding the extension functions to Observation and the other one is adding those extension function to ObservationRegistry, which takes care of creating the observation.

I initially created this PR #4772 with ObservationRegistry as I thought Observation didn't have access to ObservationRegistry, but as I learned that Observation contained a reference to ObservationRegistry I created the second PR #4823.
I'm personally tending towards closing #4772, but I wanted to allow you folks to pick which one you would rather merge.

[marcingrzejszczak]

Hey, I've totally missed the PRs, sorry. What's the difference between the two?

[ilya40umov]

Hey, I've totally missed the PRs, sorry. What's the difference between the two?

@marcingrzejszczak looks like your question from before (that I answered above) got duplicated somehow. Were you planning to write something else?

[PierrickPuimeanChieze]

First, let just say that this PR helped me a lot, even in its current state. The observeAndWait method allowed me to keep a trace with multiple span across multiple coroutine and nested call to suspend function. So again, thanks a lot @ilya40umov

But i got a question about this bit :

withContext(
openScope().use { observationRegistry.asContextElement() }
)

I don't really understand why the need to open and close a scope and retrieving the observationRegistry as a context element ? I know from my experiment prior to find this code this is needed, but I don't really understand the mechanisms behind it. Could you help me to clarify my mind on the subject ?

[GarryIV]

I don't really understand why the need to open and close a scope and retrieving the observationRegistry as a context element ?

Yeah, it looks strange but works 💯

And 👍 @ilya40umov

[ilya40umov]

@PierrickPuimeanChieze TL;DR this is needed because when coroutines suspend / resume they may end up continuing on a different thread. The concept of "current" observation is by default implement via a ThreadLocal variable, which is not handled correctly when you start using coroutines. This is why this "asContextElement" is needed, which is adding KotlinObservationContextElement to the context, which in turn implements ThreadContextElement, basically a context element that is meant to be used to save/restore ThreadLocal value(s).

I can't possibly cover this topic in depth in a Github comment, but you may find this repo I created a while ago useful. It covers actually 3 different concurrency models and different libraries for tracing (i.e. Brave vs OpenTelemetry vs using Spring & Micrometer).

[PierrickPuimeanChieze]

Thanks for your answer. I already knew about the way coroutine worked (took me a lot of brain juice a while ago to understand why my transaction got lost in oblivion when used with coroutines) but my question was more about the use of openScope. I found the definitions and explanation of the Scope quite succinct in the micrometer documentation,honestly.

I will check your repo. Maybe it will help my understanding.

And again, thanks a lot.