WSInstrumentation refactor

2 files changed, 23 insertions, 50 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 17fc772a..e671d097 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -93,8 +93,5 @@ class RequestInstrumentation {
     } else None
 
     TraceRecorder.start(defaultTraceName, token)(system)
-
-    //Necessary to force initialization of traceContext when initiating the request.
-    requestHeader.asInstanceOf[TraceContextAware].traceContext
   }
 }
diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
index 498afc4d..fce5ca13 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
@@ -16,71 +16,47 @@
 
 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.trace.TraceRecorder
 import kamon.metrics.TraceMetrics.HttpClientRequest
+import play.api.libs.ws.WS.WSRequest
+import scala.concurrent.Future
+import play.api.libs.ws.Response
+import scala.util.{ Failure, Success }
+import scala.concurrent.ExecutionContext.Implicits.global
 
 @Aspect
 class WSInstrumentation {
 
-  @Pointcut("call(* play.api.libs.ws.WS$.newClient(..))")
-  def onNewAsyncHttpClient(): Unit = {}
+  @Pointcut("execution(* play.api.libs.ws.WS$WSRequest.execute()) && this(request)")
+  def onExecuteRequest(request: WSRequest): 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))
+  @Around("onExecuteRequest(request)")
+  def aroundExecuteRequest(pjp: ProceedingJoinPoint, request: WSRequest): Any = {
+    import WSInstrumentation._
 
-    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._
+    val completionHandle = TraceRecorder.startSegment(HttpClientRequest(request.url, UserTime), basicRequestAttributes(request))
 
-  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()
-  }
+    val response = pjp.proceed().asInstanceOf[Future[Response]]
 
-  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)
+    response.onComplete {
+      case Failure(t) ⇒ completionHandle.map(_.finish(Map("completed-with-error" -> t.getMessage)))
+      case Success(_) ⇒ completionHandle.map(_.finish(Map.empty))
     }
+
+    response
   }
 }
 
-object KamonRequestFilter {
+object WSInstrumentation {
   val UserTime = "UserTime"
 
-  def basicRequestAttributes(request: Request): Map[String, String] = {
+  def basicRequestAttributes(request: WSRequest): Map[String, String] = {
     Map[String, String](
-      "host" -> request.getHeaders().getFirstValue("host"),
-      "path" -> request.getURI.getPath,
-      "method" -> request.getMethod)
+      "host" -> request.header("host").getOrElse("Unknown"),
+      "path" -> request.method,
+      "method" -> request.method)
   }
 }