Graceful shutdown of a stream for a single subscription #1201

svroonland · 2024-03-24T20:45:01Z

Implements functionality for gracefully stopping a stream for a single subscription: stop fetching records for the assigned topic-partitions but keep being subscribed so that offsets can still be committed. Intended to replace stopConsumption, which did not support multiple-subscription use cases.

A new command EndStreamsBySubscription is introduced, which calls the end method on the PartitionStreamControl of streams matching a subscription. In the method Consumer#runWithGracefulShutdown we then wait for the user's stream to complete, before removing the subscription.

This is experimental functionality, intended to replace stopConsumption at some point. Methods with this new functionality are offered besides existing methods to maintain compatibility.

~~All the fiber and scope trickery proved to be very hard to get right (the lifetime of this PR is a testimony to that), and there may still be subtle issues here.~~ This is now traced back to issue zio/zio#9288

Implements some of #941.

erikvanoosten

I didn't look at the implementation yet, only docs and tests.

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

zio-kafka-test/src/test/scala/zio/kafka/consumer/ConsumerSpec.scala

erikvanoosten

Still need more time to digest this.

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

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

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

svroonland · 2024-04-03T08:26:00Z

Hmm, should we instead of this:

Consumer.runWithGracefulShutdown(Consumer.partitionedStreamWithControl(Subscription.topics("topic150"), Serde.string, Serde.string)) { 
  stream => ... 
}

offer this:

Consumer.partitionedStreamWithGracefulShutdown(Subscription.topics("topic150"), Serde.string, Serde.string) {
  (stream, _) => stream.flatMapPar(...) 
}

The second parameter would be the SubscriptionStreamControl, which you could always manually call stop on. Or would that prevent certain use cases.. 🤔

erikvanoosten · 2024-04-03T08:55:26Z

Hmm, should we instead of this:

If I understand it correctly, the proposal allows for more use cases; with it you can also call stop for any condition you want. Is it true that after stopping, you can start consuming again?

svroonland · 2024-04-03T09:03:14Z

Well, I mean compared to just the partitionedStreamWithControl method. In both cases you would need to do something with the stream that ultimately reduces to a ZIO of Any, so I don't think the partitionedStreamWithGracefulShutdown is limiting in that regard.

stop currently doesn't support that, since the stream would then be finished. We could probably build pause and resume like in #941.

erikvanoosten · 2024-04-03T09:11:20Z

If resume after stop is not supported (and never will be), then I like the first proposal better where you don't need to call stop. What would you do after calling stop?

svroonland · 2024-04-03T09:40:37Z

Well, in both proposals you can call stop.

I don't think you want to do anything after stop, but it would give you more explicit control when to stop, instead of when the scope ends.

We probably need to decide if we want to add pause/resume in the future. If we do, we should add the control parameter like in the partitionedStreamWithGracefulShutdown example for future compatibility. If we don't, we can drop it altogether and make SubscriptionStreamControl a purely internal concept (if at all).

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

guizmaii · 2024-04-05T03:30:45Z

Hey :)

Thanks for the great work!

Here's some initial feedback:

I'm not a big fan of the SubscriptionStreamControl implementation.

To me, functions/methods returning it should return a Tuple (stream, control).
It avoids adding one more concept for our users to understand and learn (Kafka already has a lot of concepts)
It also simplifies the interface of the control type, the current one with the [S <: ZStream[_, _, _]] being complex
It also simplifies the return type of our functions/methods, avoiding this kind of type:

SubscriptionStreamControl[Stream[Throwable, Chunk[(TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])]]]

in favor of:

(Stream[Throwable, Chunk[(TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])], SubscriptionStreamControl)

Made the change in a PR to show/study how, to me, it simplifies things: https://github.com/zio/zio-kafka/pull/1207/files

zio-kafka/src/main/scala/zio/kafka/consumer/internal/RunloopAccess.scala

guizmaii · 2024-04-05T04:23:04Z

Didn't finish my review yet. I still have some parts of the code to explore/understand, but I have to go. I'll finish it later 🙂

svroonland · 2024-04-05T06:35:18Z

Thanks for the feedback Jules. Agreed about the extra concept that would be unwanted. Check out my latest interface proposal where there is only a plainStreamWithGracefulShutdown method and SubscriptionStreamControl remains hidden.

