From ac4175c2d9830dbb08a23a25497062c41d266295 Mon Sep 17 00:00:00 2001
From: "Ross A. Baker" <ross@rossabaker.com>
Date: Fri, 12 May 2023 12:09:34 -0400
Subject: [PATCH] Experimental IOLocalContextStorage

---
 .../otel4s/java/IOLocalContextStorage.scala   | 36 ++++++++++++++++
 .../otel4s/java/ContextStorageSuite.scala     | 43 +++++++++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 java/common/src/main/scala/org/typelevel/otel4s/java/IOLocalContextStorage.scala
 create mode 100644 java/common/src/test/scala/org/typelevel/otel4s/java/ContextStorageSuite.scala

diff --git a/java/common/src/main/scala/org/typelevel/otel4s/java/IOLocalContextStorage.scala b/java/common/src/main/scala/org/typelevel/otel4s/java/IOLocalContextStorage.scala
new file mode 100644
index 000000000..b4cf1f6b2
--- /dev/null
+++ b/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
+  }
+}
diff --git a/java/common/src/test/scala/org/typelevel/otel4s/java/ContextStorageSuite.scala b/java/common/src/test/scala/org/typelevel/otel4s/java/ContextStorageSuite.scala
new file mode 100644
index 000000000..7c71b26e1
--- /dev/null
+++ b/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 ()
+    }
+}