OpenTracingUtils.scala

package com.strava.tracing

import io.opentracing.{Span, Tracer}
import com.twitter.util.{Future, Throw}

import io.opentracing.util.GlobalTracer
import io.opentracing.tag.Tags

object TracingUtils {
  /** Trace an asynchronous operation. This handles creating a span, starting
    * it, setting it active, and closing it when the operation completes, as
    * well as setting the error tag if the operation fails.
    *
    * If, within the operation you are tracing, you'd like to add more tags to
    * the span or otherwise use it, retrieve the active span with
    * `tracer.activeSpan`. Alternatively, use a combination of the other helper
    * methods on this object.
    *
    * @param operationName a name for the operation to be performed; this will
    * be used as the main name for the span.
    * @param tags key/value metadata pairs describing additional properties of the span.
    * @param tracer the tracer to use for instrumentation; this should generally
    * be left as the default value except in circumstances like unit testing.
    * @param op the asynchronous operation to be instrumented.
    * @return the result of the given operation.
    * @see the
    * [[https://github.com/opentracing/specification/blob/master/specification.md#start-a-new-span
    * OpenTracing specification]] for guidance on setting operationName
    * @see
    * [[https://github.com/opentracing/specification/blob/master/semantic_conventions.md
    * OpenTracing semantic conventions]] for guidance on standard tags.
    * @see
    * [[https://github.com/opentracing/opentracing-java/blob/master/opentracing-api/src/main/java/io/opentracing/tag/Tags.java
    * io.opentracing.tag.Tags]] for constants representing many common tags.
    */
  def trace[T](
    operationName: String,
    tags: Map[String, String] = Map(),
    tracer: Tracer = GlobalTracer.get
  )(op: => Future[T]): Future[T] = {
    val span = buildSpan(operationName, tags, tracer).start

    withSpanOverOp(span, tracer) {
      withSpanActive(span, tracer)(op)
    }
  }

  def withSpanOverOp[T](span: Span, tracer: Tracer = GlobalTracer.get)(op: => Future[T]): Future[T] = {
    op.respond { result =>
      try {
        result match {
          case Throw(ex) => span.setTag(Tags.ERROR.getKey, true)
          case _ =>
        }
      } finally {
        span.finish
      }
    }
  }

  def withSpanActive[T](span: Span, tracer: Tracer = GlobalTracer.get)(f: => T): T = {
    val scope = tracer.scopeManager.activate(span)
    try {
      f
    } finally {
      scope.close
    }
  }

  def buildSpan(
    operationName: String,
    tags: Map[String, String],
    tracer: Tracer = GlobalTracer.get
  ): Tracer.SpanBuilder = {
    tags.foldLeft(tracer.buildSpan(operationName)) { (spanBuilder, tag) =>
      spanBuilder.withTag(tag._1, tag._2)
    }
  }
}