Executor that propagates Brave tracing context

TracingExecutionContext.scala

package com.bigcommerce.core.metrics.tracing

import com.github.kristofa.brave.{ServerSpan, ServerSpanThreadBinder}

import scala.concurrent.ExecutionContext

/**
  * Created by zack.angelo on 12/5/16.
  */
class TracingExecutionContext(wrapped:ExecutionContext,
                              serverSpanThreadBinder: ServerSpanThreadBinder)
  extends ExecutionContext { self =>

  override def execute(runnable: Runnable): Unit =
    wrapped.execute(runnable)

  override def reportFailure(cause: Throwable): Unit =
    wrapped.reportFailure(cause)

  protected def prepareWithSpanState(spanState: ServerSpan): ExecutionContext = {
    new ExecutionContext {
      override def reportFailure(cause: Throwable): Unit =
        self.reportFailure(cause)

      override def execute(runnable: Runnable): Unit = {
        val tracedRunnable = new Runnable {
          override def run(): Unit = {
            val oldState = serverSpanThreadBinder.getCurrentServerSpan

            //Propagate tracing context
            serverSpanThreadBinder.setCurrentSpan(spanState)
            runnable.run()

            serverSpanThreadBinder.setCurrentSpan(oldState)
          }
        }

        self.execute(tracedRunnable)
      }

      override def prepare(): ExecutionContext =
        self.prepareWithSpanState(serverSpanThreadBinder.getCurrentServerSpan)
    }
  }

  override def prepare(): ExecutionContext =
    self.prepareWithSpanState(serverSpanThreadBinder.getCurrentServerSpan)
}