Experimental IOLocalContextStorage

java/common/src/main/scala/org/typelevel/otel4s/java/IOLocalContextStorage.scala

@@ -0,0 +1,36 @@
+package org.typelevel.otel4s.java
+
+import cats.effect.IOLocal
+import cats.effect.LiftIO
+import cats.effect.std.Dispatcher
+import io.opentelemetry.context.Context
+import io.opentelemetry.context.ContextStorage
+import io.opentelemetry.context.Scope
+
+class IOLocalContextStorage[F[_]: LiftIO](
+    dispatcher: Dispatcher[F],
+    ioLocal: IOLocal[Context]
+) extends ContextStorage {
+
+  override def attach(toAttach: Context): Scope =
+    dispatcher.unsafeRunSync(
+      currentOrRoot
+        .flatMap { old =>
+          ioLocal
+            .set(toAttach)
+            .as(new Scope {
+              def close() = dispatcher.unsafeRunSync(ioLocal.set(old).to[F])
+            })
+        }
+        .to[F]
+    )
+
+  override def current(): Context = {
+    dispatcher.unsafeRunSync(currentOrRoot.to[F])
+  }
+
+  private def currentOrRoot = ioLocal.get.map {
+    case null => Context.root()
+    case ctx  => ctx
+  }
+}

java/common/src/test/scala/org/typelevel/otel4s/java/ContextStorageSuite.scala

@@ -0,0 +1,43 @@
+package org.typelevel.otel4s.java
+
+import cats.effect.IO
+import cats.effect.IOLocal
+import cats.effect.Resource
+import cats.effect.std.Dispatcher
+import io.opentelemetry.context.Context
+import io.opentelemetry.context.ContextKey
+import io.opentelemetry.context.ContextStorage
+import java.util.logging._
+
+object Run extends cats.effect.IOApp.Simple {
+
+  val key = ContextKey.named[String]("test")
+
+  def run =
+    Dispatcher.parallel[IO].use { dispatcher =>
+      for {
+        _ <- IO {
+          val rootLog = Logger.getLogger("")
+          rootLog.setLevel(Level.FINE)
+          rootLog.getHandlers().head.setLevel(Level.FINE)
+        }
+        ioLocal <- IOLocal(null: Context)
+        storage = new IOLocalContextStorage(dispatcher, ioLocal)
+        _ <- IO(ContextStorage.addWrapper(_ => storage))
+        _ <- ioLocal.set(null)
+        _ <- IO.println(ContextStorage.get().getClass)
+        ctx = Context.root()
+        _ <- Resource
+          .make(IO(ctx.`with`(key, "hello").makeCurrent()))(scope =>
+            IO(scope.close())
+          )
+          .surround {
+            for {
+              key <- IO(Context.current())
+              _ <- IO.println(key)
+            } yield ()
+          }
+        _ <- IO(Option(Context.current().get(key))).flatMap(v => IO.println(v))
+      } yield ()
+    }
+}