2 files changed, 71 insertions, 1 deletions
diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
index 76edb90c..c55f89e0 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -30,7 +30,7 @@ import scala.concurrent.ExecutionContext.Implicits.global
 @Aspect
 class RequestInstrumentation {
 
-  @DeclareMixin("play.api.mvc.Request || play.api.mvc.WrappedRequest || play.api.test.FakeRequest")
+  @DeclareMixin("play.api.mvc.RequestHeader || play.api.test.FakeRequest")
   def mixinContextAwareNewRequest: TraceContextAware = TraceContextAware.default
 
   @Pointcut("execution(* play.api.GlobalSettings+.onStart(*)) && args(application)")
diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
new file mode 100644
index 00000000..2fedcd70
--- /dev/null
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
@@ -0,0 +1,70 @@
+package kamon.play.instrumentation
+
+import javax.net.ssl.SSLContext
+import org.aspectj.lang.annotation.{ Around, Pointcut, Aspect }
+import org.aspectj.lang.ProceedingJoinPoint
+import com.ning.http.client._
+import com.ning.http.client.filter.{ RequestFilter, FilterContext }
+import kamon.trace.{ SegmentCompletionHandle, TraceRecorder }
+import kamon.metrics.TraceMetrics.HttpClientRequest
+
+@Aspect
+class WSInstrumentation {
+
+  @Pointcut("call(* play.api.libs.ws.WS$.newClient(..))")
+  def onNewAsyncHttpClient(): Unit = {}
+
+  @Around("onNewAsyncHttpClient()")
+  def aroundNewAsyncHttpClient(pjp: ProceedingJoinPoint): Any = {
+    val playConfig = play.api.Play.maybeApplication.map(_.configuration)
+    val wsTimeout = playConfig.flatMap(_.getMilliseconds("ws.timeout"))
+    val asyncHttpConfig = new AsyncHttpClientConfig.Builder()
+      .setConnectionTimeoutInMs(playConfig.flatMap(_.getMilliseconds("ws.timeout.connection")).orElse(wsTimeout).getOrElse(120000L).toInt)
+      .setIdleConnectionTimeoutInMs(playConfig.flatMap(_.getMilliseconds("ws.timeout.idle")).orElse(wsTimeout).getOrElse(120000L).toInt)
+      .setRequestTimeoutInMs(playConfig.flatMap(_.getMilliseconds("ws.timeout.request")).getOrElse(120000L).toInt)
+      .setFollowRedirects(playConfig.flatMap(_.getBoolean("ws.followRedirects")).getOrElse(true))
+      .setUseProxyProperties(playConfig.flatMap(_.getBoolean("ws.useProxyProperties")).getOrElse(true))
+
+    playConfig.flatMap(_.getString("ws.useragent")).map { useragent ⇒
+      asyncHttpConfig.setUserAgent(useragent)
+    }
+    if (!playConfig.flatMap(_.getBoolean("ws.acceptAnyCertificate")).getOrElse(false)) {
+      asyncHttpConfig.setSSLContext(SSLContext.getDefault)
+    }
+
+    asyncHttpConfig.addRequestFilter(new KamonRequestFilter())
+
+    new AsyncHttpClient(asyncHttpConfig.build())
+  }
+}
+
+class KamonRequestFilter extends RequestFilter {
+  import KamonRequestFilter._
+
+  override def filter(ctx: FilterContext[_]): FilterContext[_] = {
+    val completionHandle = TraceRecorder.startSegment(HttpClientRequest(ctx.getRequest.getRawURI.toString(), UserTime), basicRequestAttributes(ctx.getRequest()))
+    new FilterContext.FilterContextBuilder(ctx).asyncHandler(new AsyncHandlerWrapper[Response](ctx.getAsyncHandler(), completionHandle)).build()
+  }
+
+  class AsyncHandlerWrapper[T](asyncHandler: AsyncHandler[_], completionHandle: Option[SegmentCompletionHandle]) extends AsyncCompletionHandler[T] {
+    override def onCompleted(response: Response): T = {
+      completionHandle.map(_.finish(Map.empty))
+      asyncHandler.onCompleted().asInstanceOf[T]
+    }
+    override def onThrowable(t: Throwable) = {
+      asyncHandler.onThrowable(t)
+    }
+  }
+}
+
+object KamonRequestFilter {
+  val UserTime = "UserTime"
+
+  def basicRequestAttributes(request: Request): Map[String, String] = {
+    Map[String, String](
+      "host" -> request.getHeaders().getFirstValue("host"),
+      "path" -> request.getURI.getPath,
+      "method" -> request.getMethod)
+  }
+}
+