introduce the SpanCustomizer API

8 files changed, 152 insertions, 71 deletions

diff --git a/kamon-core-tests/src/test/resources/reference.conf b/kamon-core-tests/src/test/resources/reference.conf
index efbca846..862cfa32 100644
--- a/kamon-core-tests/src/test/resources/reference.conf
+++ b/kamon-core-tests/src/test/resources/reference.conf
@@ -2,11 +2,11 @@ kamon {
 
   context.codecs {
     http-headers-keys {
-      string-broadcast-key = "kamon.context.SimpleStringCodec$Headers"
+      string-broadcast-key = "kamon.testkit.SimpleStringCodec$Headers"
     }
 
     binary-keys {
-      string-broadcast-key = "kamon.context.SimpleStringCodec$Binary"
+      string-broadcast-key = "kamon.testkit.SimpleStringCodec$Binary"
     }
   }
 }
\ No newline at end of file
diff --git a/kamon-core-tests/src/test/scala/kamon/trace/SpanCustomizerSpec.scala b/kamon-core-tests/src/test/scala/kamon/trace/SpanCustomizerSpec.scala
new file mode 100644
index 00000000..39bd7cd7
--- /dev/null
+++ b/kamon-core-tests/src/test/scala/kamon/trace/SpanCustomizerSpec.scala
@@ -0,0 +1,31 @@
+package kamon.trace
+
+import kamon.Kamon
+import kamon.context.Context
+import kamon.testkit.SpanInspection
+import org.scalatest.{Matchers, WordSpec}
+
+class SpanCustomizerSpec extends WordSpec with Matchers with SpanInspection {
+
+  "a SpanCustomizer" should {
+    "default to a Noop implementation when none is in the context" in {
+      val noopCustomizer = Context.Empty.get(SpanCustomizer.ContextKey)
+      val spanBuilder = noopCustomizer.customize(Kamon.buildSpan("noop"))
+      val span = inspect(spanBuilder.start())
+
+      span.operationName() shouldBe "noop"
+      span.metricTags() shouldBe empty
+      span.spanTags() shouldBe empty
+    }
+
+    "have a simple builder for customizing the operation name" in {
+      val operationNameCustomizer = SpanCustomizer.forOperationName("myCustomOperationName")
+      val spanBuilder = operationNameCustomizer.customize(Kamon.buildSpan("noop"))
+      val span = inspect(spanBuilder.start())
+
+      span.operationName() shouldBe "myCustomOperationName"
+      span.metricTags() shouldBe empty
+      span.spanTags() shouldBe empty
+    }
+  }
+}
diff --git a/kamon-core-tests/src/test/scala/kamon/trace/TracerSpec.scala b/kamon-core-tests/src/test/scala/kamon/trace/TracerSpec.scala
index b398ee98..3ad8a171 100644
--- a/kamon-core-tests/src/test/scala/kamon/trace/TracerSpec.scala
+++ b/kamon-core-tests/src/test/scala/kamon/trace/TracerSpec.scala
@@ -3,11 +3,11 @@ package kamon.trace
 import com.typesafe.config.ConfigFactory
 import kamon.Kamon
 import kamon.context.Context
-import kamon.testkit.{SpanBuilding, SpanInspector}
+import kamon.testkit.{SpanBuilding, SpanInspection}
 import kamon.trace.Span.TagValue
 import org.scalatest.{Matchers, OptionValues, WordSpec}
 
