Netty handler for Kamon context

gistfile1.scala

package server

import io.sphere.util.Logging
import kamon.Kamon
import kamon.trace.TraceContext
import org.jboss.netty.channel._
import org.jboss.netty.handler.codec.http.HttpRequest

import scala.util.control.NonFatal

class CorrelationIdHandler extends ChannelUpstreamHandler with ChannelDownstreamHandler with Logging {

  var context: TraceContext = _

  override def handleUpstream(ctx: ChannelHandlerContext, e: ChannelEvent): Unit = e match {
    case msg: MessageEvent =>
      val request = msg.getMessage.asInstanceOf[HttpRequest]
      val correlationId = Option(request.headers().get("X-Correlation-ID"))
      context = Kamon.tracer.newContext("ws", correlationId)
      ctx.sendUpstream(e)
    case _ =>
      ctx.sendUpstream(e)
  }

  override def handleDownstream(ctx: ChannelHandlerContext, e: ChannelEvent): Unit = e match {
    case msg: MessageEvent =>
      try {
        context.finish()
      } catch {
        case NonFatal(t) ⇒
          log.warn("error while finishing the kamon context", t)
      }
      ctx.sendDownstream(e)
    case _ =>
      ctx.sendDownstream(e)
  }
}