erikvanoosten

Still reading the code...

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

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

zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala

erikvanoosten · 2024-04-07T08:17:14Z

I understand now that when graceful shutdown starts we're ending the subscribed streams. That should work nicely. Lets work out what will happen next to the runloop. The runloop would still be happily fetching records for that stream. When those are offered to the stream, PartitionStreamControl.offerRecords will probably append those records to the queue (even though it now also contains an 'end' token). Because of the 'end' token that is already in that queue, these new records will never be taken out. Back pressure will kick in (depending on the fetch strategy) and the partitions will be paused. Once we're unsubscribed, 15 seconds later, the queue will be garbage collected. So far so good.

We can do slightly better though. We're fetching and storing all these records in the queue for nothing, even potentially causing an OOM for systems that are tuned for the case where processing happens almost immediately.

My proposal is to:

stop accepting more records in PartitionStreamControl.offerRecords when the queue was ended
in Runloop.handlePoll only pass running streams to fetchStrategy.selectPartitionsToFetch so that partitions for ended streams are immediately paused

If you want, I can extend this PR with that proposal (or create a separate PR).

svroonland · 2024-04-07T11:02:54Z

@erikvanoosten If you have some time to implement those two things, by all means.

erikvanoosten · 2024-04-13T15:55:34Z

@erikvanoosten If you have some time to implement those two things, by all means.

@svroonland Done in commit 1218204.

Now I am wondering, how can we test this?

svroonland · 2024-04-14T15:45:54Z

Change looks good. Totally forgot to implement this part.

…for comprehension

svroonland · 2024-11-10T07:16:52Z

Was able to create a minimized reproducer of the issue: zio/zio#9288

svroonland · 2024-11-16T09:24:27Z

The abovementioned issue has been fixed and will probably be in the next ZIO release. For compatibility with older versions let's keep the .fork.flatMap(_.join) workaround.

…-control

erikvanoosten · 2024-11-16T12:17:37Z

The abovementioned issue has been fixed and will probably be in the next ZIO release. For compatibility with older versions let's keep the .fork.flatMap(_.join) workaround.

Great work!

Shall we slap a comment on it? E.g. something like:

// Workaround for a bug in ZIO up to v2.1.12. See https://github.com/zio/zio/issues/9288

…-control

This reverts commit f899758.

erikvanoosten

So many little comments. I'll push this now so you have something to work with. Meanwhile I'll continue reviewing.

erikvanoosten · 2025-03-09T10:46:54Z

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

+
+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 `with*Stream` variants of `plainStream`, `partitionedStream` and `partitionedAssignmentStream` for this purpose. These methods accept a parameter that describes the execution of a stream, which is gracefully ended when the method is interrupted.


We can extend this a bit more before we merge this PR.

zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala

erikvanoosten · 2025-03-09T10:50:00Z

zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala

+      case RunloopCommand.EndStreamsBySubscription(subscription, cont) =>
+        ZIO.foreachDiscard(
+          state.assignedStreams.filter(stream => Subscription.subscriptionMatches(subscription, stream.tp))
+        )(_.end) *> cont


The formatter makes some weird choices here. Would this case be more readable when written as a for comprehension?

erikvanoosten · 2025-03-09T10:55:23Z

zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala

