return the wrapped Span upon deactivation of a ActiveSpan

1 files changed, 4 insertions, 2 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/ActiveSpan.scala b/kamon-core/src/main/scala/kamon/trace/ActiveSpan.scala
index 3a46d94f..3753cd7e 100644
--- a/kamon-core/src/main/scala/kamon/trace/ActiveSpan.scala
+++ b/kamon-core/src/main/scala/kamon/trace/ActiveSpan.scala
@@ -14,7 +14,7 @@ trait ActiveSpan extends Span {
     * Sets the currently active Span to whatever Span was active when this Span was activated.
     *
     */
-  def deactivate(): Unit
+  def deactivate(): Span
 }
 
 object ActiveSpan {
@@ -22,8 +22,10 @@ object ActiveSpan {
   final class Default(wrappedSpan: Span, restoreOnDeactivate: ActiveSpan, tl: ThreadLocal[ActiveSpan])
       extends ActiveSpan {
 
-    override def deactivate(): Unit =
+    override def deactivate(): Span = {
       tl.set(restoreOnDeactivate)
+      wrappedSpan
+    }
 
     //
     //  Forward all other members to the wrapped Span.