A rebalance callback that outlives the `ToTry` timeout leaves the `TopicFlow` blocked forever

[tobiajo]

kafka-flow master at 1a062dc, skafka 21.0.3, cats-helper 3.13.1, cats-effect 3.7.1, kafka-clients 4.3.1.

Summary

Partition recovery runs inside the rebalance callback, which skafka executes through cats-helper's ToTry (one minute by default). When recovery takes longer the callback fails with a TimeoutException, and the TopicFlow keeps its semaphore permit forever: every later apply, add, remove and its release block, the error never reaches the retry, and the process stays up without polling.

Where the budget comes from

RebalanceListener lifts flow.add(partitions) into the callback (RebalanceListener.scala#L31). TopicFlow.add allocates a PartitionFlow whose acquire is the whole partition recovery (PartitionFlow.scala#L125, #L344). skafka runs each lifted effect on the poll thread through the application's ToTry[F] (ConsumerConverters.scala#L49-L73, RebalanceCallback.scala#L100). cats-helper's default is ioToTry(1.minute) (ToTry.scala), applied per lifted effect.

Why the timeout leaks the permit

ioToTry splits the effect with IO.syncStep, which walks it on the calling thread. With SyncIO's Sync instance the walk goes inside uncancelable and past onCancel, because SyncIO can never be cancelled. The IO it returns lacks both. unsafeRunTimed then cancels that IO on timeout and returns at once ("does not backpressure on the completion of finalizers").

TopicFlow.safeguard wraps each call in semaphore.permit.use { ... }.uncancelable (TopicFlow.scala#L180-L188). The permit is taken in the part the walk executes; the rest of recovery is in the IO the walk returns. That IO runs without the uncancelable and without the onCancel that would return the permit, so the timeout cancels an unprotected effect and the permit is gone.

Why the error never propagates

The stream fails, but releasing the TopicFlow needs the leaked permit inside uncancelable, so it blocks forever. The retry only sees the error after the resources are released. The member keeps its assignment until max.poll.interval.ms (default 5 min), when the heartbeat thread leaves the group.

An operator sees one kafka-clients ERROR with TimeoutException: 1 minute, then nothing.

Reproduction

The spec drives RebalanceListener through skafka's RebalanceCallback.run with ToTry.ioToTry, as the Java bridge does, against a PartitionFlowOf whose recovery takes 3 s.

• A (production path, 500 ms budget): the callback fails, recovery finishes later, then apply, the release and cancellation all block.
• B (30 s budget): everything completes.
• C (same callback as plain IO, cancelled at 500 ms, no ToTry split): the permit is returned. The split is the cause.
• D (ToTry alone on IO.sleep(3s).uncancelable): Failure at 500 ms, the sleep was cancelled despite the mask.

[repro] === A ===
[repro] 66ms recovery started
[repro] 567ms callback outcome: Failure(java.util.concurrent.TimeoutException: 500 milliseconds)
[repro] 3072ms recovery finished
[repro] 6569ms probing after 2x recovery duration
[repro] 9578ms TopicFlow.apply: still blocked after 3 seconds
[repro] 12585ms TopicFlow release: still blocked after 3 seconds
[repro] 15795ms cancel of a blocked TopicFlow.apply: still blocked after 3 seconds
[repro] === C ===
[repro] 5ms recovery started
[repro] 508ms cancelling the callback fiber
[repro] 3010ms recovery finished
[repro] 3013ms cancel returned
[repro] 3014ms callback outcome: Canceled()
[repro] 9019ms probing after 2x recovery duration
[repro] 9023ms TopicFlow.apply: completed as Succeeded(IO(())) within 3 seconds
[repro] 9025ms partition flow released
[repro] 9027ms TopicFlow release: completed as Succeeded(IO(())) within 3 seconds
[repro] 9233ms cancel of a blocked TopicFlow.apply: completed as Succeeded(IO(())) within 3 seconds
[repro] D 506ms toTry returned: Failure(java.util.concurrent.TimeoutException: 500 milliseconds)

ToTryTimeoutReproSpec.scala

package com.evolutiongaming.kafka.flow

import cats.data.{NonEmptyMap, NonEmptySet}
import cats.effect.unsafe.implicits.global
import cats.effect.{IO, Resource}
import cats.syntax.all.*
import com.evolutiongaming.catshelper.{LogOf, Runtime, ToTry}
import com.evolutiongaming.kafka.flow.kafka.{Consumer, ScheduleCommit}
import com.evolutiongaming.skafka.*
import com.evolutiongaming.skafka.consumer.{ConsumerGroupMetadata, ConsumerRecord, ConsumerRecords, RebalanceListener1}
import munit.FunSuite
import scodec.bits.ByteVector

import java.util.concurrent.TimeoutException
import scala.concurrent.duration.*
import scala.util.{Failure, Success, Try}

/** Drives kafka-flow's RebalanceListener through skafka's RebalanceCallback.run with cats-helper's ToTry.ioToTry, as
  * skafka's Java listener bridge does. Calls that may block forever run in their own fibers with a bounded join, so
  * the test itself cannot hang.
  */
class ToTryTimeoutReproSpec extends FunSuite {

  private implicit val logOf: LogOf[IO]     = LogOf.empty[IO]
  private implicit val runtime: Runtime[IO] = Runtime.lift[IO]

  override def munitTimeout: Duration = 180.seconds

  private val topic     = "topic"
  private val partition = Partition.min
  private val recovery  = 3.seconds

  private val consumer = new Consumer[IO] {
    def subscribe(topics: NonEmptySet[Topic], listener: RebalanceListener1[IO]): IO[Unit] = IO.unit
    def poll(timeout: FiniteDuration): IO[ConsumerRecords[String, ByteVector]]    = ConsumerRecords.empty.pure[IO]
    def commit(offsets: NonEmptyMap[TopicPartition, OffsetAndMetadata]): IO[Unit] = IO.unit
    def groupMetadata: IO[Option[ConsumerGroupMetadata]] = none[ConsumerGroupMetadata].pure[IO]
  }

  private def bounded[A](label: String, fa: IO[A], limit: FiniteDuration): IO[String] =
    fa.start.flatMap { fiber =>
      IO.race(fiber.join, IO.sleep(limit)).map {
        case Left(outcome) => s"$label: completed as $outcome within $limit"
        case Right(_)      => s"$label: still blocked after $limit"
      }
    }

  /** Builds a TopicFlow whose partition flow takes `recovery` to construct, then runs `drive` (the assignment
    * callback, in the shape under test) and probes whether the TopicFlow is still usable afterwards.
    */
  private def scenario(name: String)(drive: (TopicFlow[IO], String => IO[Unit]) => IO[String]): IO[(String, String, String)] =
    for {
      _     <- IO.println(s"[repro] === $name ===")
      start <- IO.monotonic
      record = (s: String) => IO.monotonic.flatMap(t => IO.println(s"[repro] ${(t - start).toMillis}ms $s"))
      partitionFlowOf = new PartitionFlowOf[IO] {
        def apply(a: PartitionAssignment[IO], sc: ScheduleCommit[IO]): Resource[IO, PartitionFlow[IO]] =
          Resource
            .make(
              record("recovery started") *> IO.sleep(recovery).onCancel(record("recovery CANCELLED")) *>
                record("recovery finished")
            )(_ => record("partition flow released"))
            .as(new PartitionFlow[IO] {
              def apply(records: List[ConsumerRecord[String, ByteVector]]): IO[Unit] = IO.unit
            })
      }
      allocated <- TopicFlow.of(consumer, topic, partitionFlowOf).allocated
      (flow, release) = allocated
      outcome        <- drive(flow, record)
      _              <- record(s"callback outcome: $outcome")
      // give the detached recovery ample time to finish (or to be cancelled) before probing
      _ <- IO.sleep(recovery * 2)
      _ <- record("probing after 2x recovery duration")
      // what the next poll would do
      applyOutcome <- bounded("TopicFlow.apply", flow.apply(ConsumerRecords.empty), 3.seconds)
      _            <- record(applyOutcome)
      // what the stream's resource teardown would do
      releaseOutcome <- bounded("TopicFlow release", release, 3.seconds)
      _              <- record(releaseOutcome)
      // what cancelling the flow's fiber on shutdown would do (KafkaFlow.resource runs the stream in .background)
      blocked       <- flow.apply(ConsumerRecords.empty).start
      _             <- IO.sleep(200.millis)
      cancelOutcome <- bounded("cancel of a blocked TopicFlow.apply", blocked.cancel, 3.seconds)
      _             <- record(cancelOutcome)
    } yield (applyOutcome, releaseOutcome, cancelOutcome)

  private def assigned(flow: TopicFlow[IO]) =
    RebalanceListener[IO](Map(topic -> flow)).onPartitionsAssigned(NonEmptySet.of(TopicPartition(topic, partition)))

  test("A: ToTry timeout shorter than recovery, via skafka's RebalanceCallback.run (the production path)") {
    val toTry = ToTry.ioToTry(500.millis)
    val (applyOutcome, releaseOutcome, cancelOutcome) = scenario("A") { (flow, _) =>
      // what skafka's RebalanceListenerJ.onPartitionsAssigned does on the Kafka poll thread
      IO.blocking(assigned(flow).run(new Consumer.NoopRebalanceConsumer)(toTry)).map(_.toString)
    }.unsafeRunSync()
    println(s"[repro] SUMMARY A apply -> $applyOutcome")
    println(s"[repro] SUMMARY A release -> $releaseOutcome")
    println(s"[repro] SUMMARY A cancel -> $cancelOutcome")
    assert(applyOutcome.contains("still blocked"), applyOutcome)
    assert(releaseOutcome.contains("still blocked"), releaseOutcome)
    assert(cancelOutcome.contains("still blocked"), cancelOutcome)
  }

  test("B: ToTry timeout longer than recovery, same path (the healthy case)") {
    val toTry = ToTry.ioToTry(30.seconds)
    val (applyOutcome, releaseOutcome, cancelOutcome) = scenario("B") { (flow, _) =>
      IO.blocking(assigned(flow).run(new Consumer.NoopRebalanceConsumer)(toTry)).map(_.toString)
    }.unsafeRunSync()
    println(s"[repro] SUMMARY B apply -> $applyOutcome")
    println(s"[repro] SUMMARY B release -> $releaseOutcome")
    println(s"[repro] SUMMARY B cancel -> $cancelOutcome")
    assert(applyOutcome.contains("completed"), applyOutcome)
    assert(releaseOutcome.contains("completed"), releaseOutcome)
    assert(cancelOutcome.contains("completed"), cancelOutcome)
  }

  test("C: the same cancellation delivered to the callback run as plain IO (mask intact)") {
    val (applyOutcome, releaseOutcome, cancelOutcome) = scenario("C") { (flow, record) =>
      // skafka's toF interprets the callback into IO without the ToTry sync/async split
      val io = assigned(flow).toF(new Consumer.NoopRebalanceConsumer)
      io.start.flatMap { fiber =>
        IO.sleep(500.millis) *> record("cancelling the callback fiber") *> fiber.cancel *>
          record("cancel returned") *> fiber.join.map(_.toString)
      }
    }.unsafeRunSync()
    println(s"[repro] SUMMARY C apply -> $applyOutcome")
    println(s"[repro] SUMMARY C release -> $releaseOutcome")
    println(s"[repro] SUMMARY C cancel -> $cancelOutcome")
    assert(applyOutcome.contains("completed"), applyOutcome)
    assert(releaseOutcome.contains("completed"), releaseOutcome)
    assert(cancelOutcome.contains("completed"), cancelOutcome)
  }

  test("D: ToTry alone on a sleep inside uncancelable: the mask is dropped by the sync/async split") {
    val toTry = ToTry.ioToTry(500.millis)
    val program = for {
      start <- IO.monotonic
      record = (s: String) => IO.monotonic.flatMap(t => IO.println(s"[repro] D ${(t - start).toMillis}ms $s"))
      io     = (IO.sleep(recovery).onCancel(record("sleep CANCELLED")) *> record("sleep finished")).uncancelable
      result <- IO.blocking(toTry(io))
      _      <- record(s"toTry returned: $result")
      _      <- IO.sleep(recovery * 2)
    } yield result
    val result: Try[Unit] = program.unsafeRunSync()
    result match {
      case Failure(_: TimeoutException) => ()
      case Success(_)                   => fail("expected a timeout")
      case Failure(other)               => fail(s"unexpected: $other")
    }
  }
}

Related

• Keep the TopicFlow permit when a rebalance callback times out [AWAIT cats-helper#425] #936 puts an async boundary in front of the guarded region so the split cannot land inside it.
• Make ioToTry's timeout respect uncancelable and run finalizers cats-helper#425 makes ioToTry stop at cancellation structure instead of walking through it, keeping the synchronous fast path.
• skafka#125 (closed) discussed deriving the timeout from max.poll.interval.ms.