get some basic sampling going

1 files changed, 77 insertions, 55 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/Tracer.scala b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
index 84aafe68..6bb5a252 100644
--- a/kamon-core/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
@@ -1,27 +1,40 @@
 package kamon.trace
 
-import java.util.concurrent.atomic.AtomicLong
+import java.util.concurrent.ThreadLocalRandom
 
-import io.opentracing.propagation.Format
+import com.typesafe.scalalogging.Logger
+import io.opentracing.propagation.{TextMap, Format}
+import io.opentracing.propagation.Format.Builtin.{BINARY, HTTP_HEADERS, TEXT_MAP}
 import io.opentracing.util.ThreadLocalActiveSpanSource
 import kamon.ReporterRegistryImpl
-import kamon.metric.RecorderRegistry
+import kamon.metric.{Entity, EntityRecorder, RecorderRegistry}
 import kamon.util.Clock
 
 class Tracer(metrics: RecorderRegistry, reporterRegistry: ReporterRegistryImpl) extends io.opentracing.Tracer {
-  private val traceCounter = new AtomicLong()
-  private val spanCounter = new AtomicLong()
+  private val logger = Logger(classOf[Tracer])
+  private val metricsRecorder = new TracerMetricsRecorder(metrics.getRecorder(Entity("tracer", "tracer", Map.empty)))
   private val activeSpanSource = new ThreadLocalActiveSpanSource()
 
+  @volatile private var sampler: Sampler = Sampler.never
+  @volatile private var textMapSpanContextCodec = SpanContextCodec.TextMap
+  @volatile private var httpHeaderSpanContextCodec = SpanContextCodec.ZipkinB3
 
   override def buildSpan(operationName: String): io.opentracing.Tracer.SpanBuilder =
-    new SpanBuilder(operationName, spanCounter.incrementAndGet())
+    new SpanBuilder(operationName)
 
-  override def extract[C](format: Format[C], carrier: C): io.opentracing.SpanContext =
-    sys.error("Extracting not implemented yet.")
+  override def extract[C](format: Format[C], carrier: C): io.opentracing.SpanContext = format match {
+    case HTTP_HEADERS => httpHeaderSpanContextCodec.extract(carrier.asInstanceOf[TextMap], sampler)
+    case TEXT_MAP     => textMapSpanContextCodec.extract(carrier.asInstanceOf[TextMap], sampler)
+    case BINARY       => null // TODO: Implement Binary Encoding
+    case _            => null
+  }
 
-  override def inject[C](spanContext: io.opentracing.SpanContext, format: Format[C], carrier: C): Unit =
-    sys.error("Injecting not implemented yet.")
+  override def inject[C](spanContext: io.opentracing.SpanContext, format: Format[C], carrier: C): Unit = format match {
+    case HTTP_HEADERS => httpHeaderSpanContextCodec.inject(spanContext.asInstanceOf[SpanContext], carrier.asInstanceOf[TextMap])
+    case TEXT_MAP     => textMapSpanContextCodec.inject(spanContext.asInstanceOf[SpanContext], carrier.asInstanceOf[TextMap])
+    case BINARY       =>
+    case _            =>
+  }
 
   override def activeSpan(): io.opentracing.ActiveSpan =
     activeSpanSource.activeSpan()
@@ -29,82 +42,91 @@ class Tracer(metrics: RecorderRegistry, reporterRegistry: ReporterRegistryImpl)
   override def makeActive(span: io.opentracing.Span): io.opentracing.ActiveSpan =
     activeSpanSource.makeActive(span)
 
+  def setTextMapSpanContextCodec(codec: SpanContextCodec[TextMap]): Unit =
+    this.textMapSpanContextCodec = codec
 
-  private[kamon] def newTraceID: Long =
-    traceCounter.incrementAndGet()
+  def setHttpHeaderSpanContextCodec(codec: SpanContextCodec[TextMap]): Unit =
+    this.httpHeaderSpanContextCodec = codec
 
-  private class SpanBuilder(operationName: String, spanID: Long) extends io.opentracing.Tracer.SpanBuilder {
-    private var traceID = 0L
+  private class SpanBuilder(operationName: String) extends io.opentracing.Tracer.SpanBuilder {
+    private var parentContext: SpanContext = _
     private var startTimestamp = 0L
-    private var parentID = 0L
     private var initialTags = Map.empty[String, String]
+    private var useActiveSpanAsParent = true
 
-    override def start(): io.opentracing.Span =
-      startManual()
-
-    override def asChildOf(parent: io.opentracing.SpanContext): io.opentracing.Tracer.SpanBuilder = {
-      parent match {
-        case kamonSpanContext: kamon.trace.SpanContext =>
-          traceID = kamonSpanContext.traceID
-          parentID = kamonSpanContext.spanID
-        case _ => sys.error("Can't extract the parent ID from a non-kamon SpanContext")
-      }
-      this
+    override def asChildOf(parent: io.opentracing.SpanContext): io.opentracing.Tracer.SpanBuilder = parent match {
+      case spanContext: kamon.trace.SpanContext =>
+        this.parentContext = spanContext
+        this
+      case _ => logger.error("Can't extract the parent ID from a non-Kamon SpanContext"); this
     }
 
-    override def asChildOf(parent: io.opentracing.BaseSpan[_]): io.opentracing.Tracer.SpanBuilder = {
-      parent.context() match {
-        case kamonSpanContext: kamon.trace.SpanContext =>
-          traceID = kamonSpanContext.traceID
-          parentID = kamonSpanContext.spanID
-        case _ => sys.error("Can't extract the parent ID from a non-kamon SpanContext")
-      }
-      this
-    }
+    override def asChildOf(parent: io.opentracing.BaseSpan[_]): io.opentracing.Tracer.SpanBuilder =
+      asChildOf(parent.context())
 
     override def addReference(referenceType: String, referencedContext: io.opentracing.SpanContext): io.opentracing.Tracer.SpanBuilder = {
       if(referenceType != null && referenceType.equals(io.opentracing.References.CHILD_OF)) {
-        referencedContext match {
-          case kamonSpanContext: kamon.trace.SpanContext =>
-            traceID = kamonSpanContext.traceID
-            parentID = kamonSpanContext.spanID
-          case _ => sys.error("Can't extract the parent ID from a non-kamon SpanContext")
-        }
-      }
-      this
+        asChildOf(referencedContext)
+      } else this
     }
 
     override def withTag(key: String, value: String): io.opentracing.Tracer.SpanBuilder = {
-      initialTags = initialTags + (key -> value)
+      this.initialTags = this.initialTags + (key -> value)
       this
     }
 
     override def withTag(key: String, value: Boolean): io.opentracing.Tracer.SpanBuilder = {
-      initialTags = initialTags + (key -> value.toString)
+      this.initialTags = this.initialTags + (key -> value.toString)
       this
     }
 
     override def withTag(key: String, value: Number): io.opentracing.Tracer.SpanBuilder = {
-      initialTags = initialTags + (key -> value.toString)
+      this.initialTags = this.initialTags + (key -> value.toString)
       this
     }
 
-    override def startManual(): Span = {
-      if(traceID == 0L) traceID = Tracer.this.newTraceID
-      val startTimestampMicros = if(startTimestamp != 0L) startTimestamp else Clock.microTimestamp()
-      new Span(new SpanContext(traceID, spanID, parentID), operationName, startTimestampMicros, metrics, reporterRegistry)
+    override def withStartTimestamp(microseconds: Long): io.opentracing.Tracer.SpanBuilder = {
+      this.startTimestamp = microseconds
+      this
     }
 
-    override def withStartTimestamp(microseconds: Long): io.opentracing.Tracer.SpanBuilder = {
-      startTimestamp = microseconds
+    override def ignoreActiveSpan(): io.opentracing.Tracer.SpanBuilder = {
+      this.useActiveSpanAsParent = false
       this
     }
 
-    override def startActive(): io.opentracing.ActiveSpan = {
-      Tracer.this.makeActive(startManual())
+    override def start(): io.opentracing.Span =
+      startManual()
+
+    override def startActive(): io.opentracing.ActiveSpan =
+      makeActive(startManual())
+
+    override def startManual(): Span = {
+      val startTimestampMicros = if(startTimestamp != 0L) startTimestamp else Clock.microTimestamp()
+
+      if(parentContext == null && useActiveSpanAsParent) {
+        val possibleParent = activeSpan()
+        if(possibleParent != null)
+          parentContext = possibleParent.context().asInstanceOf[SpanContext]
+      }
+
+      val spanContext =
+        if(parentContext != null)
+          new SpanContext(parentContext.traceID, createID(), parentContext.spanID, parentContext.sampled, initialTags)
+        else {
+          val traceID = createID()
+          new SpanContext(traceID, traceID, 0L, sampler.decide(traceID), initialTags)
+        }
+
+      metricsRecorder.createdSpans.increment()
+      new Span(spanContext, operationName, initialTags, startTimestampMicros, metrics, reporterRegistry)
     }
 
-    override def ignoreActiveSpan(): io.opentracing.Tracer.SpanBuilder = ???
+    private def createID(): Long =
+      ThreadLocalRandom.current().nextLong()
   }
 
+  private class TracerMetricsRecorder(recorder: EntityRecorder) {
+    val createdSpans = recorder.counter("created-spans")
+  }
 }