Blaze client loses context in callback

[kubukoz]

Hi!

I'm playing with kamon-http4s again and I'm trying to wrap the blaze client with the middleware. Unfortunately, my context gets lost at the point Blaze creates a connection. I tracked it down to here:

//ClientChannelFactory.scala, blaze-core 0.14.11
  def connect(
      remoteAddress: SocketAddress,
      bufferSize: Int = bufferSize): Future[HeadStage[ByteBuffer]] = {
    val p = Promise[HeadStage[ByteBuffer]]

    try {
      val ch = AsynchronousSocketChannel.open(group.orNull)

      val onTimeout = //...

      val scheduledTimeout = scheduler.schedule(onTimeout, connectTimeout)

      val completionHandler = new CompletionHandler[Void, Null] {
        //...

        def completed(result: Void, attachment: Null): Unit = {
          channelOptions.applyToChannel(ch)
          p.trySuccess(new ByteBufferHead(ch, bufferSize = bufferSize))
          scheduledTimeout.cancel()
        }
      }

      try {
        ch.connect(remoteAddress, null: Null, completionHandler)
      } catch {
         //...
      }
    } catch { case NonFatal(t) => p.tryFailure(t) }

    p.future
  }

As far as I understand, the context should somehow be passed to the completion handler at instantiation time and restored in def completed, before completing the promise (then the further transformations on the future would also have the context, which they don't - so even though all the tracing information is in the request thanks to the middleware, the Kamon context has no span after any client call).

Can this be done somehow through the configuration of the agent, or would it require special treatment in blaze (e.g. being able to inject a method that creates the completion handler)?

[kubukoz]

I'll prepare a repro...
Also, I should note that the same thing happens with the AHC client.

Also, this is when I'm not using kamon-scala-future / kamon-cats-io (the only modules of kamon I have are zipkin, logback and http4s).

I had this wrapping every ExecutionContext in the world:

final class KamonExecutionContext(underlying: ExecutionContext) extends ExecutionContext {

  def execute(runnable: Runnable): Unit = {
    val ctx = Kamon.currentContext()

    underlying.execute(() => Kamon.runWithContext(ctx)(runnable.run()))
  }
  def reportFailure(cause: Throwable): Unit = underlying.reportFailure(cause)
}

Including the EC passed to BlazeClientBuilder, the one in IO.contextShift and IO.timer.

I also tried removing that and using kamon-scala-future + kamon-cats-io, but that didn't help (same result).

[kubukoz]

Reproduction in https://github.com/kubukoz/talks/blob/e09aae0b9479c21e75a7736a18908c86d73fcb1f/compositional-tracing/src/main/scala/com/kubukoz/tracing/magic/KamonTracing.scala

1. Start the upstream project on the same host (it'll create a new span in the same trace, although it doesn't use kamon at all - on purpose)
2. Run the client project with com.kubukoz.tracing.magic.KamonTracing as the main class
3. Look at console output in client (or logs in /tmp/logs/totallykibana.log to see output from both apps)

I also prepared one in which kamon-cats-io and kamon-scala-future are used instead of custom instrumentation. Same rules to run :)

https://github.com/kubukoz/talks/blob/kamon-instrumentation/compositional-tracing/src/main/scala/com/kubukoz/tracing/magic/KamonTracing.scala

[mladens]

Hey! From glancing over it, it seems that context gets lost after CompletionHandler gets wrapped in a runnable by sun.nio.ch.Invoker and handed to whichever executor underpins the socket channel.

I would try it with kamon-executors module, it should be able to pick up the runnable submission and carry the context 🤔

[kubukoz]

Isn't that included in the futures module?

[dpsoft]

Hey guys!, my output is:

