scala cats-effects instrumentation loosing Transaction track when IOFromFuture is used

[michalkoza]

Description

Hi, It seems that the cats-effect-2 instrumentation is not robust against using IO.fromFuture. As far as I managed to figure out the core of cats-effect-2 instrumentation weaves on the IOShift mechanizm and adds a logic to the Tick creation that stores NR transaction token as tokenAndRefCount. The problem with IO.fromFuture seems to be that inside IOFromFuture.apply:

private[effect] object IOFromFuture {

  /**
   * Implementation for `IO.fromFuture`.
   */
  def apply[A](f: Future[A]): IO[A] =
    f.value match {
      case Some(result) =>
        result match {
          case Success(a) => IO.pure(a)
          case Failure(e) => IO.raiseError(e)
        }
      case _ =>
        IO.async { cb =>
          f.onComplete { r =>
            cb(r match {
              case Success(a) => Right(a)
              case Failure(e) => Left(e)
            })
          }(immediate)
        }
    }
}

if the case _ => is matched, the currently executing and properly instrumented thread (A) just schedules the f.onComplete call and finishes its life taking its tokenAndRefCount to the grave. The function passed to onComplete executes as a new runnable on some other thread (B) which does not have the tokenAndRefCount stored. As far as I understand it is the thread on which the Future computation is running, thus it has been spawned by Slick which is wrapped in NR instrumentation but in some other/old way - outside the IOShift mechanizm. When f completes and the (B) thread is executing the callback (IORunLoop) inside which new Tick is created but this time Utils.getThreadTokenAndRefCount() returns null and there is no token and the entire following execution is not linked to any Transaction.

Your Environment

Scala 2.13, cats-effects 2.5.5, Slick 3.3.3, NewRelicAgent 7.9.0

(Migrate to Jira)

[michalkoza]

Updating instrumentation/slick-2.12_3.2.0/src/main/scala/slick/util/AsyncExecutor.scala the following way + giving it Utils from cats instrumentation seems to be a working fix:

 import scala.concurrent.{ExecutionContext, Future, Promise}
 import scala.concurrent.duration._
 import slick.util.AsyncExecutor
+import slick.util.Utils

 @ScalaWeave(`type` = ScalaMatchType.Object, originalName="slick.util.AsyncExecutor")
 class WeavedAsyncExecutor {
@@ -80,9 +81,14 @@ class NewRelicExecutionContext(delegatee :ExecutionContext) extends ExecutionCon
 }

 class NewRelicRunnable(runnable :Runnable) extends Runnable {
+
+  private val tokenAndRefCount = Utils.getThreadTokenAndRefCount()
+
   @Trace(async = true)
   override def run() {
     try {
+    try {
+      Utils.setThreadTokenAndRefCount(tokenAndRefCount)
       AgentBridge.currentApiSource.set(WeavePackageType.INTERNAL)

       if(AgentBridge.getAgent().startAsyncActivity(runnable)) {
@@ -98,5 +104,8 @@ class NewRelicRunnable(runnable :Runnable) extends Runnable {
       AgentBridge.currentApiSource.remove()
     }
     runnable.run()
+    } finally {
+      Utils.clearThreadTokenAndRefCountAndTxn(tokenAndRefCount)
+    }
   }
 }

[kbford56]

@michalkoza thank you for the detailed problem statement and the analysis on a possible approach to fixing this!

While we cannot commit to when we can devote time to this issue, having this amount of detail is a good start - thanks again!

[workato-integration[bot]]

https://issues.newrelic.com/browse/NEWRELIC-4096