Otel4s with async-profiler profiling data.

[rochala]

Hey, I've been working to port https://github.com/grafana/otel-profiling-java/tree/6ddc19fe2e37ad1b60ec26baf613a049c2be30e5 to otel4s and I managed to hack a working prototype.

This is proof of concept that I've created:

  val runtime0 = {
    IORuntimeBuilder()
      .transformCompute { ec => wrapExecutionContext(ec, localSpan) }
      .transformBlocking { ec => wrapExecutionContext(ec, localSpan) }
      .build()
  }

  override protected def runtime: IORuntime = runtime

  def wrapExecutionContext(ec: ExecutionContext, localSpan: LocalSpan): ExecutionContext = {
    new ExecutionContext {
      def execute(runnable: Runnable): Unit = {
        val spanOpt = localSpan.get()

        ec.execute(new Runnable {
          def run(): Unit = {

            spanOpt.foreach { span =>
              profiler.setTracingContext(span.spanId, 0)
            }

            try {
              runnable.run()
            } finally {
              spanOpt.foreach { _ => profiler.setTracingContext(0, 0) }
            }
          }
        })
      }

      def reportFailure(cause: Throwable): Unit = ec.reportFailure(cause)
    }
  }

  case class TraceSpan(val spanId: Long)

  class LocalSpan(_ioLocal: () => IOLocal[Option[TraceSpan]]) {
    private lazy val ioLocal: IOLocal[Option[TraceSpan]] = _ioLocal()
    private lazy val unsafeThreadLocal: ThreadLocal[Option[TraceSpan]] = {
      val fiberLocal = ioLocal.unsafeThreadLocal()

      new ThreadLocal[Option[TraceSpan]] {
        override def initialValue(): Option[TraceSpan] = None

        override def get(): Option[TraceSpan] =
          if (IORuntime.isUnderFiberContext()) fiberLocal.get() else super.get()

        override def set(value: Option[TraceSpan]): Unit =
          if (IORuntime.isUnderFiberContext()) fiberLocal.set(value) else super.set(value)
      }
    }
    def get() = unsafeThreadLocal.get()
    def set(value: Option[TraceSpan]) = unsafeThreadLocal.set(value)
  }

And the span ID can then be used to find flamegraph of given trace span:

The issue is that I'm quite unsure the implementation details of how we can set profiler context to span ID.
Right now I've created IOLocal with custom IORuntime, but maybe there is a better way to achieve it ?

[iRevive]

Hey, I will take a look next week.

[iRevive]

Dependening on the platform, you can do things slightly differently.
If you use otel4s-oteljava backend, you can use IOLocalContextStorage to enable seamless context propagation between CE and OpenTelemetry. Here is the doc: https://typelevel.org/otel4s/oteljava/tracing-context-propagation.html.

Here is an example:

val runtime0 = {
    IORuntimeBuilder()
      .transformCompute { ec => wrapExecutionContext(ec, localSpan) }
      .transformBlocking { ec => wrapExecutionContext(ec, localSpan) }
      .build()
  }

override protected def runtime: IORuntime = runtime

def wrapExecutionContext(ec: ExecutionContext): ExecutionContext = {
    new ExecutionContext {
      def execute(runnable: Runnable): Unit = {
        // get the current span
        val ctx = io.opentelemetry.api.trace.Span.current().getSpanContext

        ec.execute(new Runnable {
          def run(): Unit = {
            if (ctx.isValid)
              profiler.setTracingContext(span.spanId, 0)

            try {
              runnable.run()
            } finally {
              if (ctx.isValid) profiler.setTracingContext(0, 0)
            }
          }
        })
      }

      def reportFailure(cause: Throwable): Unit = ec.reportFailure(cause)
    }
  }
  
def main: IO[Unit] = {
  // make `otel4s` use the `context-storage`-backed IOLocal
  given LocalProvider[IO, Context] = IOLocalContextStorage.localProvider[IO]

  OtelJava.autoConfigured[IO]().use { otelJava => 
    ... // do the rest
  }
}

However, if you want to use OpenTelemetry Java Agent, it won't work out of the box. In this case, you can use otel4s distribution of the OTeL Java Agent, here are the details: https://typelevel.org/otel4s/oteljava/agent.html.

[iRevive]

If you want to use otel4s-sdk backend (otel4s implementation of the OpenTelemetry spec), there are not many alternatives to your approach.

• Implement a custom span processor to set the pyroscope.profile.id attribute
• Customize IORuntime to set the span id before running the task

In this case, we should unseal SpanProcessor to make it possible.

[rochala]