2020-02-24 15:31:05.421 [INFO] [upstream] [-] [-] Started upstream
2020-02-24 15:32:05.219 [INFO] [upstream] [-] [-] Started upstream
2020-02-24 15:32:30.916 [INFO] [client  ] [cd3711992963fe55] [18dabc87d0d38177] Executing request Args(32837713-63f2-4f5b-a6d7-9451a6d183c4,hello)
2020-02-24 15:32:30.916 [INFO] [client  ] [82ff220dd9aad8d4] [b541d28a344074fb] Executing request Args(7927fca5-695d-4572-ae72-fac3b1578947,bye)
2020-02-24 15:32:31.134 [INFO] [client  ] [cd3711992963fe55] [18dabc87d0d38177] Before client call
2020-02-24 15:32:31.134 [INFO] [client  ] [82ff220dd9aad8d4] [b541d28a344074fb] Before client call
2020-02-24 15:32:32.252 [INFO] [upstream] [073de654b2b444c4] [18dabc87d0d38177] Created new span: 073de654b2b444c4
2020-02-24 15:32:32.252 [INFO] [upstream] [381403b56e4b261e] [b541d28a344074fb] Created new span: 381403b56e4b261e
2020-02-24 15:32:32.260 [INFO] [upstream] [073de654b2b444c4] [18dabc87d0d38177] Received execution request
2020-02-24 15:32:32.260 [INFO] [upstream] [381403b56e4b261e] [b541d28a344074fb] Received execution request
2020-02-24 15:32:32.579 [INFO] [client  ] [-] [-] Executed request Args(32837713-63f2-4f5b-a6d7-9451a6d183c4,hello)
2020-02-24 15:32:32.581 [INFO] [client  ] [-] [-] Executed request Args(7927fca5-695d-4572-ae72-fac3b1578947,bye)

I found this mesage:

[warn] The Kanela instrumentation can only be attached on the current JVM when using the ClassLoaderLayeringStrategy.Flat strategy but you are currently using [ScalaLibrary]. The application will run without instrumentation and might fail to behave properly.

and for that reason I ran the example with:

addSbtPlugin ("com.lightbend.sbt"% "sbt-javaagent"% "0.1.5") 

[kubukoz]

@dpsoft I'm attaching Kanela manually in bloop, but it's fine to add forking if you're running through sbt :)

The output is consistent with mine - I can't see the trace/span IDs in the "Executed request" messages.

[dpsoft]

@kubukoz this is my output:

2020-02-24 16:04:25.597 [INFO] [upstream] [-] [-] Started upstream
2020-02-24 16:04:37.283 [INFO] [client  ] [1f7d9c68586c3f19] [4fba4e3a029e9c15] Executing request Args(c55eb612-6d1c-4e96-aac2-da8cb0383b19,hello)
2020-02-24 16:04:37.283 [INFO] [client  ] [8db93dd767e5e417] [2bed914de8e17d25] Executing request Args(bb3fb985-c36f-46fb-b124-83841bdf0f1c,bye)
2020-02-24 16:04:37.504 [INFO] [client  ] [1f7d9c68586c3f19] [4fba4e3a029e9c15] Before client call
2020-02-24 16:04:37.504 [INFO] [client  ] [8db93dd767e5e417] [2bed914de8e17d25] Before client call
2020-02-24 16:04:38.601 [INFO] [upstream] [aba85a887b557f58] [4fba4e3a029e9c15] Created new span: aba85a887b557f58
2020-02-24 16:04:38.601 [INFO] [upstream] [98ef4e52d1651798] [2bed914de8e17d25] Created new span: 98ef4e52d1651798
2020-02-24 16:04:38.606 [INFO] [upstream] [98ef4e52d1651798] [2bed914de8e17d25] Received execution request
2020-02-24 16:04:38.606 [INFO] [upstream] [aba85a887b557f58] [4fba4e3a029e9c15] Received execution request
2020-02-24 16:04:38.832 [INFO] [client  ] [8db93dd767e5e417] [2bed914de8e17d25] Executed request Args(c55eb612-6d1c-4e96-aac2-da8cb0383b19,hello)
2020-02-24 16:04:38.834 [INFO] [client  ] [8db93dd767e5e417] [2bed914de8e17d25] Executed request Args(bb3fb985-c36f-46fb-b124-83841bdf0f1c,bye)

using the executor-service-capture-on-submit(disabled by default) feature from kamon-executors module:

"io.kamon" %% "kamon-executors" % "2.0.2",

kanela {
   modules {
        """Provides automatic context propagation by capturing the current context at the instant when 
       a  `Runnable/Callable` is submitted for execution by instrumenting the executors themselves, 
       including JDK classes"""
    executor-service-capture-on-submit.enabled = true
    executor-service.enabled = false
  }
}

[kubukoz]

In that output, the two concurrent calls from different traces (and spans, in fact) ended up in the same span :/

[dpsoft]

:/

[kubukoz]

