1 files changed, 17 insertions, 0 deletions
diff --git a/src/main/scala/xyz/driver/core/trace/SimpleSpanContextHandler.scala b/src/main/scala/xyz/driver/core/trace/SimpleSpanContextHandler.scala
new file mode 100644
index 0000000..4ea0e8c
--- /dev/null
+++ b/src/main/scala/xyz/driver/core/trace/SimpleSpanContextHandler.scala
@@ -0,0 +1,17 @@
+package xyz.driver.core.trace
+
+import com.google.cloud.trace.SpanContextHandler
+import com.google.cloud.trace.DetachedSpanContextHandle
+import com.google.cloud.trace.core.{SpanContext, SpanContextHandle}
+
+@SuppressWarnings(Array("org.wartremover.warts.Var"))
+class SimpleSpanContextHandler(rootSpan: SpanContext) extends SpanContextHandler {
+  private var currentSpanContext = rootSpan
+
+  override def current(): SpanContext = currentSpanContext
+
+  override def attach(context: SpanContext): SpanContextHandle = {
+    currentSpanContext = context
+    new DetachedSpanContextHandle(context)
+  }
+}