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

14 files changed, 87 insertions, 50 deletions

diff --git a/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala b/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala
index 50fbc815..7564cfd0 100644
--- a/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala
+++ b/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala
@@ -88,7 +88,7 @@ class RemotingInstrumentation {
 
       val ctx = tracer.newContext(
         remoteTraceContext.getTraceName,
-        remoteTraceContext.getTraceToken,
+        Option(remoteTraceContext.getTraceToken),
         new MilliTimestamp(remoteTraceContext.getStartMilliTime()).toRelativeNanoTimestamp,
         remoteTraceContext.getIsOpen,
         isLocal = false)
diff --git a/kamon-akka-remote/src/test/scala/kamon/akka/instrumentation/RemotingInstrumentationSpec.scala b/kamon-akka-remote/src/test/scala/kamon/akka/instrumentation/RemotingInstrumentationSpec.scala
index d8535ff8..e63e06ef 100644
--- a/kamon-akka-remote/src/test/scala/kamon/akka/instrumentation/RemotingInstrumentationSpec.scala
+++ b/kamon-akka-remote/src/test/scala/kamon/akka/instrumentation/RemotingInstrumentationSpec.scala
@@ -61,7 +61,7 @@ class RemotingInstrumentationSpec extends TestKitBase with WordSpecLike with Mat
 
   "The Remoting instrumentation" should {
     "propagate the TraceContext when creating a new remote actor" in {
-      Tracer.withContext(tracer.newContext("deploy-remote-actor", "deploy-remote-actor-1")) {
+      Tracer.withContext(tracer.newContext("deploy-remote-actor", Some("deploy-remote-actor-1"))) {
         system.actorOf(TraceTokenReplier.remoteProps(Some(testActor), RemoteSystemAddress), "remote-deploy-fixture")
       }
 
@@ -71,7 +71,7 @@ class RemotingInstrumentationSpec extends TestKitBase with WordSpecLike with Mat
     "propagate the TraceContext when sending a message to a remotely deployed actor" in {
       val remoteRef = system.actorOf(TraceTokenReplier.remoteProps(None, RemoteSystemAddress), "remote-message-fixture")
 
-      Tracer.withContext(tracer.newContext("message-remote-actor", "message-remote-actor-1")) {
+      Tracer.withContext(tracer.newContext("message-remote-actor", Some("message-remote-actor-1"))) {
         remoteRef ! "reply-trace-token"
       }
 
@@ -83,7 +83,7 @@ class RemotingInstrumentationSpec extends TestKitBase with WordSpecLike with Mat
       implicit val askTimeout = Timeout(10 seconds)
       val remoteRef = system.actorOf(TraceTokenReplier.remoteProps(None, RemoteSystemAddress), "remote-ask-and-pipe-fixture")
 
-      Tracer.withContext(tracer.newContext("ask-and-pipe-remote-actor", "ask-and-pipe-remote-actor-1")) {
+      Tracer.withContext(tracer.newContext("ask-and-pipe-remote-actor", Some("ask-and-pipe-remote-actor-1"))) {
         (remoteRef ? "reply-trace-token") pipeTo (testActor)
       }
 
@@ -95,7 +95,7 @@ class RemotingInstrumentationSpec extends TestKitBase with WordSpecLike with Mat
       remoteSystem.actorOf(TraceTokenReplier.props(None), "actor-selection-target-b")
       val selection = system.actorSelection(RemoteSystemAddress + "/user/actor-selection-target-*")
 
-      Tracer.withContext(tracer.newContext("message-remote-actor-selection", "message-remote-actor-selection-1")) {
+      Tracer.withContext(tracer.newContext("message-remote-actor-selection", Some("message-remote-actor-selection-1"))) {
         selection ! "reply-trace-token"
       }
 
@@ -107,7 +107,7 @@ class RemotingInstrumentationSpec extends TestKitBase with WordSpecLike with Mat
     "propagate the TraceContext a remotely supervised child fails" in {
       val supervisor = system.actorOf(Props(new SupervisorOfRemote(testActor, RemoteSystemAddress)))
 
-      Tracer.withContext(tracer.newContext("remote-supervision", "remote-supervision-1")) {
+      Tracer.withContext(tracer.newContext("remote-supervision", Some("remote-supervision-1"))) {
         supervisor ! "fail"
       }
 
@@ -118,7 +118,7 @@ class RemotingInstrumentationSpec extends TestKitBase with WordSpecLike with Mat
       remoteSystem.actorOf(TraceTokenReplier.props(None), "remote-routee")
       val router = system.actorOf(RoundRobinGroup(List(RemoteSystemAddress + "/user/actor-selection-target-*")).props(), "router")
 
-      Tracer.withContext(tracer.newContext("remote-routee", "remote-routee-1")) {
+      Tracer.withContext(tracer.newContext("remote-routee", Some("remote-routee-1"))) {
         router ! "reply-trace-token"
       }
 
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
diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
index 5214f420..ac73bff0 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -43,7 +43,7 @@ class RequestInstrumentation {
       requestHeader.headers.toSimpleMap.find(_._1 == playExtension.traceTokenHeaderName).map(_._2)
     } else None
 
-    val newContext = token.map(t ⇒ tracer.newContext(defaultTraceName, t)).getOrElse(tracer.newContext(defaultTraceName))
+    val newContext = tracer.newContext(defaultTraceName, token)
     Tracer.setCurrentContext(newContext)
   }
 
diff --git a/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala b/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala
index 5be156cd..d76b9f58 100644
--- a/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala
@@ -46,7 +46,7 @@ class ServerRequestInstrumentation {
       request.headers.find(_.name == sprayExtension.settings.traceTokenHeaderName).map(_.value)
     } else None
 
-    val newContext = token.map(customToken ⇒ tracer.newContext(defaultTraceName, customToken)) getOrElse (tracer.newContext(defaultTraceName))
+    val newContext = tracer.newContext(defaultTraceName, token)
     Tracer.setCurrentContext(newContext)
 
     // Necessary to force initialization of traceContext when initiating the request.
diff --git a/kamon-system-metrics/src/main/scala/kamon/system/jmx/GarbageCollectionMetrics.scala b/kamon-system-metrics/src/main/scala/kamon/system/jmx/GarbageCollectionMetrics.scala
index 7a5770d8..2e9ea954 100644
--- a/kamon-system-metrics/src/main/scala/kamon/system/jmx/GarbageCollectionMetrics.scala
+++ b/kamon-system-metrics/src/main/scala/kamon/system/jmx/GarbageCollectionMetrics.scala
@@ -19,7 +19,7 @@ package kamon.system.jmx
 import java.lang.management.{ GarbageCollectorMXBean, ManagementFactory }
 
 import kamon.Kamon
-import kamon.metric.{ EntityRecorderFactory, Entity, Metrics, GenericEntityRecorder }
+import kamon.metric._
 import kamon.metric.instrument.{ DifferentialValueCollector, Time, InstrumentFactory }
 import scala.collection.JavaConverters._
 
@@ -44,7 +44,7 @@ object GarbageCollectionMetrics {
   def sanitizeCollectorName(name: String): String =
     name.replaceAll("""[^\w]""", "-").toLowerCase
 
-  def register(metricsExtension: Metrics): Unit = {
+  def register(metricsExtension: MetricsModule): Unit = {
     ManagementFactory.getGarbageCollectorMXBeans.asScala.filter(_.isValid) map { gc ⇒
       val gcName = sanitizeCollectorName(gc.getName)
       Kamon.metrics.entity(EntityRecorderFactory("system-metric", new GarbageCollectionMetrics(gc, _)), s"$gcName-garbage-collector")
diff --git a/kamon-system-metrics/src/main/scala/kamon/system/jmx/JmxSystemMetricRecorderCompanion.scala b/kamon-system-metrics/src/main/scala/kamon/system/jmx/JmxSystemMetricRecorderCompanion.scala
index 15bd399e..2a86c13c 100644
--- a/kamon-system-metrics/src/main/scala/kamon/system/jmx/JmxSystemMetricRecorderCompanion.scala
+++ b/kamon-system-metrics/src/main/scala/kamon/system/jmx/JmxSystemMetricRecorderCompanion.scala
@@ -17,10 +17,10 @@
 package kamon.system.jmx
 
 import kamon.metric.instrument.InstrumentFactory
-import kamon.metric.{ EntityRecorderFactory, EntityRecorder, Metrics }
+import kamon.metric.{ MetricsModule, EntityRecorderFactory, EntityRecorder }
 
 abstract class JmxSystemMetricRecorderCompanion(metricName: String) {
-  def register(metricsExtension: Metrics): EntityRecorder =
+  def register(metricsExtension: MetricsModule): EntityRecorder =
     metricsExtension.entity(EntityRecorderFactory("system-metric", apply(_)), metricName)
 
   def apply(instrumentFactory: InstrumentFactory): EntityRecorder
diff --git a/kamon-system-metrics/src/main/scala/kamon/system/sigar/SigarMetricsUpdater.scala b/kamon-system-metrics/src/main/scala/kamon/system/sigar/SigarMetricsUpdater.scala
index a1bd3e01..5e5f7518 100644
--- a/kamon-system-metrics/src/main/scala/kamon/system/sigar/SigarMetricsUpdater.scala
+++ b/kamon-system-metrics/src/main/scala/kamon/system/sigar/SigarMetricsUpdater.scala
@@ -19,7 +19,7 @@ package kamon.system.sigar
 import akka.actor.{ Props, Actor }
 import kamon.Kamon
 import kamon.metric.instrument.InstrumentFactory
-import kamon.metric.{ EntityRecorderFactory, Entity, EntityRecorder, Metrics }
+import kamon.metric._
 import kamon.system.sigar.SigarMetricsUpdater.UpdateSigarMetrics
 import org.hyperic.sigar.Sigar
 
@@ -65,7 +65,7 @@ trait SigarMetric extends EntityRecorder {
 }
 
 abstract class SigarMetricRecorderCompanion(metricName: String) {
-  def register(sigar: Sigar, metricsExtension: Metrics): SigarMetric =
+  def register(sigar: Sigar, metricsExtension: MetricsModule): SigarMetric =
     metricsExtension.entity(EntityRecorderFactory("system-metric", apply(sigar, _)), metricName)
 
   def apply(sigar: Sigar, instrumentFactory: InstrumentFactory): SigarMetric