DEV Community: Linh Nguyen

OpenTelemetry Distributed Tracing with ZIO

Linh Nguyen — Mon, 01 Nov 2021 00:00:00 +0000

Introduction

This post is some quick notes on using ZIO and zio-telemetry to implement OpenTelemetry distributed tracing for Scala applications.
The source code is available here.

This is not an introduction to any of these technologies, but here are a few good reads:

NewRelic's introduction to Distributed Tracing.
Lightstep's concise summary for OpenTelemetry.

Initial implementation

For demonstration purpose, we will perform manual instrumentation on a modified version of the zio-grpc's helloworld example, in which we incorporate both gRPC and HTTP communications:

Original: hello-client sends a HelloRequest with name x and hello-server returns a HelloResponse with message Hello, x.
Modified: in addition to the original behavior, client sends an optional integer field guess and server performs an HTTP request to HTTPBin based on its value.

Add the new flag `guess`:

// The greeting service definition.
service Greeter {
  // Sends a greeting
  rpc SayHello (HelloRequest) returns (HelloReply) {}
}

// The request message containing the user's name.
message HelloRequest {
  string name = 1;
  google.protobuf.Int32Value guess = 2;
}

// The response message containing the greetings
message HelloReply {
  string message = 1;
}

Add zio-grpc dependency

resolvers += Resolver.sonatypeRepo("snapshots")

addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.4.3")

addSbtPlugin("com.thesamet" % "sbt-protoc" % "1.0.4")

val zioGrpcVersion = "0.5.1+12-93cdbe22-SNAPSHOT"

libraryDependencies ++= Seq(
  "com.thesamet.scalapb.zio-grpc" %% "zio-grpc-codegen" % zioGrpcVersion,
  "com.thesamet.scalapb"          %% "compilerplugin"   % "0.11.5"
)

Set up `build.sbt`

Generate Scala code from helloworld.proto.
Depend on sttp for HTTP client.

val grpcVersion = "1.41.0"
val sttpVersion = "3.3.15"

val scalaPBRuntime = "com.thesamet.scalapb" %% "scalapb-runtime-grpc" % scalapb.compiler.Version.scalapbVersion

val grpcRuntimeDeps = Seq(
  "io.grpc"      % "grpc-netty" % grpcVersion,
  scalaPBRuntime,
  scalaPBRuntime % "protobuf"
)

val sttpZioDeps = Seq(
  "com.softwaremill.sttp.client3" %% "async-http-client-backend-zio" % sttpVersion
)

lazy val root = Project("opentelemetry-distributed-tracing-zio", file(".")).aggregate(zio)

lazy val zio = commonProject("zio").settings(
  Compile / PB.targets := Seq(
    scalapb.gen(grpc = true)          -> (Compile / sourceManaged).value,
    scalapb.zio_grpc.ZioCodeGenerator -> (Compile / sourceManaged).value
  ),
  libraryDependencies ++= grpcRuntimeDeps ++ sttpZioDeps
)

Client implementation

Create a gRPC client pointing to localhost:9000.
Send a single HelloRequest.
Send 5 HelloRequests in parallel.
Send a single HelloRequest with an invalid guess.
Print "Done" and exit.

object ZClient extends zio.App {
  private val clientLayer = GreeterClient.live(
    ZManagedChannel(
      ManagedChannelBuilder.forAddress("localhost", 9000).usePlaintext()
    )
  )

  private val singleHello = GreeterClient.sayHello(HelloRequest("World"))

  private val multipleHellos = ZIO.collectAllParN(5)(
    List(
      GreeterClient.sayHello(HelloRequest("1", Some(1))),
      GreeterClient.sayHello(HelloRequest("2", Some(2))),
      GreeterClient.sayHello(HelloRequest("3", Some(3))),
      GreeterClient.sayHello(HelloRequest("4", Some(4))),
      GreeterClient.sayHello(HelloRequest("5", Some(5)))
    )
  )

  private val invalidHello = GreeterClient.sayHello(HelloRequest("Invalid", Some(-1))).ignore

  private def myAppLogic =
    singleHello *> multipleHellos *> invalidHello *> putStrLn("Done")

  def run(args: List[String]): URIO[ZEnv, ExitCode] =
    myAppLogic.provideCustomLayer(clientLayer).exitCode
}

Server implementation

Fail the request if guess is less than 0.
Based on the value of guess, delay for some time and then send a request to HTTPBin.

type ZGreeterEnv = Clock with Random with SttpClient

object ZGreeterImpl extends RGreeter[ZGreeterEnv] {

  def sayHello(request: HelloRequest): ZIO[ZGreeterEnv, Status, HelloReply] = {
    val guess = request.guess.getOrElse(0)
    for {
      _      <- ZIO.fail(Status.INVALID_ARGUMENT).when(guess < 0)
      code   <- ???
      delayMs = ???
      _      <- httpRequest(code)
                  .delay(delayMs.millis)
                  .mapError(ex => Status.INTERNAL.withCause(ex))
    } yield HelloReply(s"Hello, ${request.name}")
  }

  def httpRequest(code: Int): RIO[SttpClient, Unit] =
    send(basicRequest.get(uri"https://httpbin.org/status/$code")).unit
}

Run it

To run the server:

$ sbt "zio/runMain com.github.tuleism.ZServer"

[info] running (fork) com.github.tuleism.ZServer
[info] Server is running. Press Ctrl-C to stop.

To run the client:

$ sbt "zio/runMain com.github.tuleism.ZClient"

[info] running (fork) com.github.tuleism.ZClient
[info] Done
[success] Total time: 12 s

At this point, we only know that it takes roughly 12 seconds for the client to initialize and finish its work.

