Graceful shutdown of a stream for a single subscription

docs/consuming-kafka-topics-using-zio-streams.md

@@ -74,3 +74,70 @@ Each time a new partition is assigned to the consumer, a new partition stream is
 ensures only a single stream is doing commits.
 
 In this example, each partition is processed in parallel, on separate fibers, while a single fiber is doing commits.
+
+## Controlled shutdown (experimental)
+
+The examples above will keep processing records forever, or until the fiber is interrupted, typically at application shutdown. When interrupted, some records may be 'in-flight', e.g. being processed by one of the stages of your consumer stream user code. Those records will be processed partly and their offsets may not be committed. For fast shutdown in an at-least-once processing scenario this is fine.
+
+zio-kafka also supports a _graceful shutdown_, where the fetching of records for the subscribed topics/partitions is stopped, the streams are ended and all downstream stages are completed, allowing in-flight records to be fully processed.
+
+Use the `*StreamWithControl` variants of `plainStream`, `partitionedStream` and `partitionedAssignmentStream` for this purpose. These methods return a `StreamControl` object allowing you to stop fetching records and terminate the execution of the stream gracefully.
+
+The `StreamControl` can be used with `Consumer.runWithGracefulShutdown`, which can gracefully terminate the stream upon fiber interruption. This is useful for a controlled shutdown when your application is terminated:
+
+```scala
+import zio.Console.printLine
+import zio.kafka.consumer._
+import zio._
+
+ZIO.scoped {
+  for {
+    control <- consumer.partitionedStreamWithControl(
+      Subscription.topics("topic150"),
+      Serde.string,
+      Serde.string
+    )
+    _ <- consumer.runWithGracefulShutdown(control, shutdownTimeout = 30.seconds) { stream =>
+        stream.flatMapPar(Int.MaxValue) { case (topicPartition, partitionStream) =>
+          partitionStream
+            .tap(record => printLine(s"key: ${record.key}, value: ${record.value}"))
+            .map(_.offset)
+        }
+        .aggregateAsync(Consumer.offsetBatches)
+        .mapZIO(_.commit)
+        .runDrain
+      }
+  } yield ()
+}
+```
+
+For more control over when to end the stream, use the `StreamControl` construct like this:
+
+```scala
+import zio.Console.printLine
+import zio.kafka.consumer._
+
+ZIO.scoped {
+    for {
+      control <- consumer.partitionedStreamWithControl(
+          Subscription.topics("topic150"),
+          Serde.string,
+          Serde.string
+        )
+      fiber <- control.stream.flatMapPar(Int.MaxValue) { case (topicPartition, partitionStream) =>
+            partitionStream
+              .tap(record => printLine(s"key: ${record.key}, value: ${record.value}"))
+              .map(_.offset)
+          }
+          .aggregateAsync(Consumer.offsetBatches)
+          .mapZIO(_.commit)
+          .runDrain
+          .forkScoped
+      // At some point in your code, when some condition is met
+      _ <- control.end // Stop fetching more records
+      _ <- fiber.join // Wait for graceful completion of the stream
+    } yield ()
+}
+```
+
+As of zio-kafka 3.0, this functionality is experimental. If no issues are reported and the API has good usability, it will eventually be marked as stable.

zio-kafka-test/src/test/scala/zio/kafka/consumer/internal/PartitionStreamControlSpec.scala

