Potential Finagle ServerDispatcher Annotations solution

ServerDispatcher.scala

package com.twitter.finagle.dispatch

import com.twitter.finagle.context.Contexts
import com.twitter.finagle.tracing.{Annotation, Trace}
import com.twitter.finagle.transport.Transport
import com.twitter.finagle.{Service, NoStacktrace, CancelledRequestException}
import com.twitter.util._
import java.util.concurrent.atomic.AtomicReference
import scala.collection.mutable.ListBuffer

trait ServerDispatcher[Req, Rep, In, Out] extends Closable {
  protected val service: Service[Req, Rep]
  protected val trans: Transport[Req, Rep]

  def dispatch(req: Req, eos: Promise[Unit]): Future[Rep] = 
    service(req) ensure eos.setDone()
  def handle(rep: Rep): Future[Unit] = trans.write(rep)
}

class DecoratedServerDispatcher[Req, Rep, In, Out, A <: ServerDispatcher[Req, Rep, In, Out]]
  (dispatcher: A, dispatchAnnotations: List[Annotation], handleSuccessAnnotations: List[Annotation],
  handleFailureAnnotations: List[Annotation]) extends ServerDispatcher[Req, Rep, In, Out] {

  override def handle(rep: Rep): Future[Unit] =
    dispatcher.handle(rep) onSuccess { _ =>
      recordAnnotations(handleSuccessAnnotations) 
    } onFailure { ex => 
      recordAnnotations(handleFailureAnnotations)
      Future.exception(ex)
    }
  override def dispatch(req: Req, eos: Promise[Unit]): Future[Rep] = {
    recordAnnotations(dispatchAnnotations)
    dispatcher.dispatch(req, eos)
  }

  @annotation.tailrec
  private[this] def recordAnnotations(xs: List[Annotation]): Unit = xs match {
    case h :: t => Trace.record(h); recordAnnotations(t)
    case t      => Unit
  }
}

I want to make a ServerDispatcher immutable so that each time we add an annotation we're creating a new ServerDispatcher, but this is the basis of what I had in mind.

Need a better name than Decorator though.

Also, this implementation would allow us to plug in the Tracing for things like Mux or Http pretty easily - just call insert the recordDispatchAnnotations / recordHandleAnnotations at the proper site and we're good to go.