Let's add distributed tracing to gain more insights into this.

Common tracing requirements

For both client and server, we need to acquire a Tracer, an object responsible for creating and managing Spans.
Tracing data is sent to Jaeger, which acts as a standalone collector.

Add new dependencies

zio-telemetry's OpenTelemetry module.
We also depend on zio-config to read tracing config from file and zio-magic to ease ZLayer wiring.

val openTelemetryVersion = "1.6.0"
val zioConfigVersion     = "1.0.10"
val zioMagicVersion      = "0.3.9"
val zioTelemetryVersion  = "0.8.2"

val openTelemetryDeps = Seq(
  "io.opentelemetry" % "opentelemetry-exporter-jaeger"    % openTelemetryVersion,
  "io.opentelemetry" % "opentelemetry-sdk"                % openTelemetryVersion,
  "io.opentelemetry" % "opentelemetry-extension-noop-api" % s"$openTelemetryVersion-alpha"
)

val zioConfigDeps = Seq(
  "dev.zio" %% "zio-config"          % zioConfigVersion,
  "dev.zio" %% "zio-config-magnolia" % zioConfigVersion,
  "dev.zio" %% "zio-config-typesafe" % zioConfigVersion
)

val zioMagicDeps = Seq(
  "io.github.kitlangton" %% "zio-magic" % zioMagicVersion
)

val zioTelemetryDeps = Seq(
  "dev.zio"                       %% "zio-opentelemetry"                   % zioTelemetryVersion,
  "com.softwaremill.sttp.client3" %% "zio-telemetry-opentelemetry-backend" % sttpVersion
)

Add a config layer

tracing {
  enable = false
  enable = ${?TRACING_ENABLE}
  endpoint = "http://127.0.0.1:14250"
  endpoint = ${?JAEGER_ENDPOINT}
}

case class AppConfig(tracing: TracingConfig)

case class TracingConfig(enable: Boolean, endpoint: String)

object AppConfig {
  private val configDescriptor = descriptor[AppConfig]

  val live: Layer[ReadError[String], Has[AppConfig]] = TypesafeConfig.fromDefaultLoader(configDescriptor)
}

Add a `Tracer` layer

Depend on the configuration, we either create a noop Tracer or one that sends data to Jaeger.
Once we have it, we can construct a Tracing layer, which give us access to many useful operations in zio-telemetry.

object ZTracer {
  private val InstrumentationName = "com.github.tuleism"

  private def managed(serviceName: String, endpoint: String) = {
    val resource = Resource.builder().put(ResourceAttributes.SERVICE_NAME, serviceName).build()
    for {
      spanExporter   <- ZManaged.fromAutoCloseable(
                          Task(JaegerGrpcSpanExporter.builder().setEndpoint(endpoint).build())
                        )
      spanProcessor  <- ZManaged.fromAutoCloseable(UIO(SimpleSpanProcessor.create(spanExporter)))
      tracerProvider <- UIO(
                          SdkTracerProvider.builder().addSpanProcessor(spanProcessor).setResource(resource).build()
                        ).toManaged_
      openTelemetry  <- UIO(OpenTelemetrySdk.builder().setTracerProvider(tracerProvider).build()).toManaged_
      tracer         <- UIO(openTelemetry.getTracer(InstrumentationName)).toManaged_
    } yield tracer
  }

  def live(serviceName: String): RLayer[Has[TracingConfig], Has[Tracer]] =
    (
      for {
        config <- ZIO.service[TracingConfig].toManaged_
        tracer <- if (!config.enable) {
                    Task(NoopOpenTelemetry.getInstance().getTracer(InstrumentationName)).toManaged_
                  } else {
                    managed(serviceName, config.endpoint)
                  }
      } yield tracer
    ).toLayer
}

New server

Instrument the HTTP client

Use out-of-the-box sttp backend.
We also add additional HTTP specific attributes according to the OpenTelemetry's semantic convention.

object SttpTracing {
  private val wrapper = new ZioTelemetryOpenTelemetryTracer {
    def before[T](request: Request[T, Nothing]): RIO[Tracing, Unit] =
      Tracing.setAttribute(SemanticAttributes.HTTP_METHOD.getKey, request.method.method) *>
        Tracing.setAttribute(SemanticAttributes.HTTP_URL.getKey, request.uri.toString()) *>
        ZIO.unit

    def after[T](response: Response[T]): RIO[Tracing, Unit] =
      Tracing.setAttribute(SemanticAttributes.HTTP_STATUS_CODE.getKey, response.code.code) *>
        ZIO.unit
  }

  val live = AsyncHttpClientZioBackend.layer().flatMap { hasBackend =>
    ZIO
      .service[Tracing.Service]
      .map { tracing =>
        ZioTelemetryOpenTelemetryBackend(hasBackend.get, tracing, wrapper)
      }
      .toLayer
  }
}

Instrument the gRPC server

We can add Tracing without changing our server implementation with a ZTransform. For each request:

We use zio-telemetry's spanFrom, which extracts the propagated context (through gRPC Metadata, using W3C Trace Context format) and starts a new child Span right after.
We have access to a RequestContext and thus the full method name used for Span's name.
We also add additional gRPC specific attributes according to the OpenTelemetry's semantic convention.

object GrpcTracing {
  private val propagator: TextMapPropagator = W3CTraceContextPropagator.getInstance()