@@ -56,7 +56,17 @@ object PartitionStreamControlSpec extends ZIOSpecDefault {
           _             <- control.offerRecords(records)
           _             <- control.end
           pulledRecords <- control.stream.runCollect
-        } yield assertTrue(records == pulledRecords)
+          hasEnded      <- control.hasEnded
+        } yield assertTrue(records == pulledRecords && hasEnded)
+      },
+      test("offering records after end will fail") {
+        for {
+          control <- createTestControl
+          records = createTestRecords(3)
+          _    <- control.offerRecords(records)
+          _    <- control.end
+          exit <- control.offerRecords(records).exit
+        } yield assertTrue(exit.isFailure)
       },
       test("halt completes the stream with a TimeoutException") {
         for {

zio-kafka/src/main/scala/zio/kafka/consumer/Consumer.scala

@@ -122,6 +122,63 @@ trait Consumer {
     bufferSize: Int = 4
   ): ZStream[R, Throwable, CommittableRecord[K, V]]
 
+  /**
+   * Like [[plainStream]] but returns an object to control the stream.
+   *
+   * The Scope provided to the returned effect controls the lifetime of the subscription. The subscription is
+   * unsubscribed when the scope is ended. Calling [[StreamControl.end]] stops fetching data for the subscription
+   * partitions but will not unsubscribe until the Scope is ended.
+   *
+   * WARNING: this is an EXPERIMENTAL API and may disappear or change in an incompatible way without notice in any
+   * zio-kafka version.
+   */
+  def plainStreamWithControl[R, K, V](
+    subscription: Subscription,
+    keyDeserializer: Deserializer[R, K],
+    valueDeserializer: Deserializer[R, V],
+    bufferSize: Int = 4
+  ): ZIO[Scope with R, Throwable, StreamControl[R, Throwable, CommittableRecord[K, V]]]
+
+  /**
+   * Like [[partitionedStream]] but returns an object to control the stream.
+   *
+   * The Scope provided to the returned effect controls the lifetime of the subscription. The subscription is
+   * unsubscribed when the scope is ended. Calling [[StreamControl.end]] stops fetching data for the subscription
+   * partitions but will not unsubscribe until the Scope is ended.
+   *
+   * WARNING: this is an EXPERIMENTAL API and may disappear or change in an incompatible way without notice in any
+   * zio-kafka version.
+   */
+  def partitionedStreamWithControl[R, K, V](
+    subscription: Subscription,
+    keyDeserializer: Deserializer[R, K],
+    valueDeserializer: Deserializer[R, V]
+  ): ZIO[
+    Scope with R,
+    Throwable,
+    StreamControl[R, Throwable, (TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])]
+  ]
+
+  /**
+   * Like [[partitionedAssignmentStream]] but returns an object to control the stream.
+   *
+   * The Scope provided to the returned effect controls the lifetime of the subscription. The subscription is
+   * unsubscribed when the scope is ended. Calling [[StreamControl.end]] stops fetching data for the subscription
+   * partitions but will not unsubscribe until the Scope is ended.
+   *
+   * WARNING: this is an EXPERIMENTAL API and may disappear or change in an incompatible way without notice in any
+   * zio-kafka version.
+   */
+  def partitionedAssignmentStreamWithControl[R, K, V](
+    subscription: Subscription,
+    keyDeserializer: Deserializer[R, K],
+    valueDeserializer: Deserializer[R, V]
+  ): ZIO[Scope, Throwable, StreamControl[
+    R,
+    Throwable,
+    Chunk[(TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])]
+  ]]
+
   /**
    * Stops consumption of data, drains buffered records, and ends the attached streams while still serving commit
    * requests.
@@ -341,6 +398,56 @@ object Consumer {
     case object Latest   extends AutoOffsetStrategy
     case object None     extends AutoOffsetStrategy
   }
+
+  /**
+   * Takes a StreamControl for some stream and runs the given ZIO workflow on that stream such that, when interrupted,
+   * stops fetching records and gracefully waits for the ZIO workflow to complete.
+   *
+   * This is useful for running streams from within your application's Main class, such that streams are cleanly stopped
+   * when the application is shutdown (for example by your container runtime).
+   *
+   * WARNING: this is an EXPERIMENTAL API and may disappear or change in an incompatible way without notice in any
+   * zio-kafka version.
+   *
+   * @param streamControl
+   *   Result of one of the Consumer's methods returning a [[StreamControl]]
+   * @param shutdownTimeout
+   *   Timeout for the workflow to complete after initiating the graceful shutdown
+   * @param withStream
+   *   Takes the stream as input and returns a ZIO workflow that processes the stream. As in most programs the given
+   *   workflow runs until an external interruption, the result value (Any type) is meaningless. `withStream` is
+   *   typically something like `stream => stream.mapZIO(record => ZIO.debug(record)).mapZIO(_.offset.commit)`
+   */
+  def runWithGracefulShutdown[R, E, A](
+    control: StreamControl[R, E, A],
+    shutdownTimeout: Duration
+  )(
+    withStream: ZStream[R, E, A] => ZIO[R, E, Any]
+  ): ZIO[R, E, Any] =
+    ZIO.scoped[R] {
+      for {
+        fib <-
+          withStream(control.stream)
+            .onInterrupt(
+              ZIO.logError("withStream in runWithGracefulShutdown was interrupted, this should not happen")
+            )
+            .tapErrorCause(cause => ZIO.logErrorCause("Error in withStream fiber in runWithGracefulShutdown", cause))
+            .forkScoped
+        result <-
+          fib.join.onInterrupt(
+            control.end *>
+              fib.join
+                .timeout(shutdownTimeout)
+                .someOrElseZIO(
+                  ZIO.logError(
+                    "Timeout waiting for `withStream` to shut down gracefully. Not all in-flight records may have been processed."
+                  )
+                )
+                .tapErrorCause(cause => ZIO.logErrorCause("Stream failed while awaiting its graceful shutdown", cause))
+                .ignore
+          )
+      } yield result
+    }
 }
 
 private[consumer] final class ConsumerLive private[consumer] (
@@ -392,7 +499,8 @@ private[consumer] final class ConsumerLive private[consumer] (
 
     ZStream.unwrapScoped {
       for {
-        stream <- runloopAccess.subscribe(subscription)
+        streamControl <- runloopAccess.subscribe(subscription)
+        stream = streamControl.stream
       } yield stream
         .map(_.exit)
         .flattenExitOption
@@ -410,6 +518,33 @@ private[consumer] final class ConsumerLive private[consumer] (
     }
   }
 
+  override def partitionedAssignmentStreamWithControl[R, K, V](
+    subscription: Subscription,
+    keyDeserializer: Deserializer[R, K],
+    valueDeserializer: Deserializer[R, V]
+  ): ZIO[Scope, Throwable, StreamControl[
+    R,
+    Throwable,
+    Chunk[(TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])]
+  ]] = {
+    val onlyByteArraySerdes: Boolean = (keyDeserializer eq Serde.byteArray) && (valueDeserializer eq Serde.byteArray)
+    for {
+      streamControl <- runloopAccess.subscribe(subscription)
+    } yield streamControl.map(
+      _.map(_.exit).flattenExitOption.map {
+        _.collect {
+          case (tp, partitionStream) if Subscription.subscriptionMatches(subscription, tp) =>
+            val stream: ZStream[R, Throwable, CommittableRecord[K, V]] =
+              if (onlyByteArraySerdes)
+                partitionStream.asInstanceOf[ZStream[R, Throwable, CommittableRecord[K, V]]]
+              else partitionStream.mapChunksZIO(_.mapZIO(_.deserializeWith(keyDeserializer, valueDeserializer)))
+
+            tp -> stream
+        }
+      }
+    )
+  }
+
   override def partitionedStream[R, K, V](
     subscription: Subscription,
     keyDeserializer: Deserializer[R, K],
@@ -428,6 +563,32 @@ private[consumer] final class ConsumerLive private[consumer] (
       bufferSize = bufferSize
     )(_._2)
 
+  override def plainStreamWithControl[R, K, V](
+    subscription: Subscription,
+    keyDeserializer: Deserializer[R, K],
+    valueDeserializer: Deserializer[R, V],
+    bufferSize: Int = 4
+  ): ZIO[Scope with R, Throwable, StreamControl[R, Throwable, CommittableRecord[K, V]]] =
+    partitionedStreamWithControl(subscription, keyDeserializer, valueDeserializer).map(
+      _.map(
+        _.flatMapPar(
+          n = Int.MaxValue,
+          bufferSize = bufferSize
+        )(_._2)
+      )
+    )
+
+  override def partitionedStreamWithControl[R, K, V](
+    subscription: Subscription,
+    keyDeserializer: Deserializer[R, K],
+    valueDeserializer: Deserializer[R, V]
+  ): ZIO[
+    Scope with R,
+    Throwable,
+    StreamControl[R, Throwable, (TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])]
+  ] = partitionedAssignmentStreamWithControl(subscription, keyDeserializer, valueDeserializer)
+    .map(_.map(_.flattenChunks))
+
   override def stopConsumption: UIO[Unit] =
     ZIO.logDebug("stopConsumption called") *>
       runloopAccess.stopConsumption

zio-kafka/src/main/scala/zio/kafka/consumer/StreamControl.scala

@@ -0,0 +1,33 @@
+package zio.kafka.consumer
+import zio.UIO
+import zio.stream.ZStream
+
+trait StreamControl[-R, +E, +A] {
+
+  /**
+   * The stream associated with this subscription
+   *
+   * The stream should be run at most once. Running it more than once will result in chunks of records being divided
+   * over the streams. After ending, running the stream another time will not produce records
+   */
+  def stream: ZStream[R, E, A]
+
+  /**
+   * Stops fetching data for all partitions associated with this subscription. The stream will end and the effect
+   * running the stream will eventually complete.
+   *
+   * @return
+   *   Effect that will complete immediately.
+   */
+  def end: UIO[Unit]
+}
+
+object StreamControl {
+  implicit class StreamControlOps[-R, +E, +A](streamControl: StreamControl[R, E, A]) {
+    def map[R1 <: R, E1 >: E, B](f: ZStream[R, E, A] => ZStream[R1, E1, B]): StreamControl[R1, E1, B] =
+      new StreamControl[R1, E1, B] {
+        def stream = f(streamControl.stream)
+        def end    = streamControl.end
+      }
+  }
+}