Fix scala/play2-6 support

newrelic · Oct 26, 2020 · fe8b486 · fe8b486
1 parent 4817190
commit fe8b486
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 12 deletions.
diff --git a/instrumentation/play-2.6/src/main/scala/com/nr/agent/instrumentation/play26/PlayUtil.scala b/instrumentation/play-2.6/src/main/scala/com/nr/agent/instrumentation/play26/PlayUtil.scala
@@ -18,15 +18,8 @@ object PlayUtil {
 
   val newRelicToken: TypedKey[Token] = TypedKey.apply("NR-TOKEN")
 
-  def appendToken(result: Either[Future[Result], (RequestHeader, Handler, Application)], token: Token) : Either[Future[Result], (RequestHeader, Handler, Application)] = {
-    var tempResult = result
-    if (tempResult.isRight) {
-      tempResult = result.right.flatMap(values => {
-        val newRequestHeader = values._1.addAttr(newRelicToken, token)
-        Right(newRequestHeader, values._2, values._3)
-      })
-    }
-    tempResult
+  def appendToken(request: RequestHeader, token: Token) : RequestHeader = {
+    request.addAttr(newRelicToken, token)
   }
 
 }
diff --git a/instrumentation/play-2.6/src/main/scala/play/core/server/Server_Instrumentation.scala b/instrumentation/play-2.6/src/main/scala/play/core/server/Server_Instrumentation.scala
@@ -13,18 +13,42 @@ import com.newrelic.api.agent.{NewRelic, Trace}
 import com.nr.agent.instrumentation.play26.PlayUtil
 import play.api.Application
 import play.api.mvc.{Handler, RequestHeader, Result}
+import play.core.ApplicationProvider
 
 import scala.concurrent._;
 
 @ScalaWeave(`type` = ScalaMatchType.Trait, `originalName` = "play.core.server.Server")
 class Server_Instrumentation {
   @Trace(dispatcher = true)
   def getHandlerFor(request: RequestHeader): Either[Future[Result], (RequestHeader, Handler, Application)] = {
-    var result: Either[Future[Result], (RequestHeader, Handler, Application)] = Weaver.callOriginal()
-    if (result.isRight) {
+    val tempResult: Either[Future[Result], (RequestHeader, Handler, Application)] = Weaver.callOriginal()
+    val result = tempResult.right.map { case (request, handler, application) =>
       // In order to correctly track async work from here we need to stash a token into the request.
       val token = NewRelic.getAgent.getTransaction.getToken
-      result = PlayUtil.appendToken(result, token)
+      (PlayUtil.appendToken(request, token), handler, application)
+    }
+
+    val tracer = NewRelic.getAgent.getTracedMethod
+    if (tracer != null) {
+      tracer.setMetricName("Play2Routing");
+    }
+
+    result
+  }
+}
+
+@ScalaWeave(`type` = ScalaMatchType.Object, `originalName` = "play.core.server.Server")
+class OServer_Instrumentation {
+  @Trace(dispatcher = true)
+  def getHandlerFor(
+    request: RequestHeader,
+    applicationProvider: ApplicationProvider
+  ): Either[Future[Result], (RequestHeader, Handler)] = {
+    val tempResult: Either[Future[Result], (RequestHeader, Handler)] = Weaver.callOriginal()
+    val result = tempResult.right.map { case (request, handler) =>
+      // In order to correctly track async work from here we need to stash a token into the request.
+      val token = NewRelic.getAgent.getTransaction.getToken
+      (PlayUtil.appendToken(request, token), handler)
     }
 
     val tracer = NewRelic.getAgent.getTracedMethod