  private val metadataGetter: TextMapGetter[Metadata] = new TextMapGetter[Metadata] {
    override def keys(carrier: Metadata): java.lang.Iterable[String] =
      carrier.keys()

    override def get(carrier: Metadata, key: String): String =
      carrier.get(
        Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)
      )
  }

  private def withSemanticAttributes[R, A](effect: ZIO[R, Status, A]): ZIO[Tracing with R, Status, A] =
    Tracing.setAttribute(SemanticAttributes.RPC_SYSTEM.getKey, "grpc") *>
      effect
        .tapBoth(
          status =>
            Tracing.setAttribute(
              SemanticAttributes.RPC_GRPC_STATUS_CODE.getKey,
              status.getCode.value()
            ),
          _ =>
            Tracing.setAttribute(SemanticAttributes.RPC_GRPC_STATUS_CODE.getKey, Status.OK.getCode.value())
        )

  def serverTracingTransform[R]: ZTransform[R, Status, R with Tracing with Has[RequestContext]] =
    new ZTransform[R, Status, R with Tracing with Has[RequestContext]] {

      def effect[A](io: ZIO[R, Status, A]): ZIO[R with Tracing with Has[RequestContext], Status, A] =
        for {
          rc       <- ZIO.service[RequestContext]
          metadata <- rc.metadata.wrap(identity)
          result   <- withSemanticAttributes(io)
                        .spanFrom(
                          propagator,
                          metadata,
                          metadataGetter,
                          rc.methodDescriptor.getFullMethodName,
                          SpanKind.SERVER,
                          { case _ => StatusCode.ERROR }
                        )
        } yield result

      def stream[A](io: ZStream[R, Status, A]): ZStream[R with Tracing with Has[RequestContext], Status, A] =
        ???
    }
}

Update Server Main

Add required layers for Tracing.
Transform the original ZGreeterImpl.

import zio.magic._

object ZServer extends ServerMain {
  private val requirements =
    ZLayer
      .wire[ZEnv with ZGreeterEnv](
        ZEnv.live,
        AppConfig.live.narrow(_.tracing),
        ZTracer.live("hello-server"),
        Tracing.live,
        SttpTracing.live
      )
      .orDie

  def services: ServiceList[Any] =
    ServiceList
      .add(ZGreeterImpl.transform[ZGreeterEnv, Has[RequestContext]](GrpcTracing.serverTracingTransform))
      .provideLayer(requirements)
}

New client

Inject current context into gRPC Metadata for context propagation

object GrpcTracing {

  ...

  private val metadataSetter: TextMapSetter[Metadata] = (carrier, key, value) =>
    carrier.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value)

  val contextPropagationClientInterceptor: ZClientInterceptor[Tracing] = ZClientInterceptor.headersUpdater {
    (_, _, metadata) =>
      metadata.wrapM(Tracing.inject(propagator, _, metadataSetter))
  }

  ...

}

object ZClient extends zio.App {
  private val clientLayer = GreeterClient.live(
    ZManagedChannel(
      ManagedChannelBuilder.forAddress("localhost", 9000).usePlaintext(),
      Seq(GrpcTracing.contextPropagationClientInterceptor)
    )
  )

  ...
}

Start a `Span` for each request

Use ZTransform to record the relevant gRPC attributes.

object GrpcTracing {

  ...

  def clientTracingTransform[R]: ZTransform[R, Status, R with Tracing] =
    new ZTransform[R, Status, R with Tracing] {
      def effect[A](io: ZIO[R, Status, A]): ZIO[R with Tracing, Status, A] = withSemanticAttributes(io)

      def stream[A](io: ZStream[R, Status, A]): ZStream[R with Tracing, Status, A] = ???
    }
}

Unlike the server, we don't have access to a RequestContext object, so we have to set the method name manually.
We also start additional Spans.

object ZClient extends zio.App {

  ...


  private def errorToStatusCode[E]: PartialFunction[E, StatusCode] = { case _ => StatusCode.ERROR }

  private def sayHello(request: HelloRequest) =
    GreeterClient
      .sayHello(request)
      .span(
        GreeterGrpc.METHOD_SAY_HELLO.getFullMethodName,
        SpanKind.CLIENT,
        errorToStatusCode
      )

  private val singleHello = sayHello(HelloRequest("World"))
    .span("singleHello", toErrorStatus = errorToStatusCode)

  private val multipleHellos = ZIO
    .collectAllParN(5)(
      List(
        sayHello(HelloRequest("1", Some(1))),
        sayHello(HelloRequest("2", Some(2))),
        sayHello(HelloRequest("3", Some(3))),
        sayHello(HelloRequest("4", Some(4))),
        sayHello(HelloRequest("5", Some(5)))
      )
    )
    .span("multipleHellos", toErrorStatus = errorToStatusCode)

  private val invalidHello = sayHello(HelloRequest("Invalid", Some(-1))).ignore
    .span("invalidHello", toErrorStatus = errorToStatusCode)
}

Add required layers

object ZClient extends zio.App {

  ...

  private val requirements = ZLayer
    .wire[ZEnv with Tracing](
      ZEnv.live,
      AppConfig.live.narrow(_.tracing),
      ZTracer.live("hello-client"),
      Tracing.live
    ) >+> clientLayer

  def run(args: List[String]): URIO[ZEnv, ExitCode] =
    myAppLogic.provideCustomLayer(requirements).exitCode
}

Showtime

Run Jaeger through Docker

Tracing data can be sent to port 14250.
We can view Jaeger UI at http://localhost:16686.

$ docker run --rm --name jaeger \
  -p 16686:16686 \
  -p 14250:14250 \
  jaegertracing/all-in-one:1.25

Start the server

$ TRACING_ENABLE=true sbt "zio/runMain com.github.tuleism.ZServer"

[info] running (fork) com.github.tuleism.ZServer
[info] Server is running. Press Ctrl-C to stop.

Start the client

$ TRACING_ENABLE=true sbt "zio/runMain com.github.tuleism.ZClient"