-class TracerSpec extends WordSpec with Matchers with SpanBuilding with OptionValues {
+class TracerSpec extends WordSpec with Matchers with SpanBuilding with SpanInspection with OptionValues {
 
   "the Kamon tracer" should {
     "construct a minimal Span that only has a operation name" in {
@@ -100,7 +100,4 @@ class TracerSpec extends WordSpec with Matchers with SpanBuilding with OptionVal
 
   val tracer: Tracer = Kamon
 
-  def inspect(span: Span): SpanInspector =
-    SpanInspector(span)
-
 }
diff --git a/kamon-core/src/main/scala/kamon/Kamon.scala b/kamon-core/src/main/scala/kamon/Kamon.scala
index f251b1ec..e83f6a5c 100644
--- a/kamon-core/src/main/scala/kamon/Kamon.scala
+++ b/kamon-core/src/main/scala/kamon/Kamon.scala
@@ -24,7 +24,7 @@ import scala.concurrent.Future
 import java.time.Duration
 import java.util.concurrent.{Executors, ScheduledExecutorService, ScheduledThreadPoolExecutor}
 
-import kamon.context.{Codecs, Context, Storage}
+import kamon.context.{Codecs, Context, Key, Storage}
 import org.slf4j.LoggerFactory
 
 import scala.util.Try
@@ -118,6 +118,9 @@ object Kamon extends MetricLookup with ReporterRegistry with Tracer {
     }
   }
 
+  def withContextKey[T, K](key: Key[K], value: K)(f: => T): T =
+    withContext(currentContext().withKey(key, value))(f)
+
 
   override def loadReportersFromConfig(): Unit =
     _reporterRegistry.loadReportersFromConfig()
diff --git a/kamon-core/src/main/scala/kamon/trace/SpanCustomizer.scala b/kamon-core/src/main/scala/kamon/trace/SpanCustomizer.scala
new file mode 100644
index 00000000..4c6220ce
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/trace/SpanCustomizer.scala
@@ -0,0 +1,35 @@
+package kamon.trace
+
+import kamon.context.Key
+import kamon.trace.Tracer.SpanBuilder
+
+/**
+  * Allows users to customize and add additional information to Spans created by instrumentation. The typical use
+  * case for SpanCustomizer instances is to provide proper operation names in situations where the instrumentation
+  * is unable to generate a reasonable enough operation name, e.g. JDBC and HTTP Client calls, instead of having a
+  * default operation name using the statement type or target host a SpanCustomizer can be provided to assign operation
+  * names like "queryUsers" or "getUserProfile" instead.
+  *
+  * Instrumentation wanting to take advantage of SpanCustomizers should look for an instance in the current context
+  * using SpanCustomizer.ContextKey.
+  *
+  */
+trait SpanCustomizer {
+  def customize(spanBuilder: SpanBuilder): SpanBuilder
+}
+
+object SpanCustomizer {
+
+  val Noop = new SpanCustomizer {
+    override def customize(spanBuilder: SpanBuilder): SpanBuilder = spanBuilder
+  }
+
+  val ContextKey = Key.local[SpanCustomizer]("span-customizer", Noop)
+
+  def forOperationName(operationName: String): SpanCustomizer = new SpanCustomizer {
+    override def customize(spanBuilder: SpanBuilder): SpanBuilder =
+      spanBuilder.withOperationName(operationName)
+  }
+}
+
+
diff --git a/kamon-core/src/main/scala/kamon/trace/Tracer.scala b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
index 5f61f3aa..8fa5df5d 100644
--- a/kamon-core/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
@@ -91,6 +91,7 @@ object Tracer {
 
   final class SpanBuilder(operationName: String, tracer: Tracer.Default, spanSink: SpanSink) {
     private var parentSpan: Span = _
+    private var initialOperationName: String = operationName
     private var startTimestamp = 0L
     private var initialSpanTags = Map.empty[String, Span.TagValue]
     private var initialMetricTags = Map.empty[String, String]
@@ -133,6 +134,17 @@ object Tracer {
       this
     }
 
+    def withOperationName(operationName: String): SpanBuilder = {
+      this.initialOperationName = operationName
+      this
+    }
+
+    def spanTags: Map[String, Span.TagValue] =
+      this.initialSpanTags
+
+    def metricTags: Map[String, String] =
+      this.initialMetricTags
+
     def ignoreParentFromContext(): SpanBuilder = {
       this.useParentFromContext = false
       this
@@ -150,7 +162,7 @@ object Tracer {
       val samplingDecision: SamplingDecision = parentSpan
         .map(_.context.samplingDecision)
         .filter(_ != SamplingDecision.Unknown)
-        .getOrElse(tracer.sampler.decide(operationName, initialSpanTags))
+        .getOrElse(tracer.sampler.decide(initialOperationName, initialSpanTags))
 
       val spanContext = parentSpan match {
         case Some(parent) => joinParentContext(parent, samplingDecision)
@@ -158,7 +170,7 @@ object Tracer {
       }
 
       tracer.tracerMetrics.createdSpans.increment()
-      Span.Local(spanContext, nonRemoteParent, operationName, initialSpanTags, initialMetricTags, startTimestampMicros, spanSink, tracer.scopeSpanMetrics)
+      Span.Local(spanContext, nonRemoteParent, initialOperationName, initialSpanTags, initialMetricTags, startTimestampMicros, spanSink, tracer.scopeSpanMetrics)
     }
 
     private def joinParentContext(parent: Span, samplingDecision: SamplingDecision): SpanContext =
diff --git a/kamon-testkit/src/main/scala/kamon/testkit/SpanInspection.scala b/kamon-testkit/src/main/scala/kamon/testkit/SpanInspection.scala
new file mode 100644
index 00000000..c4a83a73
--- /dev/null
+++ b/kamon-testkit/src/main/scala/kamon/testkit/SpanInspection.scala
@@ -0,0 +1,64 @@
+package kamon.testkit
+
+import kamon.trace.{Span, SpanContext}
+import kamon.trace.Span.FinishedSpan
+import kamon.util.Clock
+
+import scala.reflect.ClassTag
+import scala.util.Try
+
+trait SpanInspection {
+
+  def inspect(span: Span): SpanInspection.Inspector =
+    new SpanInspection.Inspector(span)
+}
+
+object SpanInspection {
+
+  class Inspector(span: Span) {
+    private val (realSpan, spanData) = Try {
+      val realSpan = span match {
+        case _: Span.Local => span
+      }
+
+      val spanData = invoke[Span.Local, FinishedSpan](realSpan, "toFinishedSpan", classOf[Long] -> Long.box(Clock.microTimestamp()))
+      (realSpan, spanData)
+    }.getOrElse((null, null))
+
+    def isEmpty: Boolean =
+      realSpan == null
+
+    def spanTag(key: String): Option[Span.TagValue] =
+      spanData.tags.get(key)
+
+    def spanTags(): Map[String, Span.TagValue] =
+      spanData.tags
+
+    def metricTags(): Map[String, String] =
+      getField[Span.Local, Map[String, String]](realSpan, "customMetricTags")
+
+    def startTimestamp(): Long =
+      getField[Span.Local, Long](realSpan, "startTimestampMicros")
+
+    def context(): SpanContext =
+      spanData.context
+
+    def operationName(): String =
+      spanData.operationName
+
+
+    private def getField[T, R](target: Any, fieldName: String)(implicit classTag: ClassTag[T]): R = {
+      val toFinishedSpanMethod = classTag.runtimeClass.getDeclaredFields.find(_.getName.contains(fieldName)).get
+      toFinishedSpanMethod.setAccessible(true)
+      toFinishedSpanMethod.get(target).asInstanceOf[R]
+    }
+
+    private def invoke[T, R](target: Any, fieldName: String, parameters: (Class[_], AnyRef)*)(implicit classTag: ClassTag[T]): R = {
+      val parameterClasses = parameters.map(_._1)
+      val parameterInstances = parameters.map(_._2)
+      val toFinishedSpanMethod = classTag.runtimeClass.getDeclaredMethod(fieldName, parameterClasses: _*)
+      toFinishedSpanMethod.setAccessible(true)
+      toFinishedSpanMethod.invoke(target, parameterInstances: _*).asInstanceOf[R]
+    }
+  }
+}
\ No newline at end of file
diff --git a/kamon-testkit/src/main/scala/kamon/testkit/SpanInspector.scala b/kamon-testkit/src/main/scala/kamon/testkit/SpanInspector.scala
deleted file mode 100644
index 25f7ee06..00000000
--- a/kamon-testkit/src/main/scala/kamon/testkit/SpanInspector.scala
+++ /dev/null
@@ -1,61 +0,0 @@
-package kamon.testkit
-
-import kamon.trace.{Span, SpanContext}
-import kamon.trace.Span.FinishedSpan
-import kamon.util.Clock
-
-import scala.reflect.ClassTag
-import scala.util.Try
-
-class SpanInspector(span: Span) {
-  private val (realSpan, spanData) = Try {
-    val realSpan = span match {
-      case _: Span.Local => span
-    }
-
-    val spanData = invoke[Span.Local, FinishedSpan](realSpan, "toFinishedSpan", classOf[Long] -> Long.box(Clock.microTimestamp()))
-    (realSpan, spanData)
-  }.getOrElse((null, null))
-
-  def isEmpty: Boolean =
-    realSpan == null
-
-  def spanTag(key: String): Option[Span.TagValue] =
-    spanData.tags.get(key)
-
-  def spanTags(): Map[String, Span.TagValue] =
-    spanData.tags
-
-  def metricTags(): Map[String, String] =
-    getField[Span.Local, Map[String, String]](realSpan, "customMetricTags")
-
-  def startTimestamp(): Long =
-    getField[Span.Local, Long](realSpan, "startTimestampMicros")
-
-  def context(): SpanContext =
-    spanData.context
-
-  def operationName(): String =
-    spanData.operationName
-
-
-
-
-  private def getField[T, R](target: Any, fieldName: String)(implicit classTag: ClassTag[T]): R = {
-    val toFinishedSpanMethod = classTag.runtimeClass.getDeclaredFields.find(_.getName.contains(fieldName)).get
-    toFinishedSpanMethod.setAccessible(true)
-    toFinishedSpanMethod.get(target).asInstanceOf[R]
-  }
-
-  private def invoke[T, R](target: Any, fieldName: String, parameters: (Class[_], AnyRef)*)(implicit classTag: ClassTag[T]): R = {
-    val parameterClasses = parameters.map(_._1)
-    val parameterInstances = parameters.map(_._2)
-    val toFinishedSpanMethod = classTag.runtimeClass.getDeclaredMethod(fieldName, parameterClasses: _*)
-    toFinishedSpanMethod.setAccessible(true)
-    toFinishedSpanMethod.invoke(target, parameterInstances: _*).asInstanceOf[R]
-  }
-}
-
-object SpanInspector {
-  def apply(span: Span): SpanInspector = new SpanInspector(span)
-}