try an alternative approach to active span management

5 files changed, 56 insertions, 153 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/ActiveSpan.scala b/kamon-core/src/main/scala/kamon/trace/ActiveSpan.scala
deleted file mode 100644
index b6e5d5e9..00000000
--- a/kamon-core/src/main/scala/kamon/trace/ActiveSpan.scala
+++ /dev/null
@@ -1,78 +0,0 @@
-package kamon.trace
-
-/**
-  * Wraps a [[kamon.trace.Span]] that has been activated in the current Thread. By activated we really mean, it is
-  * stored in a ThreadLocal value inside the tracer until [[kamon.trace.ActiveSpan#deactivate()]] is called.
-  *
-  * When a [[kamon.trace.Span]] is activated it will keep a reference to the previously active Span on the current
-  * Thread, take it's place as the currently active Span and put the original one once this ActiveSpan gets deactivated.
-  *
-  */
-trait ActiveSpan extends Span {
-
-  /**
-    * Sets the currently active Span to whatever Span was active when this Span was activated.
-    *
-    */
-  def deactivate(): Span
-}
-
-object ActiveSpan {
-
-  final class Default(wrappedSpan: Span, restoreOnDeactivate: ActiveSpan, tl: ThreadLocal[ActiveSpan])
-      extends ActiveSpan {
-
-    override def deactivate(): Span = {
-      tl.set(restoreOnDeactivate)
-      wrappedSpan
-    }
-
-    //
-    //  Forward all other members to the wrapped Span.
-    //
-
-    override def annotate(annotation: Span.Annotation): Span =
-      wrappedSpan.annotate(annotation)
-
-    override def addSpanTag(key: String, value: String): Span =
-      wrappedSpan.addSpanTag(key, value)
-
-    override def addSpanTag(key: String, value: Long): Span =
-      wrappedSpan.addSpanTag(key, value)
-
-    override def addSpanTag(key: String, value: Boolean): Span =
-      wrappedSpan.addSpanTag(key, value)
-
-    override def addMetricTag(key: String, value: String): Span =
-      wrappedSpan.addMetricTag(key, value)
-
-    override def addBaggage(key: String, value: String): Span =
-      wrappedSpan.addBaggage(key, value)
-
-    override def getBaggage(key: String): Option[String] =
-      wrappedSpan.getBaggage(key)
-
-    override def disableMetricsCollection(): Span =
-      wrappedSpan.disableMetricsCollection()
-
-    override def context(): SpanContext =
-      wrappedSpan.context()
-
-    override def setOperationName(operationName: String): Span =
-      wrappedSpan.setOperationName(operationName)
-
-    override def finish(finishMicros: Long): Unit =
-      wrappedSpan.finish(finishMicros)
-
-    override def capture(): Continuation =
-      wrappedSpan.capture()
-
-    override def capture(activeSpanSource: ActiveSpanSource): Continuation =
-      wrappedSpan.capture(activeSpanSource)
-  }
-
-  object Default {
-    def apply(wrappedSpan: Span, restoreOnDeactivate: ActiveSpan, tl: ThreadLocal[ActiveSpan]): Default =
-      new Default(wrappedSpan, restoreOnDeactivate, tl)
-  }
-}
\ No newline at end of file
diff --git a/kamon-core/src/main/scala/kamon/trace/ActiveSpanSource.scala b/kamon-core/src/main/scala/kamon/trace/ActiveSpanSource.scala
new file mode 100644
index 00000000..f4a363a6
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/trace/ActiveSpanSource.scala
@@ -0,0 +1,46 @@
+package kamon.trace
+
+
+trait Scope extends AutoCloseable {
+  def close(): Unit
+}
+
+trait ActiveSpanSource {
+  def activeSpan(): Span
+
+  def activate(span: Span): Scope
+  def activate(span: Span, finishOnClose: Boolean): Scope
+}
+
+object ActiveSpanSource {
+
+  final class ThreadLocalBased extends ActiveSpanSource {
+    private val emptySpan = Span.Empty(this)
+    private val storage: ThreadLocal[Span] = new ThreadLocal[Span] {
+      override def initialValue(): Span = emptySpan
+    }
+
+    override def activeSpan(): Span =
+      storage.get()
+
+    override def activate(span: Span): Scope =
+      activate(span, finishOnClose = false)
+
+    override def activate(span: Span, finishOnClose: Boolean): Scope = {
+      val previouslyActiveSpan = storage.get()
+      storage.set(span)
+
+      new Scope {
+        override def close(): Unit = {
+          storage.set(previouslyActiveSpan)
+          if (finishOnClose && span != null)
+            span.finish()
+        }
+      }
+    }
+  }
+
+  object ThreadLocalBased {
+    def apply(): ThreadLocalBased = new ThreadLocalBased()
+  }
+}
\ No newline at end of file
diff --git a/kamon-core/src/main/scala/kamon/trace/Continuation.scala b/kamon-core/src/main/scala/kamon/trace/Continuation.scala
deleted file mode 100644
index 8029b838..00000000
--- a/kamon-core/src/main/scala/kamon/trace/Continuation.scala
+++ /dev/null
@@ -1,39 +0,0 @@
-/* =========================================================================================
- * Copyright © 2013-2017 the kamon project <http://kamon.io/>
- *
- * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
- * except in compliance with the License. You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software distributed under the
- * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
- * either express or implied. See the License for the specific language governing permissions
- * and limitations under the License.
- * =========================================================================================
- */
-
-package kamon.trace
-
-
-
-trait Continuation {
-  def activate(): ActiveSpan
-}
-
-object Continuation {
-
-  /**
-    *
-    * @param span
-    * @param activeSpanSource
-    */
-  final class Default(span: Span, activeSpanSource: ActiveSpanSource) extends Continuation {
-    override def activate(): ActiveSpan =
-      activeSpanSource.makeActive(span)
-  }
-
-  object Default {
-    def apply(span: Span, activeSpanSource: ActiveSpanSource): Default = new Default(span, activeSpanSource)
-  }
-}
diff --git a/kamon-core/src/main/scala/kamon/trace/Span.scala b/kamon-core/src/main/scala/kamon/trace/Span.scala
index 113ec3de..6b38ae48 100644
--- a/kamon-core/src/main/scala/kamon/trace/Span.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Span.scala
@@ -29,10 +29,6 @@ trait BaseSpan {
 
   def context(): SpanContext
 
-  def capture(): Continuation
-
-  def capture(activeSpanSource: ActiveSpanSource): Continuation
-
   def annotate(annotation: Span.Annotation): Span
 
   def addSpanTag(key: String, value: String): Span
@@ -77,8 +73,6 @@ object Span {
 
   final class Empty(activeSpanSource: ActiveSpanSource) extends Span {
     override val context: SpanContext = SpanContext.EmptySpanContext
-    override def capture(): Continuation = Continuation.Default(this, activeSpanSource)
-    override def capture(activeSpanSource: ActiveSpanSource): Continuation = Continuation.Default(this, activeSpanSource)
 
     override def annotate(annotation: Annotation): Span = this
     override def addSpanTag(key: String, value: String): Span = this
@@ -93,7 +87,7 @@ object Span {
   }
 
   object Empty {
-    def apply(tracer: Tracer): Empty = new Empty(tracer)
+    def apply(activeSpanSource: ActiveSpanSource): Empty = new Empty(activeSpanSource)
   }
 
   /**
@@ -182,12 +176,6 @@ object Span {
       }
     }
 
-    override def capture(): Continuation =
-      Continuation.Default(this, activeSpanSource)
-
-    override def capture(activeSpanSource: ActiveSpanSource): Continuation =
-      Continuation.Default(this, activeSpanSource)
-
     private def toFinishedSpan(endTimestampMicros: Long): Span.FinishedSpan =
       Span.FinishedSpan(spanContext, operationName, startTimestampMicros, endTimestampMicros, spanTags, annotations)
 
diff --git a/kamon-core/src/main/scala/kamon/trace/Tracer.scala b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
index 71201871..737a8b8d 100644
--- a/kamon-core/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
@@ -13,7 +13,6 @@
  * =========================================================================================
  */
 
-
 package kamon.trace
 
 import java.nio.ByteBuffer
@@ -30,13 +29,7 @@ import org.slf4j.LoggerFactory
 import scala.collection.immutable
 import scala.util.Try
 
-
-trait ActiveSpanSource {
-  def activeSpan(): ActiveSpan
-  def makeActive(span: Span): ActiveSpan
-}
-
-trait Tracer extends ActiveSpanSource{
+trait Tracer extends ActiveSpanSource {
   def buildSpan(operationName: String): SpanBuilder
 
   def extract[C](format: SpanContextCodec.Format[C], carrier: C): Option[SpanContext]
@@ -48,10 +41,7 @@ object Tracer {
 
   final class Default(metrics: MetricLookup, reporterRegistry: ReporterRegistryImpl, initialConfig: Config) extends Tracer {
     private val logger = LoggerFactory.getLogger(classOf[Tracer])
-    private val emptySpan = Span.Empty(this)
-    private val activeSpanStorage: ThreadLocal[ActiveSpan] = new ThreadLocal[ActiveSpan] {
-      override def initialValue(): ActiveSpan = ActiveSpan.Default(emptySpan, null, activeSpanStorage)
-    }
+    private val activeSpanSource = ActiveSpanSource.ThreadLocalBased()
 
     private[Tracer] val tracerMetrics = new TracerMetrics(metrics)
     @volatile private[Tracer] var joinRemoteParentsWithSameSpanID: Boolean = true
@@ -83,15 +73,14 @@ object Tracer {
       case SpanContextCodec.Format.Binary      => ByteBuffer.allocate(0) // TODO: Implement binary encoding.
     }
 
-    override def activeSpan(): ActiveSpan =
-      activeSpanStorage.get()
+    override def activeSpan(): Span =
+      activeSpanSource.activeSpan()
 
-    override def makeActive(span: Span): ActiveSpan = {
-      val currentlyActiveSpan = activeSpanStorage.get()
-      val newActiveSpan = ActiveSpan.Default(span, currentlyActiveSpan, activeSpanStorage)
-      activeSpanStorage.set(newActiveSpan)
-      newActiveSpan
-    }
+    override def activate(span: Span): Scope =
+      activeSpanSource.activate(span)
+
+    override def activate(span: Span, finishOnClose: Boolean): Scope =
+      activeSpanSource.activate(span, finishOnClose)
 
     def sampler: Sampler =
       configuredSampler
@@ -232,9 +221,6 @@ object Tracer {
         baggage = SpanContext.Baggage(),
         source = Source.Local
       )
-
-    def startActive(): ActiveSpan =
-      tracer.makeActive(start())
   }
 
   private final class TracerMetrics(metricLookup: MetricLookup) {