[info] running (fork) com.github.tuleism.ZClient
[info] Done
[success] Total time: 12 s

Distributed Tracing in action!

Now we can see the details for multipleHellos:

And which guess is causing the longest delay.

Integration with Logging

Let's add tracing context into log messages following the specification.
We're going to use izumi logstage, our favorite logging library.
See the diffs.

Add logging dependency

val izumiVersion         = "1.0.8"

val loggingDeps = Seq(
  "io.7mind.izumi" %% "logstage-core"          % izumiVersion,
  "io.7mind.izumi" %% "logstage-adapter-slf4j" % izumiVersion
)

Setup logging

Add trace_id, span_id to logging context if current trace context is valid.

object Logging {
  private def baseLogger = IzLogger()

  val live: ZLayer[Has[Tracing.Service], Nothing, Has[LogZIO.Service]] =
    (
      for {
        tracing <- ZIO.service[Tracing.Service]
      } yield LogZIO.withDynamicContext(baseLogger)(
        Tracing.getCurrentSpanContext
          .map(spanContext =>
            if (spanContext.isValid)
              CustomContext(
                "trace_id"    -> spanContext.getTraceId,
                "span_id"     -> spanContext.getSpanId,
                "trace_flags" -> spanContext.getTraceFlags.asHex()
              )
            else
              CustomContext.empty
          )
          .provide(Has(tracing))
      )
    ).toLayer
}

Add a few log messages

E.g for singleHello.

object ZClient extends zio.App {

  ...

  private val singleHello = (
    for {
      _ <- log.info("singleHello")
      _ <- sayHello(HelloRequest("World"))
    } yield ()
  ).span("singleHello", toErrorStatus = errorToStatusCode)

}

Sample Logs

[info] running (fork) com.github.tuleism.ZClient
[info] I 2021-11-01T22:59:10.881 (ZClient.scala:37)  …tuleism.ZClient.singleHello [24:zio-default-async-11] trace_id=9c8a7ebb87381293bc8937a5f7673cb9, span_id=cb7c9a440472e1be, trace_flags=01 singleHello
[info] I 2021-11-01T22:59:14.064 (ZClient.scala:44)  …eism.ZClient.multipleHellos [21:zio-default-async-8 ] trace_id=fe405246fbaa5f876c19f14fa649a99f, span_id=bef19494bef4106e, trace_flags=01 multipleHellos
[info] I 2021-11-01T22:59:18.171 (ZClient.scala:60)  …uleism.ZClient.invalidHello [26:zio-default-async-13] trace_id=be5ccd425e0cfb01fd97274abd0c4d72, span_id=ea6499fb9a7c8d28, trace_flags=01 invalidHello
[info] I 2021-11-01T22:59:18.272 (ZClient.scala:66)  ….tuleism.ZClient.myAppLogic [15:zio-default-async-2 ] Done
[success] Total time: 12 s

Extra notes

If we receive an HTTP 5xx response, we should set the Span status to error according to the semantic convention. However, it is currently not possible with zio-telemetry.
We need a better way to implement tracing for zio-grpc client.

Parallel, Back-pressured Kafka Consumer

Linh Nguyen — Tue, 14 Sep 2021 00:00:00 +0000

Introduction

Almost every Kafka Consumer tutorial structure their code like this:

KafkaConsumer<String, Payment> consumer = new KafkaConsumer<>(props)
// Subscribe to Kafka topics
consumer.subscribe(topics);

while (true) {
    // Poll Kafka for new messages
    ConsumerRecords<String, String> records = consumer.poll(100);
    // Processing logic
    for (ConsumerRecord<String, String> record : records) {
        doSomething(record);
    }
  }
}

Basically we set up our Kafka Consumer, subscribe it to some Kafka topics and then go into an infinite loop where, on each iteration, we poll some messages from these topics and process them one by one. We can call this the poll-then-process loop.

This is fairly simple, easy to put into practice, and people may have been using it in production without any issue. However, there are various problems with this model, which we're going into more details in the next section.

The problems with the poll-then-process loop

1. It is not the "expected" way to poll

Looking at the code above, we developers might think that poll acts as a way to signal demand to Kafka. Our consumer only polls to pull in more messages when it has finished working on previous ones. If its processing rate is slow, Kafka would act as the shock absorber, ensuring we don't lose any message even when the producing rate is much higher.

On the other hand, when processing rate is slow, the interval between consecutive polls also increases. This is problematic since there is a default (5 minutes) upper bound on it with the max.poll.interval.ms configuration:

max.poll.interval.ms

The maximum delay between invocations of poll() when using consumer group management. This places an upper bound on the amount of time that the consumer can be idle before fetching more records. If poll() is not called before expiration of this timeout, then the consumer is considered failed and the group will rebalance in order to reassign the partitions to another member.

In other words, if our consumer doesn't call poll at least once every max.poll.interval.ms, to Kafka, it's as good as dead. When this happens, Kafka follows up with a rebalance process to distribute the current work of the dead consumer to the rest of its consumer group. This introduces more overhead and delay into an already slow processing rate.

