= core: add commodity trace context manipulation methods to the Tracer.

7 files changed, 72 insertions, 35 deletions

diff --git a/kamon-core/src/main/scala/kamon/Kamon.scala b/kamon-core/src/main/scala/kamon/Kamon.scala
index f8253875..d819588c 100644
--- a/kamon-core/src/main/scala/kamon/Kamon.scala
+++ b/kamon-core/src/main/scala/kamon/Kamon.scala
@@ -18,13 +18,12 @@ import _root_.akka.actor
 import _root_.akka.actor._
 import com.typesafe.config.{ ConfigFactory, Config }
 import kamon.metric._
-import kamon.metric.instrument.Gauge
-import kamon.trace.{ TracerImpl, Tracer }
+import kamon.trace.{ TracerModuleImpl, TracerModule }
 
 object Kamon {
   trait Extension extends actor.Extension
 
-  private case class KamonCoreComponents(metrics: Metrics, tracer: Tracer)
+  private case class KamonCoreComponents(metrics: MetricsModule, tracer: TracerModule)
 
   @volatile private var _system: ActorSystem = _
   @volatile private var _coreComponents: Option[KamonCoreComponents] = None
@@ -40,8 +39,8 @@ object Kamon {
     }
 
     if (_coreComponents.isEmpty) {
-      val metrics = MetricsImpl(config)
-      val tracer = TracerImpl(metrics, config)
+      val metrics = MetricsModuleImpl(config)
+      val tracer = TracerModuleImpl(metrics, config)
 
       _coreComponents = Some(KamonCoreComponents(metrics, tracer))
       _system = ActorSystem("kamon", resolveInternalConfig)
@@ -62,10 +61,10 @@ object Kamon {
     _system = null
   }
 
-  def metrics: Metrics =
+  def metrics: MetricsModule =
     ifStarted(_.metrics)
 
-  def tracer: Tracer =
+  def tracer: TracerModule =
     ifStarted(_.tracer)
 
   def apply[T <: Kamon.Extension](key: ExtensionId[T]): T =
diff --git a/kamon-core/src/main/scala/kamon/metric/Metrics.scala b/kamon-core/src/main/scala/kamon/metric/MetricsModule.scala
index 9fd9e771..9f2bbbef 100644
--- a/kamon-core/src/main/scala/kamon/metric/Metrics.scala
+++ b/kamon-core/src/main/scala/kamon/metric/MetricsModule.scala
@@ -29,7 +29,7 @@ import scala.concurrent.duration.FiniteDuration
 
 case class EntityRegistration[T <: EntityRecorder](entity: Entity, recorder: T)
 
-trait Metrics {
+trait MetricsModule {
   def settings: MetricsSettings
 
   def shouldTrack(entity: Entity): Boolean
@@ -228,7 +228,7 @@ trait Metrics {
   def instrumentFactory(category: String): InstrumentFactory
 }
 
-private[kamon] class MetricsImpl(config: Config) extends Metrics {
+private[kamon] class MetricsModuleImpl(config: Config) extends MetricsModule {
   import kamon.util.TriemapAtomicGetOrElseUpdate.Syntax
 
   private val _trackedEntities = TrieMap.empty[Entity, EntityRecorder]
@@ -356,9 +356,9 @@ private[kamon] class MetricsImpl(config: Config) extends Metrics {
   }
 }
 
-private[kamon] object MetricsImpl {
+private[kamon] object MetricsModuleImpl {
 
   def apply(config: Config) =
-    new MetricsImpl(config)
+    new MetricsModuleImpl(config)
 }
 
diff --git a/kamon-core/src/main/scala/kamon/metric/SubscriptionsDispatcher.scala b/kamon-core/src/main/scala/kamon/metric/SubscriptionsDispatcher.scala
index 3b6be70b..9f8c7be3 100644
--- a/kamon-core/src/main/scala/kamon/metric/SubscriptionsDispatcher.scala
+++ b/kamon-core/src/main/scala/kamon/metric/SubscriptionsDispatcher.scala
@@ -24,7 +24,7 @@ import scala.concurrent.duration.FiniteDuration
 /**
  *  Manages subscriptions to metrics and dispatch snapshots on every tick to all subscribers.
  */
-private[kamon] class SubscriptionsDispatcher(interval: FiniteDuration, metricsExtension: MetricsImpl) extends Actor {
+private[kamon] class SubscriptionsDispatcher(interval: FiniteDuration, metricsExtension: MetricsModuleImpl) extends Actor {
   var lastTick = MilliTimestamp.now
   var oneShotSubscriptions = Map.empty[ActorRef, SubscriptionFilter]
   var permanentSubscriptions = Map.empty[ActorRef, SubscriptionFilter]
@@ -81,7 +81,7 @@ private[kamon] class SubscriptionsDispatcher(interval: FiniteDuration, metricsEx
 }
 
 object SubscriptionsDispatcher {
-  def props(interval: FiniteDuration, metricsExtension: MetricsImpl): Props =
+  def props(interval: FiniteDuration, metricsExtension: MetricsModuleImpl): Props =
     Props(new SubscriptionsDispatcher(interval, metricsExtension))
 
   case object Tick
diff --git a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
index 0f09b4be..17be661b 100644
--- a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
@@ -20,7 +20,7 @@ import java.util.concurrent.ConcurrentLinkedQueue
 
 import akka.event.LoggingAdapter
 import kamon.Kamon
-import kamon.metric.{ SegmentMetrics, Metrics, TraceMetrics }
+import kamon.metric.{ SegmentMetrics, MetricsModule, TraceMetrics }
 import kamon.util.{ NanoInterval, RelativeNanoTimestamp }
 
 import scala.annotation.tailrec
diff --git a/kamon-core/src/main/scala/kamon/trace/Tracer.scala b/kamon-core/src/main/scala/kamon/trace/TracerModule.scala
index 472c1d65..416af20e 100644
--- a/kamon-core/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TracerModule.scala
@@ -21,15 +21,16 @@ import java.util.concurrent.atomic.AtomicLong
 
 import akka.actor._
 import com.typesafe.config.Config
-import kamon.metric.Metrics
+import kamon.Kamon
+import kamon.metric.MetricsModule
 import kamon.util._
 
 import scala.util.Try
 
-trait Tracer {
+trait TracerModule {
   def newContext(name: String): TraceContext
-  def newContext(name: String, token: String): TraceContext
-  def newContext(name: String, token: String, timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext
+  def newContext(name: String, token: Option[String]): TraceContext
+  def newContext(name: String, token: Option[String], timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext
 
   def subscribe(subscriber: ActorRef): Unit
   def unsubscribe(subscriber: ActorRef): Unit
@@ -55,9 +56,45 @@ object Tracer {
 
     try code finally _traceContextStorage.set(oldContext)
   }
+
+  // Java variant.
+  def withContext[T](context: TraceContext, code: Supplier[T]): T =
+    withContext(context)(code.get)
+
+  def withNewContext[T](traceName: String, traceToken: Option[String], autoFinish: Boolean)(code: ⇒ T): T = {
+    withContext(Kamon.tracer.newContext(traceName, traceToken)) {
+      val codeResult = code
+      if (autoFinish)
+        currentContext.finish()
+
+      codeResult
+    }
+  }
+
+  def withNewContext[T](traceName: String)(code: ⇒ T): T =
+    withNewContext(traceName, None, false)(code)
+
+  def withNewContext[T](traceName: String, traceToken: Option[String])(code: ⇒ T): T =
+    withNewContext(traceName, traceToken, false)(code)
+
+  def withNewContext[T](traceName: String, autoFinish: Boolean)(code: ⇒ T): T =
+    withNewContext(traceName, None, autoFinish)(code)
+
+  // Java variants.
+  def withNewContext[T](traceName: String, traceToken: Option[String], autoFinish: Boolean, code: Supplier[T]): T =
+    withNewContext(traceName, traceToken, autoFinish)(code.get)
+
+  def withNewContext[T](traceName: String, code: Supplier[T]): T =
+    withNewContext(traceName, None, false)(code.get)
+
+  def withNewContext[T](traceName: String, traceToken: Option[String], code: Supplier[T]): T =
+    withNewContext(traceName, traceToken, false)(code.get)
+
+  def withNewContext[T](traceName: String, autoFinish: Boolean, code: Supplier[T]): T =
+    withNewContext(traceName, None, autoFinish)(code.get)
 }
 
-private[kamon] class TracerImpl(metricsExtension: Metrics, config: Config) extends Tracer {
+private[kamon] class TracerModuleImpl(metricsExtension: MetricsModule, config: Config) extends TracerModule {
   private val _settings = TraceSettings(config)
   private val _hostnamePrefix = Try(InetAddress.getLocalHost.getHostName).getOrElse("unknown-localhost")
   private val _tokenCounter = new AtomicLong
@@ -69,26 +106,29 @@ private[kamon] class TracerImpl(metricsExtension: Metrics, config: Config) exten
     _hostnamePrefix + "-" + String.valueOf(_tokenCounter.incrementAndGet())
 
   def newContext(name: String): TraceContext =
-    createTraceContext(name)
+    createTraceContext(name, None)
 
-  def newContext(name: String, token: String): TraceContext =
+  def newContext(name: String, token: Option[String]): TraceContext =
     createTraceContext(name, token)
 
-  def newContext(name: String, token: String, timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext =
+  def newContext(name: String, token: Option[String], timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext =
     createTraceContext(name, token, timestamp, isOpen, isLocal)
 
-  private def createTraceContext(traceName: String, token: String = newToken, startTimestamp: RelativeNanoTimestamp = RelativeNanoTimestamp.now,
+  private def createTraceContext(traceName: String, token: Option[String], startTimestamp: RelativeNanoTimestamp = RelativeNanoTimestamp.now,
     isOpen: Boolean = true, isLocal: Boolean = true): TraceContext = {
 
-    def newMetricsOnlyContext = new MetricsOnlyContext(traceName, token, isOpen, _settings.levelOfDetail, startTimestamp, null)
+    def newMetricsOnlyContext(token: String): TraceContext =
+      new MetricsOnlyContext(traceName, token, isOpen, _settings.levelOfDetail, startTimestamp, null)
+
+    val traceToken = token.getOrElse(newToken)
 
     if (_settings.levelOfDetail == LevelOfDetail.MetricsOnly || !isLocal)
-      newMetricsOnlyContext
+      newMetricsOnlyContext(traceToken)
     else {
       if (!_settings.sampler.shouldTrace)
-        newMetricsOnlyContext
+        newMetricsOnlyContext(traceToken)
       else
-        new TracingContext(traceName, token, true, _settings.levelOfDetail, isLocal, startTimestamp, null, dispatchTracingContext)
+        new TracingContext(traceName, traceToken, true, _settings.levelOfDetail, isLocal, startTimestamp, null, dispatchTracingContext)
     }
   }
 
@@ -122,10 +162,10 @@ private[kamon] class TracerImpl(metricsExtension: Metrics, config: Config) exten
   }
 }
 
-private[kamon] object TracerImpl {
+private[kamon] object TracerModuleImpl {
 
-  def apply(metricsExtension: Metrics, config: Config) =
-    new TracerImpl(metricsExtension, config)
+  def apply(metricsExtension: MetricsModule, config: Config) =
+    new TracerModuleImpl(metricsExtension, config)
 }
 
 case class TraceInfo(name: String, token: String, timestamp: NanoTimestamp, elapsedTime: NanoInterval, metadata: Map[String, String], segments: List[SegmentInfo])
diff --git a/kamon-core/src/main/scala/kamon/trace/TracingContext.scala b/kamon-core/src/main/scala/kamon/trace/TracingContext.scala
index 9708d25f..9269a99e 100644
--- a/kamon-core/src/main/scala/kamon/trace/TracingContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TracingContext.scala
@@ -21,7 +21,7 @@ import java.util.concurrent.atomic.AtomicInteger
 
 import akka.event.LoggingAdapter
 import kamon.util.{ NanoInterval, RelativeNanoTimestamp, NanoTimestamp }
-import kamon.metric.Metrics
+import kamon.metric.MetricsModule
 
 import scala.collection.concurrent.TrieMap
 
diff --git a/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala b/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala
index e7b18770..cb909ad9 100644
--- a/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala
+++ b/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala
@@ -20,13 +20,11 @@ import akka.testkit.{ ImplicitSender, TestKitBase }
 import akka.actor.ActorSystem
 import com.typesafe.config.{ Config, ConfigFactory }
 import kamon.Kamon
-import kamon.metric.{ Entity, SubscriptionsDispatcher, EntitySnapshot, MetricsImpl }
+import kamon.metric.{ Entity, SubscriptionsDispatcher, EntitySnapshot }
 import kamon.trace.TraceContext
 import kamon.util.LazyActorRef
 import org.scalatest.{ BeforeAndAfterAll, Matchers, WordSpecLike }
 
-import scala.reflect.ClassTag
-
 abstract class BaseKamonSpec(actorSystemName: String) extends TestKitBase with WordSpecLike with Matchers with ImplicitSender with BeforeAndAfterAll {
   lazy val collectionContext = Kamon.metrics.buildDefaultCollectionContext
   implicit lazy val system: ActorSystem = {
@@ -41,7 +39,7 @@ abstract class BaseKamonSpec(actorSystemName: String) extends TestKitBase with W
     Kamon.tracer.newContext(name)
 
   def newContext(name: String, token: String): TraceContext =
-    Kamon.tracer.newContext(name, token)
+    Kamon.tracer.newContext(name, Option(token))
 
   def takeSnapshotOf(name: String, category: String): EntitySnapshot = {
     val recorder = Kamon.metrics.find(name, category).get