From 172ac6dcc3ca403293f8c6e2d65c66ea0e7e8522 Mon Sep 17 00:00:00 2001 From: Dan Di Spaltro Date: Thu, 6 Jun 2024 11:12:31 -0700 Subject: [PATCH 1/2] Do a better job closing scopes --- .../zio2/ZIO2Instrumentation.scala | 20 ++++++++++++------- .../zio2/ZIO2InstrumentationSpec.scala | 6 ++++++ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala b/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala index f0fa0e41d..57fceadda 100644 --- a/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala +++ b/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala @@ -54,8 +54,10 @@ object HasStorage { } } + class NewSupervisor extends Supervisor[Any] { + override def value(implicit trace: zio.Trace): UIO[Any] = ZIO.unit override def onStart[R, E, A_]( @@ -64,21 +66,25 @@ class NewSupervisor extends Supervisor[Any] { parent: Option[Fiber.Runtime[Any, Any]], fiber: Fiber.Runtime[E, A_] )(implicit unsafe: Unsafe): Unit = { - fiber.asInstanceOf[HasContext].setContext(Kamon.currentContext()) + val fi = fiber.asInstanceOf[HasContext with HasStorage] + fi.setContext(Kamon.currentContext()) } override def onSuspend[E, A_](fiber: Fiber.Runtime[E, A_])(implicit unsafe: Unsafe): Unit = { - fiber.asInstanceOf[HasContext].setContext(Kamon.currentContext()) + val fi = fiber.asInstanceOf[HasContext with HasStorage] + fi.setContext(Kamon.currentContext()) + fi.kamonScope.close() } override def onResume[E, A_](fiber: Fiber.Runtime[E, A_])(implicit unsafe: Unsafe): Unit = { - val fiberInstance = fiber.asInstanceOf[HasContext with HasStorage] - val ctx = fiberInstance.context - fiberInstance.setKamonScope(Kamon.storeContext(ctx)) + val fi = fiber.asInstanceOf[HasContext with HasStorage] + val ctx = fi.context + fi.setKamonScope(Kamon.storeContext(ctx)) } override def onEnd[R, E, A_](value: Exit[E, A_], fiber: Fiber.Runtime[E, A_])(implicit unsafe: Unsafe): Unit = { - val fiberInstance = fiber.asInstanceOf[HasContext with HasStorage] - fiberInstance.kamonScope.close() + val fi = fiber.asInstanceOf[HasContext with HasStorage] + fi.kamonScope.close() + () } } diff --git a/instrumentation/kamon-zio-2/src/test/scala/kamon/instrumentation/zio2/ZIO2InstrumentationSpec.scala b/instrumentation/kamon-zio-2/src/test/scala/kamon/instrumentation/zio2/ZIO2InstrumentationSpec.scala index 6236e6f54..907043336 100644 --- a/instrumentation/kamon-zio-2/src/test/scala/kamon/instrumentation/zio2/ZIO2InstrumentationSpec.scala +++ b/instrumentation/kamon-zio-2/src/test/scala/kamon/instrumentation/zio2/ZIO2InstrumentationSpec.scala @@ -167,6 +167,12 @@ class ZIO2InstrumentationSpec extends AnyWordSpec with Matchers with ScalaFuture } } + override protected def afterEach(): Unit = { + super.afterEach() + + kamon.context.Storage.Debug.printNonEmptyThreads() + } + private def getKey: UIO[String] = { ZIO.succeed(Kamon.currentContext().getTag(plain("key"))) } From 4f463495b85f18a6d587fd7422bd1b0fef1e91b7 Mon Sep 17 00:00:00 2001 From: Dan Di Spaltro Date: Tue, 11 Jun 2024 08:51:05 -0700 Subject: [PATCH 2/2] Format code correctly --- .../scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala | 2 -- 1 file changed, 2 deletions(-) diff --git a/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala b/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala index 57fceadda..a2f19328e 100644 --- a/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala +++ b/instrumentation/kamon-zio-2/src/main/scala/kamon/instrumentation/zio2/ZIO2Instrumentation.scala @@ -54,10 +54,8 @@ object HasStorage { } } - class NewSupervisor extends Supervisor[Any] { - override def value(implicit trace: zio.Trace): UIO[Any] = ZIO.unit override def onStart[R, E, A_](