Worse yet, if the processing causes slowness in one consumer, chances are, it would cause the same problem for other consumers taking over its work. Moreover, the presumed dead consumer can also cause rebalance when it attempts to rejoin the group on its next poll (remember it's an infinite loop!). Both of these makes rebalance happen again and again, slowing consumption even more.

Now, there is another configuration which can help with this situation:

max.poll.records

The maximum number of records returned in a single call to poll(). Note, that max.poll.records does not impact the underlying fetching behavior. The consumer will cache the records from each fetch request and returns them incrementally from each poll.

With this set to a lower value, our consumer will process fewer messages per poll. Thus, the poll interval will decrease. Alternatively, we can also increase max.poll.interval.ms to a bigger value. This should solve the problem temporarily if we can't move away from the poll-then-process loop. Nevertheless, it is not ideal.

Firstly, these configurations are set when we start our consumers, but whether they work or not depends on the messages or the applications. We may set them specifically for each application, but in the end, we're playing a guessing game and pray that we're lucky.

Secondly, in the worst case, it may take twice the max.poll.interval.ms duration for the rebalance process to start:

Kafka has to wait max.poll.interval.ms to detect that our consumer is not polling anymore.
When Kafka decides to rebalance the group, other consumers are only made aware of this decision on their next poll.

We never want rebalance to take even more time, so setting a higher max.poll.interval.ms is not great.

Finally, these configurations implies that our consumer is "expected" to poll frequently, at least once every max.poll.interval.ms, no matter what kind of processing it is doing. By not incorporating this expectation, the poll-then-process loop is not only misleading to developers but also doomed to fail.

2. Message processing is synchronous

Kafka only guarantees the order of messages within one partition. Messages from different partitions are unrelated and can be processed in parallel. That's why in Kafka, the number of partition in a topic is the unit of parallelism.

In theory, we could easily achieve maximum parallelism by having as many consumers running as the number of partitions on a topic. However, in reality, this is too much overhead, not to mention its impact on increasing the chance for rebalance, since there are more consumers that can come and go.

If we look at our consumer code again, it can subscribe to multiple topics and possibly receive messages from multiple partitions. Yet, when it comes to processing these messages, it does so one by one. This is not optimal.

Now, assuming our processing logic is very simple, could we just use a thread pool to parallelize it? For example, by submitting one processing task to the thread pool, for each message?

Well, kind of. It only works if we don't care about processing ordering and guarantees like at-most-once, at-least-once, etc. So it is not very useful in practice.

A better model

Overview

The many setbacks of the poll-then-process loop come from the fact that different concerns - polling, processing, offsets committing - are mixed together. As a result, when we divide them into separated components, we end up with an improved model which supports parallel processing and back-pressure properly. Each component is described in more details below.

Work Queues

The Work Queues is the communication channel between Poller and Executor:

There is a one-to-one mapping of assigned TopicPartitions to work queues. After each poll, Poller pushes new messages from each partition into its corresponding work queue, preserving the original ordering. Each work queue is also bounded with a configurable size. When full, it back-pressures Poller so that it can follow up with the appropriate actions.
The work queues are asynchronous, which decouples polling and message processing, allowing them to happen independently. This is in contrast to the poll-then-process loop, where they are two sequential steps within a loop.

Poller

In short, Poller encapsulates everything related to poll in Kafka:

It watches out for rebalance events - e.g by registering a ConsumerRebalanceListener - and coordinates other units to handle them.
- For each newly assigned TopicPartition, it set up a new work queue.
- For each revoked (or lost) TopicPartition, it commands both Executor and Offset Manager to wrap up related works and tears down the corresponding work queue.
It polls Kafka periodically using a short (e.g 50ms), configurable interval. Since this is many times lower than the default max.poll.interval.ms, while also not affected by message processing, we avoid the "rebalance storm" that plagues the poll-then-process loop. Kafka won't mistaken our consumer as dead for not polling often enough. In addition, we would know sooner if another rebalance is going to happen.
When we poll more often, we can also use a lower max.poll.interval.ms to speed up the rebalance process.
For each TopicPartition that the Executor cannot keep up with the rate of incoming messages, its corresponding work queue will become full and back-pressure the Poller. The Poller would need to selectively pause this TopicPartition, so that subsequent polls won't pull in more messages from it. When the queue is freed up again, it would resume the same TopicPartition to get new messages starting from the next poll. That's why we can keep polling. That's also why we use a short interval, so that we can "resume" faster.

pause(Collection<TopicPartition> partitions)

Suspend fetching from the requested partitions. Future calls to poll(Duration) will not return any records from these partitions until they have been resumed using resume(Collection).

Executor

Executor acts like a thread pool, where it maintains multiple workers to process the messages:

Executor and the number of workers is tunable to optimize for different workloads e.g CPU bound, I/O bound, etc.
Each work queue is processed by one worker.
One worker can be responsible for multiple work queues.
For each work queue, the worker processes its messages one by one.

With this setup, messages within one partition is processed in order, while messages from different partitions are processed in parallel.

Offset Manager

Each message in Kafka is associated with an offset - an integer number denoting its position in the current partition. By storing this number, we essentially provide a checkpoint for our consumer. If it fails and comes back, it knows from where to continue. As such, it is vital for implementing various processing guarantees in Kafka:

For at-most-once, we need to save $offset + 1 before processing $offset. If our consumer fails before successfully process $offset and restarts, it will continue from $offset + 1 and not reprocess $offset.
For at-least-once, we need to successfully process $offset before saving $offset + 1. If our consumer fails before saving $offset + 1 and restarts, it will continue from and reprocess $offset.
For exactly-once using an external transactional storage - we need to process $offset and save $offset + 1 within one transaction and roll back if anything goes wrong.

With that in mind, we have Offset Manager providing a consistent interface for other components to work with these saved offsets:

If we store offsets within Kafka, it is responsible for manually committing offsets.
If we decide to manage offsets using an external storage, it is responsible for retrieving from and saving into that storage.
It allows Poller and Executor to save offsets either synchronously or asynchronously - in fire and forget fashion.
Offset Manager storing behavior can be configured: in batched, recurring with a timer, etc...

What about Kafka's auto commit? Confluent claims that:

Using auto-commit gives you “at least once” delivery: Kafka guarantees that no messages will be missed, but duplicates are possible.

This is true for delivery, however, it doesn't provide any guarantee for processing:

It's not at-most-once: If some messages are successfully processed, and our consumer crashes before the next auto commit event, these messages are reprocessed.
It's not at-least-once: If auto commit kicks in, and our consumer crashes right afterward, some messages are lost.

Due to this, we always set enable.auto.commit to false and have Offset Manager manage offsets manually.

Achieving processing guarantees

Let's go through a few example use cases to see how the components work together to satisfy different processing guarantees.

At-most-once, offsets managed by Kafka

For at-most-once, we just need to commit offsets before processing the messages. We can do it right before processing each message. However, it doesn't give us a stronger guarantee while introducing more costs. Therefore, Poller is responsible for it. After each poll, it will tell Offset Manager to save these offsets and wait for a success acknowledgement from Kafka before queuing the messages for processing.

Prior to a rebalance event, it just needs to send a fire-and-forget signal to Executor to stop processing. It then takes down the work queues and gets back to wait for rebalance. The lost messages are those still sitting in the queues or in the middle of processing. If we want to optimize for fewer lost without affecting the duration of rebalance, we can use a smaller queue size.

At-least-once, offsets managed by Kafka

For at-least-once, we just need to make sure offsets are only saved after the messages have been processed successfully. Consequently, if we were to process 10 messages, we wouldn't need to save offsets for all of them but only the last one.

In this setup, Executor will emit signals to Offset Manager each time it completes processing for a message. Offset Manager keeps track of the latest offset for each partition - which in total is not that many - and decide when to commit them to Kafka. For example, we can set Offset Manager to commit once every 5 seconds. This happens regardless of whether new messages are coming or not. (Interestingly, this is important for older Kafka if we don't want to lose committed offsets in a low activity Kafka topic, see KAFKA-4682).

Prior to a rebalance event, Poller set a hard deadline and notifies Executor to wrap up its in-flight processing and Offset Manager to follow up with a last commit. If the deadline has passed, or Poller has received responses from others, it takes down the work queues and gets back to wait for rebalance.

To optimize for fewer duplicated processing, we can:

Use a looser deadline, allowing more time for "wrapping up". However, it also increases timing for rebalance.
Set Offset Manager to commit more often.

Exactly-once, offsets managed externally

In this case, offset saving and message processing needs to happen within one transaction. This means Executor and Offset Manager working closely together using synchronous calls to make it happen.

Following a rebalance event, Poller asks Offset Manager for the saved offsets of current assignments. It then seeks to restore the saved positions before resuming to poll.

public void seek(TopicPartition partition, long offset)

Overrides the fetch offsets that the consumer will use on the next poll(timeout).

Prior to a rebalance event, Poller notifies Executor and waits for its response. Executor rolls back its in-flight transactions and return to Poller. Poller then takes down the work queues and gets back to wait for rebalance.

Conclusion

We analyze the various issues with the loop-then-process loop and come up with a more proper model for understanding and implementing Kafka Consumer. The downside is that it is much more complicated, and probably not easy for beginners. We blame this complexity on Kafka and its low level API.

In practice, we probably won't do it ourselves but use a ready-made library that may or may not base on similar models: Alpakka Kafka, Spring for Kafka, zio-kafka, etc... Even then, the proposed model can be useful for evaluating these solutions or implementing new ones.

Zero Downtime Reindex in Elasticsearch

Linh Nguyen — Sun, 22 Aug 2021 00:00:00 +0000

Why reindexing data in Elasticsearch?

Reindexing data is a common operation in working with Elasticsearch. When do we need to do it? Here are a few examples:

Mapping change: when we change how data is analyzed and indexed. For example: changing a field's analyzer, changing an analyzer's definition, etc.
Cluster upgrade or migration: Elasticsearch can only read indices created in the previous major version since 5.x. So if we want to upgrade our Elasticsearch that is several major versions behind, we have to resort to reindexing. For example, Elasticsearch support reindexing from a remote cluster.

Prior arts

We take a look at some available solutions to implement reindexing and analyze their pros and cons.

Elasticsearch's official Reindex API

Pros

Easy: not much work to do since it's readily available.
Fastest solution.
Work for both: reindex within the same cluster or reindex from a remote cluster.
Works well if we can stop indexing data for a while (to wait for the reindex process to finish).

Cons

Hard to reason about consistency if writes still come to the old index during the reindex process.
Meant for indices that must be "paused", so not really zero downtime.
Only work with data that are already in Elasticsearch.

Index Alias Wizardry

In Elasticsearch, alias works as a pointer to one or multiple indices. Many have written about using it for zero downtime reindexing: using one alias or using two, one for read and one for write, etc.

Basically it works like this:

Applications don't query the index directly, but through an alias index_alias.
The alias points to the current index current_index.
When we want to reindex data, we create a new index new_index, probably with a new mapping.
Then we duplicate write: we write to both current_index and new_index at the same time.
Once new_index has caught up with current_index, we update the alias to point to new_index instead of current_index. This can be done atomically.
If anything goes wrong, we just have to revert the final alias change.

Pros

Straightforward.
Revertible.
Can include new fields from the source of truth outside of Elasticsearch.

Cons

Only work within one Elasticsearch cluster.
Requires quite some work, but handles only one use case.
Most solutions don't mention how to keep data consistent between current_index and new_index.

A generic solution

Above solutions are purely about data coming into Elasticsearch. However, not only data, but the applications that read and write these data are also important aspects to consider when we want to rollout a reindex operation.

For example, when we do a major mapping change, we might need to update application code to be able to query from the new index.

For upgrading Elasticsearch server version, we usually have to update the application code since the API might have changed. If we were to use a library like elastic4s, upgrading server means upgrading client library at the same time. To guarantee zero downtime, we need to be able to release application code and the new cluster version independently. How can we achieve that if we don't care about application code?

By incorporating these considerations, we can achieve a generic reindexing solution that works for all mentioned use cases, while having no downtime.

If we call this a recipe for reindexing, then here are the required ingredients:

A data source that supports deterministic replay and multiple concurrent readers. Example: a Kafka topic, a change data capture stream, etc. Why? Because it allows us to have multiple independent writers/indexers that write/index the same set of data into different Elasticsearch indices in possibly different clusters.
Decoupling between readers (applications that query from Elasticsearch) and writers (applications that write/index into Elasticsearch). Why? Because it allows us to separate their deployment (think microservices' independently deployable). As a result, we can scale and upgrade them independently. We will also see later that it allows us to easily upgrade Elasticsearch code from any arbitrary version.
Writers that guarantee idempotent processing under at-least-once delivery. Why? Because combining this with the data source above, we can guarantee data consistency - all old and new indices will eventually have the same data.

Here's how the most basic architecture looks like in normal operation:

We have one writer/indexer current_writer that writes/indexes into current_index.
A separate reader current_reader query from that index and forward the results to the clients. Alternatively, the reader can be a shared library embedded in the client applications.

Let's see how it evolves when we want to reindex data.

Example #1: reindex into the same cluster e.g when mapping changes

First, we create a new index new_index with the updated mapping. Then we deploy a new writer new_writer (either the same code as current_writer or newer code if needed) that writes data into this index.

Once new_index has caught up with current_index, we redirect query requests from current_reader to new_index. This can be a configuration change on current_reader or simply an alias update on Elasticsearch.

If anything goes wrong, we just need to revert the final redirect change since current_index is still up to date.

For complex mapping changes, we might also need to update the reader code. In these cases, we can release a new version of reader - new_reader that queries directly from new_index. Once everything looks good, we redirect requests to new_reader instead of current_reader. Again, we just need to revert the last step if anything goes wrong.

Example #2: reindex into another cluster e.g when upgrading or doing migration

First, we create a new index new_index in the new cluster. Then we deploy a new writer new_writer that writes data into this index. new_writer's code has been updated with newer Elasticsearch client library to talk to the new cluster.

Next, similar to example #1, we release a new reader new_reader that query from new_index. new_reader's code has been updated with newer Elasticsearch client library. If things looks good, we switch clients to use new_reader. Again, this is fully revertible.

Final steps

Once we're happy with the result, we just need to retire those "current" indices, writers and readers. The diagram now looks exactly like the original.

Comparison to other solutions

Pros

Not more complicated than using Index Alias.
Generic solution that works for all use cases.
All steps are incremental and revertible.
Zero downtime. There is a clear plan for rolling back.
Guarantee eventual consistency because of the requirements on the data source and writers.
Freedom in deciding when the switch/redirect happens. If implemented properly, data is kept "in sync" in near real time.

Cons

Not the most performance method. However, it's very flexible. For example, we can scale writers up if we want to speed up reindexing.
It requires organizational maturity in CI/CD/Ops to automate the process successfully. With that said, manually doing it is also not that hard since the steps are all laid out.

Smart Constructors in Scala

Linh Nguyen — Sun, 16 Aug 2020 10:10:20 +0000

Introduction

Sometimes, we need guarantees about the values in our program beyond what can be accomplished with the usual type system checks. Take for example an Email type, which contains an email address:

case class Email(address: String)

The problem with Email is that address can contain any value of type String, even if they are not valid email addresses. Due to that, code that uses Email cannot assume that it's always a valid email address. Thus, we want to:

Constrain the set of possible value for address to be the set of valid email addresses.
Make sure every instance of Email satisfies this constraint. In other words, every Email instance contains a valid email address.

Smart Constructors is one solution for this: instead of normal constructors, we force construction through "smart" functions that only return Email instances when the input passes validation.

In our case, we can use a function that returns Option[Email], Either[Error, Email], etc depends on the validation result. For demonstration, I use a very simple regex:

def fromString(v: String): Option[Email] =
  NaiveEmailRegex.findFirstIn(v).filter(_ == v).map(_ => new Email(v))

Currently, I'm aware of 3 ways to implement Smart Constructors in Scala.

The Straightforward One: `sealed trait`

Using trait, it is very straightforward to implement:

sealed trait Email {
  def value: String
}

object Email {
  def fromString(v: String): Option[Email] =
    NaiveEmailRegex.findFirstIn(v).filter(_ == v).map { _ =>
      new Email {
        override def value: String = v
      }
    }
}

With sealed, we disallow attempts to extends Email from outside of this source file.
Email.fromString is the only way to construct an Email from outside of this source file.

The Confusing One: `final case class private`

When it comes to modeling data, especially those that are immutable, case classes is the method of choice for Scala. Compared to traits, case classes come with many useful features out of the box: pattern matching, equality comparison, copying, etc.

On the other hand, some of these make it harder to implement smart constructors. The implementations also differ depends on which Scala version (and compiler flags) you are using, causing confusion for beginners.

Let's begin with an initial implementation in Scala 2.11:

final case class Email private (value: String)

object Email {
  def fromString(v: String): Option[Email] = ...
}

With final, we disallow attempts to extends Email.
With private, construction through new Email(value) can only takes place within the current source file.

However, it is still possible to create an invalid Email instance:

By using copy(value = "badValue") to create a shallow copy of an instance of Email with a bad value.
By using Email(value = "badValue"), which translates to calling apply() on the Email's companion object to construct an Email with a bad value.

To fix these problems, we need to either hide or override the copy() and apply() functions. For example, with apply(), there are 2 choices:

Override it with private modifier to hide the apply() function, disallowing construction through Email(value).
Override it with an alternative return type and implementation: Email(value) is possible, but let's say, return Option[Email] instead.

2.11.x and 2.12.x

It turns out it is not possible to override apply() in vanilla Scala 2.11.x. See this StackOverflow question for more details.

To do that, we need to add the compiler flag -Xsource:2.12, which is only available after 2.11.11:

Allow custom apply and unapply methods in case class companions. Also allows case class constructors and apply methods to be private. (In 2.11.11, -Xsource:2.12 is needed to enable these changes. In Scala 2.12.2, they are on by default.)

With this, our code becomes:

final case class Email private (value: String) {
  private def copy(): Unit = ()
}

object Email {
  private def apply(value: String): Email = ???

  def fromString(v: String): Option[Email] = ...
}

This is the same for 2.12.x without any compiler flag.

2.13.x and Dotty

With vanilla 2.13.x, the code is the same as 2.12.x. However, from 2.13.2, we can reduce the boilerplate by enabling -Xsource:3. From scala/scala#7702:

Backport from dotty:

If a case class constructor is private or private[foo]: the synthesized copy and apply methods will have the same access modifier.

If a case class constructor is protected or protected[foo]: the synthesized copy method will have the same access modifier. The synthesized apply method will remain public, because protected doesn't make sense in an object. Obviously, if a user defines a custom copy or apply method, that one—including its access modifier—will take precedence.

We end up with the original code that we begin with:

final case class Email private (value: String)

object Email {
  def fromString(v: String): Option[Email] = ...
}

Phew! It only takes a few Scala major releases to get what we want.

The Tricky One: `sealed abstract case class private`

We don't normally see abstract being used on a case class. However, for smart constructors, it works wonderfully:

sealed abstract case class Email private (value: String)

object Email {
  def fromString(v: String): Option[Email] =
    NaiveEmailRegex.findFirstIn(v).filter(_ == v)
      .map(_ => new Email(v) {})
}

With sealed, we disallow attempts to extends Email from outside this source file.
With abstract, we disallow construction through new Email().
With abstract, copy() and apply() are not automatically generated. We don't have to worry about them.
It just work™! See the original gist for more details.

Testing

We verify that our smart constructors implementation works by using a test suite to check that allowed code Compiles and the rest DoesNotCompile.

Below is the test suite for sealed abstract case class private. The rest can be found in the code repository.

"AbstractCaseClassExample" in {
  import AbstractCaseClassExample.Email
  val email = Email.fromString(exampleEmail).value // apply now return Option
  email.value mustBe exampleEmail // access ok

  // case class's unapply()
  assertCompiles(
    """
      |email match { case Email(value) => value }
      |""".stripMargin
  )
  // companion's apply()
  assertDoesNotCompile(
    """
      |Email(exampleEmail)
      |""".stripMargin
  )
  // public constructor
  assertDoesNotCompile(
    """
      |new Email(exampleEmail)
      |""".stripMargin
  )
  // extends trait
  assertDoesNotCompile(
    """
      |new Email {
      |  override def value: String = exampleEmail
      |}
      |""".stripMargin
  )
  // case class's copy()
  assertDoesNotCompile(
    """
      |email.copy(value = exampleEmail)
      |""".stripMargin
  )
}

Summary

All of these allow us to implement smart constructors. My favorite one is sealed abstract case class private because:

It allows us to use case class (instead of trait, which is more verbose) and most of its goodness.
It looks the same in all Scala version. Cross-building is a breeze.

DEV Community: Linh Nguyen

OpenTelemetry Distributed Tracing with ZIO

Introduction

Initial implementation

Add the new flag guess:

Add zio-grpc dependency

Set up build.sbt

Client implementation

Server implementation

Run it

Common tracing requirements

Add new dependencies

Add a config layer

Add a Tracer layer

New server

Instrument the HTTP client

Instrument the gRPC server

Update Server Main

New client

Inject current context into gRPC Metadata for context propagation

Start a Span for each request

Add required layers

Showtime

Run Jaeger through Docker

Start the server

Start the client

Distributed Tracing in action!

Integration with Logging

Add logging dependency

Setup logging

Add a few log messages

Sample Logs

Extra notes

Parallel, Back-pressured Kafka Consumer

Introduction

The problems with the poll-then-process loop

1. It is not the "expected" way to poll

2. Message processing is synchronous

A better model

Overview

Work Queues

Poller

Executor

Offset Manager

Achieving processing guarantees

At-most-once, offsets managed by Kafka

At-least-once, offsets managed by Kafka

Exactly-once, offsets managed externally

Conclusion

Zero Downtime Reindex in Elasticsearch

Why reindexing data in Elasticsearch?

Prior arts

Elasticsearch's official Reindex API

Pros

Cons

Index Alias Wizardry

Pros

Cons

A generic solution

Example #1: reindex into the same cluster e.g when mapping changes

Example #2: reindex into another cluster e.g when upgrading or doing migration

Final steps

Comparison to other solutions

Pros

Cons

Smart Constructors in Scala

Introduction

The Straightforward One: sealed trait

The Confusing One: final case class private

2.11.x and 2.12.x

2.13.x and Dotty

The Tricky One: sealed abstract case class private

Testing

Summary

Add the new flag `guess`:

Set up `build.sbt`

Add a `Tracer` layer

Start a `Span` for each request

The Straightforward One: `sealed trait`

The Confusing One: `final case class private`

The Tricky One: `sealed abstract case class private`