Implement Graceful Shutdown

src/main/scala/zio/kafka/client/Consumer.scala

@@ -32,6 +32,9 @@ class Consumer private (
   ): BlockingTask[Map[TopicPartition, Long]] =
     consumer.withConsumer(_.endOffsets(partitions.asJava, timeout.asJava).asScala.mapValues(_.longValue()).toMap)
 
+  def gracefulShutdown: UIO[Unit] =
+    runloop.deps.gracefulShutdown
+
   def listTopics(timeout: Duration = Duration.Infinity): BlockingTask[Map[String, List[PartitionInfo]]] =
     consumer.withConsumer(_.listTopics(timeout.asJava).asScala.mapValues(_.asScala.toList).toMap)
 

src/main/scala/zio/kafka/client/Runloop.scala

@@ -46,7 +46,8 @@ object Runloop {
     partitions: Queue[Take[Throwable, (TopicPartition, ZStreamChunk[Any, Throwable, ByteArrayCommittableRecord])]],
     rebalancingRef: Ref[Boolean],
     rebalanceListener: ConsumerRebalanceListener,
-    diagnostics: Diagnostics
+    diagnostics: Diagnostics,
+    shutdownRef: Ref[Boolean]
   ) {
     def commit(cmd: Command.Commit)   = commitQueue.offer(cmd).unit
     def commits                       = ZStream.fromQueue(commitQueue)
@@ -60,6 +61,14 @@ object Runloop {
       partitions.offer(Take.Value(tp -> data)).unit
 
     def emitIfEnabledDiagnostic(event: DiagnosticEvent) = diagnostics.emitIfEnabled(event)
+
+    val isShutdown = shutdownRef.get
+
+    def gracefulShutdown: UIO[Unit] =
+      for {
+        shutdown <- shutdownRef.modify((_, true))
+        _        <- partitions.offer(Take.End).when(!shutdown)
+      } yield ()
   }
   object Deps {
     def make(
@@ -100,6 +109,7 @@ object Runloop {
                        }
                      }
                      .toManaged_
+        shutdownRef <- Ref.make(false).toManaged_
       } yield Deps(
         consumer,
         pollFrequency,
@@ -109,7 +119,8 @@ object Runloop {
         partitions,
         rebalancingRef,
         listener,
-        diagnostics
+        diagnostics,
+        shutdownRef
       )
   }
 
@@ -384,17 +395,27 @@ object Runloop {
                    else doCommit(List(cmd)).as(state)
       } yield newState
 
+    def handleShutdown(state: State, cmd: Command): BlockingTask[State] = cmd match {
+      case Command.Poll()              => UIO.succeed(state)
+      case req @ Command.Request(_, _) => req.cont.fail(None).as(state)
+      case cmd @ Command.Commit(_, _)  => handleCommit(state, cmd)
+    }
+
     ZStream
       .mergeAll(3, 32)(
         deps.polls,
         deps.requests,
         deps.commits
       )
       .foldM(State.initial) { (state, cmd) =>
-        cmd match {
-          case Command.Poll()              => handlePoll(state)
-          case req @ Command.Request(_, _) => handleRequest(state, req)
-          case cmd @ Command.Commit(_, _)  => handleCommit(state, cmd)
+        deps.isShutdown.flatMap { shutdown =>
+          if (shutdown) handleShutdown(state, cmd)
+          else
+            cmd match {
+              case Command.Poll()              => handlePoll(state)
+              case req @ Command.Request(_, _) => handleRequest(state, req)
+              case cmd @ Command.Commit(_, _)  => handleCommit(state, cmd)
+            }
         }
       }
       .onError { cause =>

src/test/scala/zio/kafka/client/ConsumerTest.scala

@@ -232,5 +232,24 @@ class ConsumerTest extends WordSpecLike with Matchers with LazyLogging with Defa
         } yield consumeResult.fold(_ => succeed, _ => fail("Expected consumeWith to fail"))
       }
     }
+
+    "shutting down gracefully" should {
+      "not receive messages after shutting down" in runWithConsumer("group150", "client150") { consumer =>
+        for {
+          kvs <- ZIO((1 to 5).toList.map(i => (s"key$i", s"msg$i")))
+          _   <- produceMany("topic150", kvs)
+          _   <- consumer.gracefulShutdown
+          records <- consumer
+                      .subscribeAnd(Subscription.Topics(Set("topic150")))
+                      .plainStream(Serde.string, Serde.string)
+                      .flattenChunks
+                      .take(5)
+                      .runCollect
+          _ <- ZIO.effectTotal(records.map { r =>
+                (r.record.key, r.record.value)
+              } shouldBe empty)
+        } yield ()
+      }
+    }
   }
 }