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.