Is there an easy way to instrument CompletionHandler to capture the context at the point of creation and restore it in the completed method?

[mladens]

I tried by reenabling instrumentation on sun.* which is excluded by defaullt but it misses it since its loaded by bootstrap classloader. After enabling bootstrapInjection it starts failing when loading instrumentation since kamon.* is not available in bootstrap cl, @dpsoft any hints there? (tried adding dependencies as helper-class-names param in config but didnt get far :/ )

[dpsoft]

@mladens you are reenabling the instrumentation on sun.* in executor-service instrumentation? You need a little of more work if you want to instrument classes in the bootstrap, i.g. in the case of the instrumentation of the executors we need add a new indirection in order to avoid the kamon dependency. https://github.com/kamon-io/kamon-executors/blob/master/kamon-executors/src/main/java/kamon/instrumentation/executor/CaptureContextOnSubmitInstrumentation.java#L111-L125. and here other example:

https://github.com/kamon-io/kanela/blob/master/agent-test/src/test/scala/app/kanela/instrumentation/BootstrapInstrumentation.scala#L24

https://github.com/kamon-io/kanela/blob/master/agent-test/src/test/scala/app/kanela/specs/BootstrapInstrumentationSpec.scala#L24

https://github.com/kamon-io/kanela/blob/master/agent-test/src/test/resources/application.conf#L17-L25

[dpsoft]

@kubukoz @mladens In order to instrument CompletionHandler we can do the following:

package kamon.instrumentation.handler

import kamon.instrumentation.context.{CaptureCurrentContextOnExit, HasContext, InvokeWithCapturedContext}
import kanela.agent.api.instrumentation.InstrumentationBuilder

class CompletionHandlerInstrumentation extends InstrumentationBuilder {
  onSubTypesOf("java.nio.channels.CompletionHandler")
    .mixin(classOf[HasContext.Mixin])
    .advise(isConstructor, CaptureCurrentContextOnExit)
    .advise(method("completed"), InvokeWithCapturedContext)
}

in application.conf

kanela {
  modules {
    completion-handler {
      name = "Completion Handler Instrumentation"

      instrumentations = [
        "kamon.instrumentation.handler.CompletionHandlerInstrumentation"
      ]

      within = [
        "org.http4s.blaze.channel..*"
      ]
    }
  }
}

and the output with this instrumentation is:

2020-02-24 18:49:10.176 [INFO] [upstream] [-] [-] Started upstream
2020-02-24 18:49:16.828 [INFO] [client  ] [40fe56078f5c8442] [e0844e682ecd9df8] Executing request Args(96a8df1c-70b5-43ce-a07b-6adf3e1d4690,hello)
2020-02-24 18:49:16.828 [INFO] [client  ] [144a0ffcac16f617] [5ac8fd044386495f] Executing request Args(a68c8c8b-016d-482d-b321-25f312ae0886,bye)
2020-02-24 18:49:17.023 [INFO] [client  ] [40fe56078f5c8442] [e0844e682ecd9df8] Before client call
2020-02-24 18:49:17.023 [INFO] [client  ] [144a0ffcac16f617] [5ac8fd044386495f] Before client call
2020-02-24 18:49:18.148 [INFO] [upstream] [e9f1dbda4e47ffb6] [5ac8fd044386495f] Created new span: e9f1dbda4e47ffb6
2020-02-24 18:49:18.148 [INFO] [upstream] [cc9e80d40c367f0e] [e0844e682ecd9df8] Created new span: cc9e80d40c367f0e
2020-02-24 18:49:18.154 [INFO] [upstream] [cc9e80d40c367f0e] [e0844e682ecd9df8] Received execution request
2020-02-24 18:49:18.154 [INFO] [upstream] [e9f1dbda4e47ffb6] [5ac8fd044386495f] Received execution request
2020-02-24 18:49:18.402 [INFO] [client  ] [144a0ffcac16f617] [5ac8fd044386495f] Executed request Args(a68c8c8b-016d-482d-b321-25f312ae0886,bye)
2020-02-24 18:49:18.403 [INFO] [client  ] [144a0ffcac16f617] [5ac8fd044386495f] Executed request Args(96a8df1c-70b5-43ce-a07b-6adf3e1d4690,hello)

[Ben-Woolley]

Hello! I just encountered this issue myself.
Is this going to be fixed in kamon-http4s?