Testkit: IOLocalContextStorage is not configured for use as the ContextStorageProvider

[AlixBa]

discord

having otel4s-oteljava-testkit as a dependency, pulls opentelemetry-sdk-testing which forces the ContextStorage to be SettableContextStorage (through https://github.com/open-telemetry/opentelemetry-java/blob/5c4a0ad78f021eef925e9ed69b67885b1b3cad3a/sdk/testing/src/main/resources/META-INF/services/io.opentelemetry.context.ContextStorageProvider)

In a context I want to test (in my test module) java propagation (with IOLocalContextStorage), it fails because

java.lang.IllegalStateException: IOLocalContextStorage is not configured for use as the ContextStorageProvider. The current ContextStorage is: io.opentelemetry.sdk.testing.context.SettableContextStorageProvider$SettableContextStorage

if I display SettableContextStorageProvider.getContextStorage is properly links to org.typelevel.otel4s.oteljava.context.IOLocalContextStorage@1ce5131f.

cause

[AlixBa]

as a quick fix on my side (and because it's only for a test module), I did a super dirty LocalProvider

import cats.effect.IOLocal
import cats.effect.LiftIO
import cats.effect.MonadCancelThrow
import cats.mtl.Local
import io.opentelemetry.context.ContextStorage
import io.opentelemetry.sdk.testing.context.SettableContextStorageProvider
import org.typelevel.otel4s.context.LocalProvider
import org.typelevel.otel4s.oteljava.context.Context
import org.typelevel.otel4s.oteljava.context.LocalContext

object IOLocalContextStorage {

  def localProvider[F[_]](implicit
      L: LiftIO[F],
      F: MonadCancelThrow[F]
  ): LocalProvider[F, Context] =
    new LocalProvider[F, Context] {
      def local: F[Local[F, Context]] =
        ContextStorage.get() match {
          case storage
              if storage.getClass.getName == "org.typelevel.otel4s.oteljava.context.IOLocalContextStorage" =>
            val local = storage.getClass.getMethod(
              "local",
              Class.forName("cats.effect.kernel.MonadCancel"),
              Class.forName("cats.effect.LiftIO")
            )
            whenPropagationEnabled(local.invoke(storage, F, L).asInstanceOf[LocalContext[F]])

          case other
              if other.getClass.getName == "io.opentelemetry.sdk.testing.context.SettableContextStorageProvider$SettableContextStorage" &&
                SettableContextStorageProvider.getContextStorage.getClass.getName == "org.typelevel.otel4s.oteljava.context.IOLocalContextStorage" =>
            val storage = SettableContextStorageProvider.getContextStorage
            val local = storage.getClass.getMethod(
              "local",
              Class.forName("cats.effect.kernel.MonadCancel"),
              Class.forName("cats.effect.LiftIO")
            )
            whenPropagationEnabled(local.invoke(storage, F, L).asInstanceOf[LocalContext[F]])

          case other =>
            F.raiseError(
              new IllegalStateException(
                "IOLocalContextStorage is not configured for use as the ContextStorageProvider. " +
                  s"The current ContextStorage is: ${other.getClass.getName}"
              )
            )
        }

      private def whenPropagationEnabled[A](whenEnabled: => A): F[A] =
        if (IOLocal.isPropagating) {
          F.pure(whenEnabled)
        } else {
          F.raiseError(
            new IllegalStateException(
              "IOLocal propagation must be enabled with: -Dcats.effect.trackFiberContext=true"
            )
          )
        }
    }

}

[AlixBa]

@iRevive do you have any lead on how this should be tackled? Should it be included in IOLocalContextStorage or should I find a way to override this from a SBT perspective (and have a doc section about this)?

[iRevive]

@AlixBa what if we provide an additional implementation in the oteljava-tracing-testkit? This provider can properly handle the SettableContextStorageProvider.

To make it work, users should bind the ContextProvider as late as possible, for example, around the entry point:

class Components[F[_]](tracer: Tracer[F], ...) {
  def startHttpServer: Resource[F, Unit] = ...
}
object Components {
  def create[F[_]: LocalContextProvider](...): Resource[F, Components[F]] = ...
}

object Main extends IOApp.Simple {
  private given LocalProvider[IO, Context] = IOLocalContextStorage.localProvider[IO]
  def run: IO[Unit] = 
    Components.create[IO].flatMap(_.startHttpServer).useForever
}

Then, in tests you can use a different context storage:

class SomeTest extends ... {
  private given LocalProvider[IO, Context] = IOLocalContextTestStorage.localProvider[IO]
  ...
}

[AlixBa]

I'd have done something like this to avoid duplicating (and maintaining twice) the same code

package org.typelevel.otel4s.oteljava.testkit.context

import cats.effect.IOLocal
import cats.effect.LiftIO
import cats.effect.MonadCancelThrow
import cats.effect.std.Console
import cats.mtl.Local
import io.opentelemetry.sdk.testing.context.SettableContextStorageProvider
import org.typelevel.otel4s.context.LocalProvider
import org.typelevel.otel4s.oteljava.context.Context
import org.typelevel.otel4s.oteljava.context.IOLocalContextStorage

object IOLocalTestContextStorage {

  def localProvider[F[_]: Console: LiftIO](implicit F: MonadCancelThrow[F]): LocalProvider[F, Context] =
    new LocalProvider[F, Context] {

      val delegate = IOLocalContextStorage.localProvider[F]

      override def local: F[Local[F, Context]] =
        SettableContextStorageProvider.getContextStorage() match {
          case storage: IOLocalContextStorage => whenPropagationEnabled(storage.local)
          case _                              => delegate.local
        }

       // to extract to IOLocalContextStorage to avoid duplication
      private def whenPropagationEnabled[A](whenEnabled: => A): F[A] =
        if (IOLocal.isPropagating) {
          F.pure(whenEnabled)
        } else {
          F.raiseError(
            new IllegalStateException(
              "IOLocal propagation must be enabled with: -Dcats.effect.trackFiberContext=true"
            )
          )
        }

    }

}

but this requires to add a sbt module dependency from oteljava-trace-testkit to oteljava-context and switch from private to protected[oteljava] the IOLocalContextStorage. Is this acceptable?