Okay, so I'm using otel4s-sdk and I've managed to make it work here: simple-scala-tooling/sls#55

But we've found the issue with IOLocals and I'm not sure if this is intended or is it a bug:

//> using scala 3.7.3
//> using dep org.typelevel::cats-effect:3.6.3

import scala.concurrent.ExecutionContext
import java.util.UUID
import cats.effect.IOFiber
import cats.effect.IOApp
import cats.effect.IO
import cats.effect.unsafe.IORuntime
import cats.syntax.all.*
import cats.effect.IOLocal

object EEE extends IOApp.Simple {
  var variabl: Option[Int] = None
  var loc: IOLocal[Int] = null

  def wrap(ec: ExecutionContext): ExecutionContext = new ExecutionContext {

    def execute(runnable: Runnable): Unit = {
      val res = try { loc.unsafeThreadLocal().get() } catch { case _: Throwable => 2 }
      ec.execute(() => {
        variabl = res.some
        runnable.run()
      })

    }
    def reportFailure(cause: Throwable): Unit = ec.reportFailure(cause)
  }

  override protected def runtime: IORuntime = IORuntime.builder().transformCompute(wrap).transformBlocking(wrap).build()

  def run: IO[Unit] = IOLocal(0).flatMap { loc0 =>
    this.loc = loc0
    loc.set(1) *> IO(println(variabl)) *> // prints 2 which is default 
    IO.cede *> IO(println(variabl)) *> // prints 1 because IO.cede created new fiber
    loc.set(-1) *> IO(println(variabl)) *> // still prints 1 because IOLocal.set does not create new fiber
    IO.cede *> IO(println(variabl)) *> // prints -1 because IO.cede reschedules the fiber
  }
}

if you run this with scala-cli run -Dcats.effect.trackFiberContext=true you will see following output:

Compiled project (Scala 3.7.3, JVM (21))
Some(2)
Some(1)
Some(1)
Some(-1)

What I was expecting here to happen was to have:

Compiled project (Scala 3.7.3, JVM (21))
Some(1)
Some(1)
Some(-1)
Some(-1)

This basically makes it impossible to track the tracing context which was the initial point of this discussion via custom executor. The current workaround is to create a wrapper on Tracer that calls IO.cede after and before the actual thunk of the trace.

The reason why this breaks tracing is if we run:

        ec.execute(new Runnable {
          def run(): Unit = {
            if (ctx.isValid)
              profiler.setTracingContext(span.spanId, 0)

span.spanId is not actually set correctly when the computations are happening.
This is case loc.set(1) *> IO(println(variabl)) at which println(variabl) still returns 2, which in this case would be no span, and for later cases it would return incorrect one so profiler tracing context would be incorrect.

[iRevive]

Yes, that doesn't look okay at first glance. I will double-check soon.

[iRevive]

I tried to use https://github.com/grafana/otel-profiling-java in my otel4s demo: https://github.com/iRevive/otel4s-showcase/pull/1/files.
In this demo, the services run with otel4s-oteljava and https://github.com/iRevive/otel4s-opentelemetry-java - our distribution of OTeL Java Agent, which includes CE and otel4s instrumentation for seamless context propagation.

It seems to be working. pyroscope.profile.id is always set and is correct. However, the flame graph data is different for the same endpoint.

Even though the profile.

1. No data

2. CE calls

3. fs2 calls

4. gRPC calls

[iRevive]

I tried using single-threaded executors and explicitly set tracing info:

PyroscopeAgent.start(Config.build())
private val runtime0 =
  IORuntimeBuilder()
    .setCompute(ExecutionContext.fromExecutor(Executors.newSingleThreadExecutor()), () => ())
    .setBlocking(ExecutionContext.fromExecutor(Executors.newSingleThreadExecutor()), () => ())
    .transformCompute(ec => wrap(ec))
    .transformBlocking(ec => wrap(ec))
    .build()
private def wrap(ec: ExecutionContext): ExecutionContext = new ExecutionContext {
  def execute(runnable: Runnable): Unit = {
    val ap = PyroscopeAsyncProfiler.getAsyncProfiler
    val spanContext = Span.current().getSpanContext
    ec.execute { () =>
      if spanContext.isValid then ap.setTracingContext(java.lang.Long.parseUnsignedLong(spanContext.getSpanId, 16), 0)
      runnable.run()
      if spanContext.isValid then ap.setTracingContext(0, 0)
    }
  }
  def reportFailure(cause: Throwable): Unit = ec.reportFailure(cause)
}

The behavior remains the same - the flame graph is still inconsistent.