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..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 @@ -64,21 +64,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"))) }