Correlation IDs in Scala using Monix by Adam Warski

.executeWithOptions(_.localContextPropagation)

-Dmonix.environment.localContextPropagation=1

The first three work with any datatype for which there’s an implementation of the cats-effect typeclasses. Monix’s 'Task' is one of them.

Warning: the code below includes global state and unwrapped, side-effecting method calls. Proceed with caution!

As we said earlier, we’ll be using a 'Local' instance. It will store an 'Option[String]' with the current correlation ID(if any):

object CorrelationId extends StrictLogging {
  private val localCid = new Local[Option[String]](() => {
    None
  })

  def apply(): Option[String] = localCid()
}

When receiving an http request (via http4s), we try to read the correlation ID from a 'X-Correlation-ID' header, or if it’s not set, create a new one — a random string. This can be implemented with a http4s middleware:

object CorrelationId {
  val CorrelationIdHeader = "X-Correlation-ID"
  
  def setCorrelationIdMiddleware(service: HttpRoutes[Task]): HttpRoutes[Task] = 
    Kleisli { req: Request[Task] =>
      val cid = req.headers.get(CaseInsensitiveString(CorrelationIdHeader)) match {
        case None => newCorrelationId()
        case Some(cidHeader) => cidHeader.value
      }

      localCid.update(Some(cid))
      logger.info("Starting request")
      service(req)
    }
}

Services that want to use correlation IDs must be wrapped with the middleware, for example as follows:

 

object Server1 extends App with StrictLogging {
  val dsl = Http4sDsl[Task]
  import dsl._

  val service = HttpRoutes.of[Task] {
    case GET -> Root / "test1" => Ok()
  }

  BlazeServerBuilder[Task]
    .bindHttp(8081)
    .withHttpApp(CorrelationId.setCorrelationIdMiddleware(service).orNotFound)
    .serve
    .compile
    .drain
    .runSyncUnsafe()
}

This takes care of the request-receiving side (the server). What about request sending (the client)? As we mentioned earlier, we’ll be using sttp to send HTTP requests. Here, we need to set the correlation ID header on outgoing requests. To do this, we create a wrapper for an sttp backend:

class SetCorrelationIdBackend(delegate: SttpBackend[Task, Nothing]) 
  extends SttpBackend[Task, Nothing] {
    
  override def send[T](request: sttp.Request[T, Nothing]): Task[Response[T]] = {
    // suspending the calculation of the correlation id until the request 
    // send is evaluated
    Task {
      CorrelationId() match {
        case Some(cid) => request.header(CorrelationIdHeader, cid)
        case None => request
      }
    }.flatMap(delegate.send)
  }

  override def close(): Unit = delegate.close()

  override def responseMonad: MonadError[Task] = delegate.responseMonad
}

Having such a wrapped implicit backend in scope, when an HTTP request is sent, it will include the current correlation ID. This is as simple as invoking e.g. 'sttp.get(uri"http://localhost:8082/test2").send()', which yields a 'Task[Response[String]]'.

Finally, what about logging? The “official” way to pass additional data to be used in log statements is using MDC, however this mechanism relies on thread locals, so isn’t usable in our case. I couldn’t find another way to intercept log events, so we’ll abuse logback’s 'Filter' and set the thread-local 'MDC' variable there. This, in turn, will make them automatically available in logging patterns:

class SetCorrelationIdInMDCFilter extends Filter[ILoggingEvent] {
  override def decide(event: ILoggingEvent): FilterReply = {
    CorrelationId().foreach(MDC.put("cid", _)) 
    FilterReply.NEUTRAL
  }
}

The 'cid' field can be then used in the log message pattern, if it is set:

<appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
  <encoder>
    <pattern>%d{HH:mm:ss.SSS} %highlight(%replace([%X{cid}] ){'\[\] ', ''})[%thread] %-5level %logger{5} - %msg%n</pattern>
  </encoder>

  <filter class="com.softwaremill.monix.SetCorrelationIdInMDCFilter" />
</appender>

And that’s it. Given that you start the servers from the PoC repository using the:

-Dmonix.environment.localContextPropagation=1

option, the correlation ID will be automatically propagated across 'Task's, both in synchronous and asynchronous operations — hence also involving thread boundaries.

Server1 in action. Note that a single request is served by multiple threads.

To verify this for yourself, you can start the 'Server1', 'Server2' and 'Server3' applications (which bind to ports '8081', '8082' and '8083', respectively — but be sure to add the context propagation option!), and then invoke a single request to 'Server1' by running the 'Client' class. You should see log messages with the (same) correlation ID highlighted in each server; the response to the request will consist of correlation IDs as seen at various stages of evaluating the call chain: propagating http calls, running parallel operations (using 'Task.gatherUnordered') and running database queries. Making another request will use another correlation ID. 

Moreover, you can run a test which issues a number of parallel requests using the 'StressClient' class. The result of each call will be verified to check that it contains only the appropriate correlation IDs.

Finally, the 'SingleServer' class contains a more contained setup, which starts a single server and makes a single http request.

 

Notes/TODOs

The code is of PoC-quality, and as such has some improvement possibilities, notably:

• in 'SetCorrelationIdInMDCFilter', the thread-local MDC variable is set but never cleared, as the filter is only run before a message is logged. Probably a different logback mechanism would have to be used to do this properly.
• similarly, in 'setCorrelationIdMiddleware', the correlation ID is set in the 'Local', but never cleared
• only enabling local context propagation globally, using the JVM option works. Evaluating the main task with the modified options (setting them just before 'runSyncUnsafe') seems to have no effect. My guess is that http4s also evaluates 'Task's to 'Future's somewhere internally, this time with the default options.
• finally, the demo is using http4s version 0.20.0-M3, which is obviously not production-ready. A stable verson can’t be used, as Monix 3.0.0-RC2, which includes a number of fixes for 'Local's, depends on a newer version of cats-effect, which is not compatible with the one used by http4s 0.18.