@@ -203,17 +213,20 @@ private[consumer] final class Runloop private (

  private def handlePoll(state: State): Task[State] = {
    for {
-      partitionsToFetch  <- settings.fetchStrategy.selectPartitionsToFetch(state.assignedStreams)
-      pendingCommitCount <- committer.pendingCommitCount
+      runningStreamsBeforePoll <- ZIO.filter(state.assignedStreams)(_.isRunning)


The isRunning is based on whether the stream control ended or not. However, we can/could do better by not providing data for ended subscriptions instead. WDYT?

erikvanoosten · 2025-03-09T12:08:10Z

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

+/**
+ * Allows graceful shutdown of a stream, where no more records are being fetched but the in-flight records can continue
+ * to be processed and their offsets committed.
+ *
+ * @param stream
+ *   The stream of partitions / records for this subscription
+ * @param stop
+ *   Stop fetching records for the subscribed topic-partitions and end the associated streams, while allowing commits to
+ *   proceed (consumer remains subscribed)
+ */
+final private[consumer] case class SubscriptionStreamControl[S <: ZStream[_, _, _]](stream: S, stop: UIO[Unit])


This case class models a stream that can be stopped from outside (with a stop method). It is not tied to subscriptions or anything else from zio-kafka. Therefore, I propose we rename this to something like StoppableStream, or even simpler: StreamControl.

Suggested change

/**

* Allows graceful shutdown of a stream, where no more records are being fetched but the in-flight records can continue

* to be processed and their offsets committed.

*

* @param stream

* The stream of partitions / records for this subscription

* @param stop

* Stop fetching records for the subscribed topic-partitions and end the associated streams, while allowing commits to

* proceed (consumer remains subscribed)

*/

final private[consumer] case class SubscriptionStreamControl[S <: ZStream[_, _, _]](stream: S, stop: UIO[Unit])

/**

* Models a stream with a graceful shutdown.

*

* In a graceful shutdown, the stream stops pulling elements from its source, but completes processing of already pulled stream elements.

*

* @param stream

* A stream that supports graceful shutdown

* @param stop

* Initiate a graceful shutdown of the stream

*/

final private[consumer] case class StreamControl[S <: ZStream[_, _, _]](stream: S, stop: UIO[Unit])

I am wondering why this it is not like this:

case class StreamControl[R, E, A]](stream: ZStream[R, E, A], stop: UIO[Unit])

There is an opportunity to clean up the code that is using this with methods like map:

case class StreamControl[R, E, A]](stream: ZStream[R, E, A], stop: UIO[Unit]) { def map[R1 <: R, E1 >: E, B](f: ZStream[R, E, A] => ZStream[R1, E1, B]): StreamControl[R1, E1, B] = StreamControl(f(stream), stop) }

Nice, I like it

erikvanoosten · 2025-03-09T12:17:43Z

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

+   * Like [[partitionedAssignmentStream]] but wraps the stream in a construct that ensures graceful shutdown.
+   *
+   * When this effect is interrupted, all partition streams are closed upstream, allowing the stream created by
+   * `withStream` to complete gracefully all stream stages, thereby fully processing all buffered and/or in-flight
+   * messages.
+   *
+   * EXPERIMENTAL API


Suggested change

* Like [[partitionedAssignmentStream]] but wraps the stream in a construct that ensures graceful shutdown.

*

* When this effect is interrupted, all partition streams are closed upstream, allowing the stream created by

* `withStream` to complete gracefully all stream stages, thereby fully processing all buffered and/or in-flight

* messages.

*

* EXPERIMENTAL API

* Like [[partitionedAssignmentStream]] but wraps the stream in an effect that allows graceful shutdown.

*

* When this effect is interrupted, the stream of assigned partitions ends, allowing the streams created by

* `withStream` to complete gracefully, thereby fully processing all buffered and/or in-flight

* stream elements.

*

* WARNING: this is an EXPERIMENTAL API and may disappear or change in an incompatible way without notice in any zio-kafka version.

erikvanoosten · 2025-03-09T12:59:56Z

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

  override def plainStream[R, K, V](
    subscription: Subscription,
    keyDeserializer: Deserializer[R, K],
    valueDeserializer: Deserializer[R, V],
-    bufferSize: Int
+    bufferSize: Int = 4


Shall we extract this change to its own PR? This looks like something that is useful on its own.

Note to self: it may also affect documentation.

erikvanoosten · 2025-03-09T13:14:29Z

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

+                .timeout(shutdownTimeout)
+                .someOrElseZIO(
+                  ZIO.logError(
+                    "Timeout joining withStream fiber in runWithGracefulShutdown. Not all pending commits may have been processed."


Here is my attempt at rewriting the message more towards the user's view. Not sure if this is entirely correct though.

Suggested change

"Timeout joining withStream fiber in runWithGracefulShutdown. Not all pending commits may have been processed."

"Timeout waiting for `withStream` to gracefully shut down. Not all in-flight records may have been processed."

erikvanoosten · 2025-03-09T13:15:23Z

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

+                  )
+                )
+                .tapErrorCause(cause =>
+                  ZIO.logErrorCause("Error joining withStream fiber in runWithGracefulShutdown", cause)


If I understand this correctly, the problem is not that the join fails, but more likely the stream itself.

Suggested change

ZIO.logErrorCause("Error joining withStream fiber in runWithGracefulShutdown", cause)

ZIO.logErrorCause("Stream failed while awaiting its graceful shutdown", cause)

erikvanoosten · 2025-03-09T13:17:09Z

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

+                // The fork and join is a workaround for https://github.com/zio/zio/issues/9288 for ZIO <= 2.1.12
+                .forkDaemon
+                .flatMap(_.join)
+                .tapErrorCause(cause =>


Question: when we remove this workaround (I am actually in favor of that, we can document that we require zio 2.1.13+), do we then still need this second tapErrorCause?

…-control

erikvanoosten · 2025-03-09T13:42:12Z

zio-kafka/src/main/scala/zio/kafka/consumer/internal/RunloopAccess.scala

               diagnostics.emit(Finalization.SubscriptionFinalized)
           }
-    } yield stream
+    } yield SubscriptionStreamControl(
+      stream = stream.merge(ZStream.fromZIO(end.await).as(Take.end)),


When I read (line 72) this some alarm bells went of in my head.

PartitionStreamControl registers which records (offsets) were pulled. This information is later used by the rebalance listener when rebalanceSafeCommit is enabled.

When the 'end' gets merged into the stream provided by PartitionStreamControl, it could be that the latter just pulled a chunk of records, but before the chunk was given downstream, the end is merged in. This causes downstream to never see those records, even though PartitionStreamControl thinks those have been pulled. When this happens, the ending rebalance will wait for maxRebalanceDuration.

A solution could be to stop offering data from the runloop (see also https://github.com/zio/zio-kafka/pull/1201/files#r1986290475).

svroonland changed the title ~~Subscription stream control~~ Graceful shutdown of a single subscription Mar 30, 2024

svroonland marked this pull request as ready for review March 30, 2024 11:07

svroonland requested a review from erikvanoosten as a code owner March 30, 2024 11:07

erikvanoosten reviewed Mar 30, 2024

View reviewed changes

svroonland mentioned this pull request Apr 4, 2024

Slow (timeout) shutdown and rebalancing with withRebalanceSafeCommits #1132

Closed

guizmaii reviewed Apr 5, 2024

View reviewed changes

zio-kafka/src/main/scala/zio/kafka/consumer/Consumer.scala Outdated Show resolved Hide resolved

guizmaii reviewed Apr 5, 2024

View reviewed changes

zio-kafka/src/main/scala/zio/kafka/consumer/internal/RunloopAccess.scala Outdated Show resolved Hide resolved

guizmaii reviewed Apr 5, 2024

View reviewed changes

zio-kafka/src/main/scala/zio/kafka/consumer/internal/RunloopAccess.scala Outdated Show resolved Hide resolved

erikvanoosten reviewed Apr 7, 2024

View reviewed changes

zio-kafka/src/main/scala/zio/kafka/consumer/internal/Runloop.scala Outdated Show resolved Hide resolved

erikvanoosten mentioned this pull request Apr 10, 2024

[Idea] pause/resume/stop Stream and remove stopConsumption #941

Open

svroonland added 6 commits April 27, 2024 09:42

Draft of interface changes

7a9a6f8

Remove deprecated for now

c8196b3

Draft implementation

583d112

Fix implementation + first test

f55c898

More tests copied from stopConsumption

6f86958

Alternative interface, workaround inability to deconstruct tuples in …

9eb989b

…for comprehension

Fix test

5ef97ef

Update comment

bec0e25

svroonland added 2 commits November 16, 2024 10:32

Merge remote-tracking branch 'origin/master' into subscription-stream…

c7b6879

…-control

Comment workaround

5e2031e

svroonland added this to the 3.0.0 milestone Nov 16, 2024

svroonland and others added 18 commits November 16, 2024 14:50

Merge branch 'master' into subscription-stream-control

4c02d08

Merge remote-tracking branch 'origin/master' into subscription-stream…

e4f4e11

…-control

Tweak

46f2142

Merge remote-tracking branch 'origin/master' into subscription-stream…

43a0c8a

…-control

Fix merge error + test

a346981

Add experimental notes

585ee42

Allow returning a value from with*Stream methods

f899758

More logging

07c0bf3

Merge remote-tracking branch 'origin/master' into subscription-stream…

00d868d

…-control

Make debugging easier

4e822fa

No infinite retry

af901eb

Restore log level

298c293

Use KafkaTestUtils for creating topic

f9f31ed

Unused imports

d87acd1

Remove usage of stopConsumption

efe709a

Rearrange all stopConsumption tests in one suite

83c48fd

Revert "Allow returning a value from with*Stream methods"

f1610f9

This reverts commit f899758.

Small refactoring

be05e25

erikvanoosten reviewed Mar 9, 2025

View reviewed changes

svroonland added 2 commits March 9, 2025 14:20

Merge remote-tracking branch 'origin/master' into subscription-stream…

09a6375

…-control

Alternative type parameters for StreamControl

adf3965

erikvanoosten reviewed Mar 9, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Graceful shutdown of a stream for a single subscription #1201

Graceful shutdown of a stream for a single subscription #1201

svroonland commented Mar 24, 2024 •

edited

Loading

erikvanoosten left a comment

erikvanoosten left a comment

svroonland commented Apr 3, 2024 •

edited

Loading

erikvanoosten commented Apr 3, 2024

svroonland commented Apr 3, 2024

erikvanoosten commented Apr 3, 2024

svroonland commented Apr 3, 2024

guizmaii commented Apr 5, 2024 •

edited

Loading

guizmaii commented Apr 5, 2024

svroonland commented Apr 5, 2024

erikvanoosten left a comment

erikvanoosten commented Apr 7, 2024 •

edited

Loading

svroonland commented Apr 7, 2024

erikvanoosten commented Apr 13, 2024 •

edited

Loading

svroonland commented Apr 14, 2024 •

edited

Loading

svroonland commented Nov 10, 2024

svroonland commented Nov 16, 2024

erikvanoosten commented Nov 16, 2024

erikvanoosten left a comment

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

svroonland Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025

erikvanoosten Mar 9, 2025


		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 `with*Stream` variants of `plainStream`, `partitionedStream` and `partitionedAssignmentStream` for this purpose. These methods accept a parameter that describes the execution of a stream, which is gracefully ended when the method is interrupted.

	"Timeout joining withStream fiber in runWithGracefulShutdown. Not all pending commits may have been processed."
	"Timeout waiting for `withStream` to gracefully shut down. Not all in-flight records may have been processed."

	ZIO.logErrorCause("Error joining withStream fiber in runWithGracefulShutdown", cause)
	ZIO.logErrorCause("Stream failed while awaiting its graceful shutdown", cause)

Graceful shutdown of a stream for a single subscription #1201

Are you sure you want to change the base?

Graceful shutdown of a stream for a single subscription #1201

Conversation

svroonland commented Mar 24, 2024 • edited Loading

erikvanoosten left a comment

Choose a reason for hiding this comment

erikvanoosten left a comment

Choose a reason for hiding this comment

svroonland commented Apr 3, 2024 • edited Loading

erikvanoosten commented Apr 3, 2024

svroonland commented Apr 3, 2024

erikvanoosten commented Apr 3, 2024

svroonland commented Apr 3, 2024

guizmaii commented Apr 5, 2024 • edited Loading

guizmaii commented Apr 5, 2024

svroonland commented Apr 5, 2024

erikvanoosten left a comment

Choose a reason for hiding this comment

erikvanoosten commented Apr 7, 2024 • edited Loading

svroonland commented Apr 7, 2024

erikvanoosten commented Apr 13, 2024 • edited Loading

svroonland commented Apr 14, 2024 • edited Loading

svroonland commented Nov 10, 2024

svroonland commented Nov 16, 2024

erikvanoosten commented Nov 16, 2024

erikvanoosten left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

svroonland commented Mar 24, 2024 •

edited

Loading

svroonland commented Apr 3, 2024 •

edited

Loading

guizmaii commented Apr 5, 2024 •

edited

Loading

erikvanoosten commented Apr 7, 2024 •

edited

Loading

erikvanoosten commented Apr 13, 2024 •

edited

Loading

svroonland commented Apr 14, 2024 •

edited

Loading