Switched to DynamicVariables, solve context passing issue produced by runnable batching

1 files changed, 22 insertions, 19 deletions

diff --git a/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala b/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
index 789c7c77..6010a185 100644
--- a/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
@@ -23,39 +23,41 @@ class RunnableInstrumentationSpec extends WordSpec with Matchers with ScalaFutur
             }
         }
 
-        "should be available during the execution of onComplete callbacks" in { new FutureWithContextFixture {
-            val onCompleteContext = Promise[TraceContext]()
+        "should be available during the execution of onComplete callbacks" in new FutureWithContextFixture {
 
+          val onCompleteContext = Promise[Option[TraceContext]]()
+
+          Tracer.traceContext.withValue(Some(testContext)) {
             futureWithContext.onComplete({
-              case _ => onCompleteContext.complete(Success(Tracer.context.get))
+              case _ => println("Completing second promise from: "+Thread.currentThread().getName + " With Context: " + Tracer.traceContext.value); onCompleteContext.complete(Success(Tracer.traceContext.value))
             })
+          }
 
-            whenReady(onCompleteContext.future) { result =>
-              result should equal(testContext)
-            }
-        }}
+          whenReady(onCompleteContext.future) { result =>
+            result should equal(Some(testContext))
+          }
+        }
       }
     }
 
     "created in a thread that doest have a TraceContext" must {
-      "not capture any TraceContext for the body execution" in { new FutureWithoutContextFixture{
-
+      "not capture any TraceContext for the body execution" in new FutureWithoutContextFixture{
           whenReady(futureWithoutContext) { result =>
             result should equal(None)
           }
-      }}
+      }
 
-      "not make any TraceContext available during the onComplete callback" in { new FutureWithoutContextFixture {
+      "not make any TraceContext available during the onComplete callback" in new FutureWithoutContextFixture {
         val onCompleteContext = Promise[Option[TraceContext]]()
 
-        futureWithoutContext.onComplete({
-          case _ => onCompleteContext.complete(Success(Tracer.context))
-        })
+        futureWithoutContext.onComplete {
+          case _ => onCompleteContext.complete(Success(Tracer.traceContext.value))
+        }
 
         whenReady(onCompleteContext.future) { result =>
           result should equal(None)
         }
-      }}
+      }
     }
   }
 
@@ -68,14 +70,15 @@ class RunnableInstrumentationSpec extends WordSpec with Matchers with ScalaFutur
 
   class FutureWithContextFixture {
     val testContext = TraceContext()
-    Tracer.set(testContext)
 
-    val futureWithContext = Future { Tracer.context }
+    var futureWithContext: Future[Option[TraceContext]] = _
+    Tracer.traceContext.withValue(Some(testContext)) {
+      futureWithContext = Future { Tracer.traceContext.value }
+    }
   }
 
   trait FutureWithoutContextFixture {
-    Tracer.clear // Make sure no TraceContext is available
-    val futureWithoutContext = Future { Tracer.context }
+    val futureWithoutContext = Future { Tracer.traceContext.value }
   }
 }