This is an automated email from the ASF dual-hosted git repository.
pjfanning pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/pekko-connectors-kafka.git
The following commit(s) were added to refs/heads/main by this push:
new 80170dff Clear lastRevoked after use to avoid wiping buffers under
cooperative rebalancing (#617)
80170dff is described below
commit 80170dfff3ab0ffc237b855c5fe7cdaf0bbde6d8
Author: Arvind Raghavan <[email protected]>
AuthorDate: Sun Aug 9 07:00:15 2026 -0500
Clear lastRevoked after use to avoid wiping buffers under cooperative
rebalancing (#617)
* Clear lastRevoked after use to avoid wiping buffers under cooperative
rebalancing
* With the cooperative rebalance protocol onPartitionsRevoked is only
invoked on
members that actually revoke partitions, while onPartitionsAssigned is
invoked
on every member in every rebalance. A stale lastRevoked set from an
earlier
rebalance was re-applied on such assignment callbacks and filtered
buffered
records of partitions that remained assigned, losing records without them
being processed or committed.
* Reset lastRevoked after the flush in SingleSourceLogic and SubSourceLogic.
* Run the RebalanceSpec buffer tests against both the eager and cooperative
protocols using a cooperative variant of the test assignor.
* Add CooperativeRebalanceSpec reproducing the record loss for plain and
partitioned sources; both tests fail without the fix.
Fixes #616
Co-Authored-By: Claude Fable 5 <[email protected]>
* Address review: standard ASF header for new file, drop final vals and
test logging, apply scalafmt
Co-Authored-By: Claude Fable 5 <[email protected]>
---------
Co-authored-by: Claude Fable 5 <[email protected]>
---
.../pekko/kafka/internal/SingleSourceLogic.scala | 4 +-
.../pekko/kafka/internal/SubSourceLogic.scala | 4 +-
.../kafka/scaladsl/CooperativeRebalanceSpec.scala | 174 ++++++++
.../pekko/kafka/scaladsl/RebalanceSpec.scala | 439 ++++++++++++---------
4 files changed, 439 insertions(+), 182 deletions(-)
diff --git
a/core/src/main/scala/org/apache/pekko/kafka/internal/SingleSourceLogic.scala
b/core/src/main/scala/org/apache/pekko/kafka/internal/SingleSourceLogic.scala
index 283a0c4c..351848f6 100644
---
a/core/src/main/scala/org/apache/pekko/kafka/internal/SingleSourceLogic.scala
+++
b/core/src/main/scala/org/apache/pekko/kafka/internal/SingleSourceLogic.scala
@@ -95,8 +95,10 @@ import scala.concurrent.{ Future, Promise }
override def onRevoke(revokedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
lastRevoked = revokedTps
- override def onAssign(assignedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
+ override def onAssign(assignedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit = {
filterRevokedPartitionsCB.invoke(lastRevoked -- assignedTps)
+ lastRevoked = Set.empty
+ }
override def onLost(lostTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
filterRevokedPartitionsCB.invoke(lostTps)
diff --git
a/core/src/main/scala/org/apache/pekko/kafka/internal/SubSourceLogic.scala
b/core/src/main/scala/org/apache/pekko/kafka/internal/SubSourceLogic.scala
index c09acb31..ec70d534 100644
--- a/core/src/main/scala/org/apache/pekko/kafka/internal/SubSourceLogic.scala
+++ b/core/src/main/scala/org/apache/pekko/kafka/internal/SubSourceLogic.scala
@@ -298,11 +298,13 @@ private class SubSourceLogic[K, V, Msg](
override def onRevoke(revokedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
lastRevoked = revokedTps
- override def onAssign(assignedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
+ override def onAssign(assignedTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit = {
for {
tp <- lastRevoked -- assignedTps
control <- subSources.get(tp)
} control.filterRevokedPartitionsCB.invoke(Set(tp))
+ lastRevoked = Set.empty
+ }
override def onLost(lostTps: Set[TopicPartition], consumer:
RestrictedConsumer): Unit =
for {
diff --git
a/tests/src/test/scala/org/apache/pekko/kafka/scaladsl/CooperativeRebalanceSpec.scala
b/tests/src/test/scala/org/apache/pekko/kafka/scaladsl/CooperativeRebalanceSpec.scala
new file mode 100644
index 00000000..482e5122
--- /dev/null
+++
b/tests/src/test/scala/org/apache/pekko/kafka/scaladsl/CooperativeRebalanceSpec.scala
@@ -0,0 +1,174 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.pekko.kafka.scaladsl
+
+import org.apache.pekko
+import pekko.Done
+import pekko.kafka._
+import pekko.kafka.scaladsl.Consumer.Control
+import pekko.kafka.testkit.scaladsl.TestcontainersKafkaLike
+import pekko.stream.scaladsl.{ Keep, Source }
+import pekko.stream.testkit.scaladsl.StreamTestKit.assertAllStagesStopped
+import pekko.stream.testkit.scaladsl.TestSink
+import pekko.testkit.TestProbe
+import org.apache.kafka.clients.consumer.{ ConsumerConfig, ConsumerRecord }
+import org.apache.kafka.common.TopicPartition
+import org.scalatest.Inside
+
+import scala.concurrent.duration._
+import scala.util.Random
+
+/**
+ * With the cooperative rebalance protocol `onPartitionsRevoked` is only
invoked on members that
+ * actually revoke partitions, while `onPartitionsAssigned` is invoked on
every member in every
+ * rebalance, possibly with an empty set. These tests reproduce rebalances
where the revoke
+ * callback is skipped and assert that buffered records of partitions that
remain assigned are
+ * still delivered (at-least-once).
+ *
+ * Reproduces the scenario of
https://github.com/apache/pekko-connectors-kafka/issues/616:
+ * 1. a partition is revoked from consumer 1 (its `lastRevoked` state becomes
non-empty),
+ * 2. the partition is later re-assigned to consumer 1,
+ * 3. an unrelated member joins, completing a rebalance in which consumer 1
revokes nothing
+ * (`onPartitionsRevoked` is not invoked) and gains nothing
(`onPartitionsAssigned` with an
+ * empty set) while records for the re-assigned partition sit in the
source stage buffer.
+ *
+ * The test runs against `plainSource` and `plainPartitionedSource` (flattened
with
+ * `flatMapMerge` so both shapes expose a single stream of records) as the
underlying stages
+ * implement the revoked-buffer bookkeeping independently. The merge stage of
the flattened
+ * variant buffers a few records of its own that survive a revoke as
duplicates, so the tests
+ * assert complete delivery rather than an exact sequence.
+ */
+class CooperativeRebalanceSpec extends SpecBase with TestcontainersKafkaLike
with Inside {
+
+ implicit override val patienceConfig: PatienceConfig =
PatienceConfig(30.seconds, 500.millis)
+
+ val partition1 = 1
+ val consumerClientId1 = "consumer-1"
+ val consumerClientId2 = "consumer-2"
+ val consumerClientId3 = "consumer-3"
+
+ private def cooperativeSettings(group: String) =
+ consumerDefaults
+ .withProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "500")
+ .withProperty(
+ ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
+ classOf[CooperativePekkoConnectorsAssignor].getName)
+ .withGroupId(group)
+
+ private def awaitAssigned(rebalanceActor: TestProbe,
+ subscription: AutoSubscription,
+ tps: Set[TopicPartition]): Unit =
+ rebalanceActor.fishForMessage(10.seconds) {
+ case TopicPartitionsAssigned(`subscription`, assigned) if assigned ==
tps => true
+ case TopicPartitionsAssigned(`subscription`, assigned) if
assigned.isEmpty => false
+ }
+
+ sealed trait SourceCase {
+ def label: String
+ def source(settings: ConsumerSettings[String, String],
+ subscription: AutoSubscription): Source[ConsumerRecord[String,
String], Control]
+ }
+
+ case object PlainCase extends SourceCase {
+ override val label = "plain source"
+ override def source(settings: ConsumerSettings[String, String],
+ subscription: AutoSubscription): Source[ConsumerRecord[String,
String], Control] =
+ Consumer.plainSource(settings, subscription)
+ }
+
+ case object PartitionedCase extends SourceCase {
+ override val label = "partitioned source"
+ override def source(settings: ConsumerSettings[String, String],
+ subscription: AutoSubscription): Source[ConsumerRecord[String,
String], Control] =
+ Consumer
+ .plainPartitionedSource(settings, subscription)
+ .flatMapMerge(breadth = 8, { case (_, records) => records })
+ }
+
+ "Buffered records of partitions that stay assigned" must {
+
+ List(PlainCase, PartitionedCase).foreach { mode =>
+ s"be delivered after a rebalance without revocation (${mode.label})" in
assertAllStagesStopped {
+ val count = 100L
+ val topicSuffix = Random.nextInt()
+ val topic1 = createTopic(topicSuffix, partitions = 2)
+ val group1 = createGroupId(1)
+ val tp0 = new TopicPartition(topic1, partition0)
+ val tp1 = new TopicPartition(topic1, partition1)
+ val consumerSettings = cooperativeSettings(group1)
+
+ def joinConsumer(clientId: String) = {
+ val rebalanceActor = TestProbe()
+ val subscription =
Subscriptions.topics(topic1).withRebalanceListener(rebalanceActor.ref)
+ val (control, probe) = Consumer
+ .plainSource(consumerSettings.withClientId(clientId), subscription)
+ .toMat(TestSink())(Keep.both)
+ .run()
+ (control, probe, rebalanceActor, subscription)
+ }
+
+ awaitProduce(produce(topic1, 0 to count.toInt, partition1))
+
+
PekkoConnectorsAssignor.clientIdToPartitionMap.set(Map(consumerClientId1 ->
Set(tp0, tp1)))
+
+ val probe1rebalanceActor = TestProbe()
+ val probe1subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe1rebalanceActor.ref)
+ val (control1, probe1) = mode
+ .source(consumerSettings.withClientId(consumerClientId1),
probe1subscription)
+ .toMat(TestSink())(Keep.both)
+ .run()
+
+
probe1rebalanceActor.expectMsg(TopicPartitionsAssigned(probe1subscription,
Set(tp0, tp1)))
+ probe1.requestNext()
+
+ PekkoConnectorsAssignor.clientIdToPartitionMap.set(
+ Map(consumerClientId1 -> Set(tp0), consumerClientId2 -> Set(tp1)))
+ val (control2, probe2, probe2rebalanceActor, probe2subscription) =
joinConsumer(consumerClientId2)
+
+
probe1rebalanceActor.expectMsg(TopicPartitionsRevoked(probe1subscription,
Set(tp1)))
+ awaitAssigned(probe2rebalanceActor, probe2subscription, Set(tp1))
+
+
PekkoConnectorsAssignor.clientIdToPartitionMap.set(Map(consumerClientId1 ->
Set(tp0, tp1)))
+ probe2.cancel()
+ control2.isShutdown.futureValue shouldBe Done
+ awaitAssigned(probe1rebalanceActor, probe1subscription, Set(tp1))
+
+ probe1.requestNext()
+
+ PekkoConnectorsAssignor.clientIdToPartitionMap.set(
+ Map(consumerClientId1 -> Set(tp0, tp1), consumerClientId3 ->
Set.empty[TopicPartition]))
+ val (control3, probe3, probe3rebalanceActor, probe3subscription) =
joinConsumer(consumerClientId3)
+
+
probe1rebalanceActor.expectMsg(TopicPartitionsAssigned(probe1subscription,
Set.empty))
+
probe3rebalanceActor.expectMsg(TopicPartitionsAssigned(probe3subscription,
Set.empty))
+
+ // give the asynchronous buffer filter of the rebalance a chance to
apply before demanding
+ probe1.expectNoMessage(500.millis)
+
+ probe1.request(count * 3)
+ val values = probe1.receiveWithin(5.seconds).map(_.value)
+ values should contain allElementsOf (1 to count.toInt).map(_.toString)
+
+ probe1.cancel()
+ probe3.cancel()
+ control1.isShutdown.futureValue shouldBe Done
+ control3.isShutdown.futureValue shouldBe Done
+ }
+ }
+ }
+}
diff --git
a/tests/src/test/scala/org/apache/pekko/kafka/scaladsl/RebalanceSpec.scala
b/tests/src/test/scala/org/apache/pekko/kafka/scaladsl/RebalanceSpec.scala
index 21fca2b2..c3e476f4 100644
--- a/tests/src/test/scala/org/apache/pekko/kafka/scaladsl/RebalanceSpec.scala
+++ b/tests/src/test/scala/org/apache/pekko/kafka/scaladsl/RebalanceSpec.scala
@@ -45,197 +45,241 @@ class RebalanceSpec extends SpecBase with
TestcontainersKafkaLike with Inside {
final val consumerClientId1 = "consumer-1"
final val consumerClientId2 = "consumer-2"
+ sealed trait ProtocolCase {
+ def label: String
+ def assignor: Class[?]
+
+ /** Expected rebalance listener events after the second consumer joins and
takes tp1. */
+ def expectSecondConsumerJoin(
+ probe1rebalanceActor: TestProbe,
+ probe1subscription: AutoSubscription,
+ probe2rebalanceActor: TestProbe,
+ probe2subscription: AutoSubscription,
+ tp0: TopicPartition,
+ tp1: TopicPartition): Unit
+ }
+
+ case object EagerCase extends ProtocolCase {
+ override val label = "eager"
+ override val assignor: Class[?] = classOf[PekkoConnectorsAssignor]
+
+ override def expectSecondConsumerJoin(
+ probe1rebalanceActor: TestProbe,
+ probe1subscription: AutoSubscription,
+ probe2rebalanceActor: TestProbe,
+ probe2subscription: AutoSubscription,
+ tp0: TopicPartition,
+ tp1: TopicPartition): Unit = {
+
probe1rebalanceActor.expectMsg(TopicPartitionsRevoked(probe1subscription,
Set(tp0, tp1)))
+
probe1rebalanceActor.expectMsg(TopicPartitionsAssigned(probe1subscription,
Set(tp0)))
+
probe2rebalanceActor.expectMsg(TopicPartitionsAssigned(probe2subscription,
Set(tp1)))
+ }
+ }
+
+ case object CooperativeCase extends ProtocolCase {
+ override val label = "cooperative"
+ override val assignor: Class[?] =
classOf[CooperativePekkoConnectorsAssignor]
+
+ override def expectSecondConsumerJoin(
+ probe1rebalanceActor: TestProbe,
+ probe1subscription: AutoSubscription,
+ probe2rebalanceActor: TestProbe,
+ probe2subscription: AutoSubscription,
+ tp0: TopicPartition,
+ tp1: TopicPartition): Unit = {
+ // tp1 is revoked from consumer 1 and assigned to consumer 2 in a
follow-up rebalance.
+ // Intermediate empty assignments are not asserted: members that fail
the intermediate
+ // generation's sync with REBALANCE_IN_PROGRESS skip its assignment
callback entirely.
+
probe1rebalanceActor.expectMsg(TopicPartitionsRevoked(probe1subscription,
Set(tp1)))
+ probe2rebalanceActor.fishForMessage(10.seconds) {
+ case TopicPartitionsAssigned(`probe2subscription`, assigned) if
assigned == Set(tp1) => true
+ case TopicPartitionsAssigned(`probe2subscription`, assigned) if
assigned.isEmpty => false
+ }
+ }
+ }
+
"Fetched records" must {
// The `max.poll.records` controls how many records Kafka fetches
internally during a poll.
// issue explained in https://github.com/akka/alpakka-kafka/issues/872
// this test added with https://github.com/akka/alpakka-kafka/pull/865
- "be removed from the source stage buffer when a partition is revoked" in
assertAllStagesStopped {
- val count = 20L
- // de-coupling consecutive test runs with crossScalaVersions on build
- val topicSuffix = Random.nextInt()
- val topic1 = createTopic(topicSuffix, partitions = 2)
- val group1 = createGroupId(1)
- val tp0 = new TopicPartition(topic1, partition0)
- val tp1 = new TopicPartition(topic1, partition1)
- val consumerSettings = consumerDefaults
- .withProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "500") // 500 is
the default value
- .withProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
classOf[PekkoConnectorsAssignor].getName)
- .withGroupId(group1)
-
- awaitProduce(produce(topic1, 0 to count.toInt, partition1))
-
- PekkoConnectorsAssignor.clientIdToPartitionMap.set(
- Map(
- consumerClientId1 -> Set(tp0, tp1)))
-
- log.debug("Subscribe to the topic (without demand)")
- val probe1rebalanceActor = TestProbe()
- val probe1subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe1rebalanceActor.ref)
- val (control1, probe1) = Consumer
- .plainSource(consumerSettings.withClientId(consumerClientId1),
probe1subscription)
- .toMat(TestSink())(Keep.both)
- .run()
-
- log.debug("Await initial partition assignment")
- probe1rebalanceActor.expectMsg(
- TopicPartitionsAssigned(probe1subscription,
- Set(new TopicPartition(topic1, partition0), new
TopicPartition(topic1, partition1))))
-
- log.debug("read one message from probe1 with partition 1")
- probe1.requestNext()
-
- PekkoConnectorsAssignor.clientIdToPartitionMap.set(
- Map(
- consumerClientId1 -> Set(tp0),
- consumerClientId2 -> Set(tp1)))
-
- log.debug("Subscribe to the topic (without demand)")
- val probe2rebalanceActor = TestProbe()
- val probe2subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe2rebalanceActor.ref)
- val (control2, probe2) = Consumer
- .plainSource(consumerSettings.withClientId(consumerClientId2),
probe2subscription)
- .toMat(TestSink())(Keep.both)
- .run()
-
- log.debug("Await a revoke to consumer 1")
- probe1rebalanceActor.expectMsg(
- TopicPartitionsRevoked(probe1subscription,
- Set(new TopicPartition(topic1, partition0), new
TopicPartition(topic1, partition1))))
-
- log.debug("the rebalance finishes")
- probe1rebalanceActor.expectMsg(
- TopicPartitionsAssigned(probe1subscription, Set(new
TopicPartition(topic1, partition0))))
- probe2rebalanceActor.expectMsg(
- TopicPartitionsAssigned(probe2subscription, Set(new
TopicPartition(topic1, partition1))))
-
- log.debug("resume demand on both consumers")
- probe1.request(count)
- probe2.request(count)
-
- val probe2messages = probe2.expectNextN(count)
-
- log.debug("no further messages enqueued on probe1 as partition 1 is
balanced away")
- probe1.expectNoMessage(500.millis)
-
- probe2messages should have size count
-
- probe1.cancel()
- probe2.cancel()
-
- control1.isShutdown.futureValue shouldBe Done
- control2.isShutdown.futureValue shouldBe Done
+ List(EagerCase, CooperativeCase).foreach { mode =>
+ s"be removed from the source stage buffer when a partition is revoked
(${mode.label})" in assertAllStagesStopped {
+ val count = 20L
+ // de-coupling consecutive test runs with crossScalaVersions on build
+ val topicSuffix = Random.nextInt()
+ val topic1 = createTopic(topicSuffix, partitions = 2)
+ val group1 = createGroupId(1)
+ val tp0 = new TopicPartition(topic1, partition0)
+ val tp1 = new TopicPartition(topic1, partition1)
+ val consumerSettings = consumerDefaults
+ .withProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "500") // 500
is the default value
+ .withProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
mode.assignor.getName)
+ .withGroupId(group1)
+
+ awaitProduce(produce(topic1, 0 to count.toInt, partition1))
+
+ PekkoConnectorsAssignor.clientIdToPartitionMap.set(
+ Map(
+ consumerClientId1 -> Set(tp0, tp1)))
+
+ log.debug("Subscribe to the topic (without demand)")
+ val probe1rebalanceActor = TestProbe()
+ val probe1subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe1rebalanceActor.ref)
+ val (control1, probe1) = Consumer
+ .plainSource(consumerSettings.withClientId(consumerClientId1),
probe1subscription)
+ .toMat(TestSink())(Keep.both)
+ .run()
+
+ log.debug("Await initial partition assignment")
+ probe1rebalanceActor.expectMsg(
+ TopicPartitionsAssigned(probe1subscription,
+ Set(new TopicPartition(topic1, partition0), new
TopicPartition(topic1, partition1))))
+
+ log.debug("read one message from probe1 with partition 1")
+ probe1.requestNext()
+
+ PekkoConnectorsAssignor.clientIdToPartitionMap.set(
+ Map(
+ consumerClientId1 -> Set(tp0),
+ consumerClientId2 -> Set(tp1)))
+
+ log.debug("Subscribe to the topic (without demand)")
+ val probe2rebalanceActor = TestProbe()
+ val probe2subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe2rebalanceActor.ref)
+ val (control2, probe2) = Consumer
+ .plainSource(consumerSettings.withClientId(consumerClientId2),
probe2subscription)
+ .toMat(TestSink())(Keep.both)
+ .run()
+
+ log.debug("Await the rebalance to complete")
+ mode.expectSecondConsumerJoin(probe1rebalanceActor,
probe1subscription, probe2rebalanceActor,
+ probe2subscription, tp0, tp1)
+
+ log.debug("resume demand on both consumers")
+ probe1.request(count)
+ probe2.request(count)
+
+ val probe2messages = probe2.expectNextN(count)
+
+ log.debug("no further messages enqueued on probe1 as partition 1 is
balanced away")
+ probe1.expectNoMessage(500.millis)
+
+ probe2messages should have size count
+
+ probe1.cancel()
+ probe2.cancel()
+
+ control1.isShutdown.futureValue shouldBe Done
+ control2.isShutdown.futureValue shouldBe Done
+ }
}
- "be removed from the partitioned source stage buffer when a partition is
revoked" in assertAllStagesStopped {
- def subSourcesWithProbes(
- partitions: Int,
- probe: TestSubscriber.Probe[(TopicPartition,
Source[ConsumerRecord[String, String], NotUsed])])
- : Seq[(TopicPartition, TestSubscriber.Probe[ConsumerRecord[String,
String]])] =
- probe
- .expectNextN(partitions.toLong)
- .map {
- case (tp, subSource) =>
- (tp, subSource.toMat(TestSink())(Keep.right).run())
+ List(EagerCase, CooperativeCase).foreach { mode =>
+ s"be removed from the partitioned source stage buffer when a partition
is revoked (${mode.label})" in
+ assertAllStagesStopped {
+ def subSourcesWithProbes(
+ partitions: Int,
+ probe: TestSubscriber.Probe[(TopicPartition,
Source[ConsumerRecord[String, String], NotUsed])])
+ : Seq[(TopicPartition, TestSubscriber.Probe[ConsumerRecord[String,
String]])] =
+ probe
+ .expectNextN(partitions.toLong)
+ .map {
+ case (tp, subSource) =>
+ (tp, subSource.toMat(TestSink())(Keep.right).run())
+ }
+
+ def runForSubSource(
+ partition: Int,
+ subSourcesWithProbes: Seq[(TopicPartition,
TestSubscriber.Probe[ConsumerRecord[String, String]])])(
+ fun: TestSubscriber.Probe[ConsumerRecord[String, String]] => Unit)
=
+ subSourcesWithProbes
+ .find { case (tp, _) => tp.partition() == partition }
+ .foreach { case (_, probe) => fun(probe) }
+
+ val count = 20L
+ // de-coupling consecutive test runs with crossScalaVersions on build
+ val topicSuffix = Random.nextInt()
+ val topic1 = createTopic(topicSuffix, partitions = 2)
+ val group1 = createGroupId(1)
+ val tp0 = new TopicPartition(topic1, partition0)
+ val tp1 = new TopicPartition(topic1, partition1)
+ val consumerSettings = consumerDefaults
+ .withProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "500") // 500
is the default value
+ .withProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
mode.assignor.getName)
+ .withGroupId(group1)
+
+ awaitProduce(produce(topic1, 0 to count.toInt, partition1))
+
+ PekkoConnectorsAssignor.clientIdToPartitionMap.set(
+ Map(
+ consumerClientId1 -> Set(tp0, tp1)))
+
+ log.debug("Subscribe to the topic (without demand)")
+ val probe1rebalanceActor = TestProbe()
+ val probe1subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe1rebalanceActor.ref)
+ val (control1, probe1) = Consumer
+
.plainPartitionedSource(consumerSettings.withClientId(consumerClientId1),
probe1subscription)
+ .toMat(TestSink())(Keep.both)
+ .run()
+
+ log.debug("Await initial partition assignment")
+ probe1rebalanceActor.expectMsg(
+ TopicPartitionsAssigned(probe1subscription,
+ Set(new TopicPartition(topic1, partition0), new
TopicPartition(topic1, partition1))))
+
+ log.debug("read 2 sub sources returned by partitioned source")
+ probe1.request(2)
+ val probe1RunningSubSourceProbes = subSourcesWithProbes(partitions =
2, probe1)
+
+ log.debug("read one message from probe1 sub source for partition 1")
+ probe1RunningSubSourceProbes
+ .find { case (tp, _) => tp.partition() == partition1 }
+ .foreach { case (_, probe) => probe.requestNext() }
+
+ PekkoConnectorsAssignor.clientIdToPartitionMap.set(
+ Map(
+ consumerClientId1 -> Set(tp0),
+ consumerClientId2 -> Set(tp1)))
+
+ log.debug("Subscribe to the topic (without demand)")
+ val probe2rebalanceActor = TestProbe()
+ val probe2subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe2rebalanceActor.ref)
+ val (control2, probe2) = Consumer
+
.plainPartitionedSource(consumerSettings.withClientId(consumerClientId2),
probe2subscription)
+ .toMat(TestSink())(Keep.both)
+ .run()
+
+ probe2.request(1)
+ val probe2RunningSubSourceProbes = subSourcesWithProbes(partitions =
1, probe2)
+
+ log.debug("Await the rebalance to complete")
+ mode.expectSecondConsumerJoin(probe1rebalanceActor,
probe1subscription, probe2rebalanceActor,
+ probe2subscription, tp0, tp1)
+
+ log.debug("resume demand on both consumers")
+ runForSubSource(partition = 1,
probe1RunningSubSourceProbes)(_.request(count))
+ runForSubSource(partition = 1,
probe2RunningSubSourceProbes)(_.request(count))
+
+ log.debug("no further messages enqueued on probe1 as partition 1 is
balanced away")
+ runForSubSource(partition = 1,
probe1RunningSubSourceProbes)(_.expectComplete())
+
+ val probe2messages = probe2RunningSubSourceProbes
+ .find { case (tp, _) => tp.partition() == partition1 }
+ .toList
+ .flatMap {
+ case (_, probe) =>
+ probe.expectNextN(count)
}
- def runForSubSource(
- partition: Int,
- subSourcesWithProbes: Seq[(TopicPartition,
TestSubscriber.Probe[ConsumerRecord[String, String]])])(
- fun: TestSubscriber.Probe[ConsumerRecord[String, String]] => Unit) =
- subSourcesWithProbes
- .find { case (tp, _) => tp.partition() == partition }
- .foreach { case (_, probe) => fun(probe) }
-
- val count = 20L
- // de-coupling consecutive test runs with crossScalaVersions on build
- val topicSuffix = Random.nextInt()
- val topic1 = createTopic(topicSuffix, partitions = 2)
- val group1 = createGroupId(1)
- val tp0 = new TopicPartition(topic1, partition0)
- val tp1 = new TopicPartition(topic1, partition1)
- val consumerSettings = consumerDefaults
- .withProperty(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, "500") // 500 is
the default value
- .withProperty(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG,
classOf[PekkoConnectorsAssignor].getName)
- .withGroupId(group1)
-
- awaitProduce(produce(topic1, 0 to count.toInt, partition1))
-
- PekkoConnectorsAssignor.clientIdToPartitionMap.set(
- Map(
- consumerClientId1 -> Set(tp0, tp1)))
-
- log.debug("Subscribe to the topic (without demand)")
- val probe1rebalanceActor = TestProbe()
- val probe1subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe1rebalanceActor.ref)
- val (control1, probe1) = Consumer
-
.plainPartitionedSource(consumerSettings.withClientId(consumerClientId1),
probe1subscription)
- .toMat(TestSink())(Keep.both)
- .run()
-
- log.debug("Await initial partition assignment")
- probe1rebalanceActor.expectMsg(
- TopicPartitionsAssigned(probe1subscription,
- Set(new TopicPartition(topic1, partition0), new
TopicPartition(topic1, partition1))))
-
- log.debug("read 2 sub sources returned by partitioned source")
- probe1.request(2)
- val probe1RunningSubSourceProbes = subSourcesWithProbes(partitions = 2,
probe1)
-
- log.debug("read one message from probe1 sub source for partition 1")
- probe1RunningSubSourceProbes
- .find { case (tp, _) => tp.partition() == partition1 }
- .foreach { case (_, probe) => probe.requestNext() }
-
- PekkoConnectorsAssignor.clientIdToPartitionMap.set(
- Map(
- consumerClientId1 -> Set(tp0),
- consumerClientId2 -> Set(tp1)))
-
- log.debug("Subscribe to the topic (without demand)")
- val probe2rebalanceActor = TestProbe()
- val probe2subscription =
Subscriptions.topics(topic1).withRebalanceListener(probe2rebalanceActor.ref)
- val (control2, probe2) = Consumer
-
.plainPartitionedSource(consumerSettings.withClientId(consumerClientId2),
probe2subscription)
- .toMat(TestSink())(Keep.both)
- .run()
-
- probe2.request(1)
- val probe2RunningSubSourceProbes = subSourcesWithProbes(partitions = 1,
probe2)
-
- log.debug("Await a revoke to consumer 1")
- probe1rebalanceActor.expectMsg(
- TopicPartitionsRevoked(probe1subscription,
- Set(new TopicPartition(topic1, partition0), new
TopicPartition(topic1, partition1))))
-
- log.debug("the rebalance finishes")
- probe1rebalanceActor.expectMsg(
- TopicPartitionsAssigned(probe1subscription, Set(new
TopicPartition(topic1, partition0))))
- probe2rebalanceActor.expectMsg(
- TopicPartitionsAssigned(probe2subscription, Set(new
TopicPartition(topic1, partition1))))
-
- log.debug("resume demand on both consumers")
- runForSubSource(partition = 1,
probe1RunningSubSourceProbes)(_.request(count))
- runForSubSource(partition = 1,
probe2RunningSubSourceProbes)(_.request(count))
-
- log.debug("no further messages enqueued on probe1 as partition 1 is
balanced away")
- runForSubSource(partition = 1,
probe1RunningSubSourceProbes)(_.expectComplete())
-
- val probe2messages = probe2RunningSubSourceProbes
- .find { case (tp, _) => tp.partition() == partition1 }
- .toList
- .flatMap {
- case (_, probe) =>
- probe.expectNextN(count)
- }
-
- probe2messages should have size count
+ probe2messages should have size count
- probe1.cancel()
- probe2.cancel()
+ probe1.cancel()
+ probe2.cancel()
- control1.isShutdown.futureValue shouldBe Done
- control2.isShutdown.futureValue shouldBe Done
+ control1.isShutdown.futureValue shouldBe Done
+ control2.isShutdown.futureValue shouldBe Done
+ }
}
}
}
@@ -290,3 +334,38 @@ class PekkoConnectorsAssignor extends
AbstractPartitionAssignor {
assignments.toMap.asJava
}
}
+
+/**
+ * Variant of [[PekkoConnectorsAssignor]] that uses the cooperative rebalance
protocol.
+ *
+ * The cooperative protocol requires that a partition never moves directly
from one member to
+ * another within a single rebalance: it must be absent from all assignments
for one generation
+ * (revoking it from its previous owner, which triggers a follow-up rebalance)
before it may be
+ * assigned to its new owner.
+ */
+class CooperativePekkoConnectorsAssignor extends PekkoConnectorsAssignor {
+
+ override def name(): String = "pekko-connector-kafka-test-cooperative"
+
+ override def supportedProtocols():
util.List[ConsumerPartitionAssignor.RebalanceProtocol] =
+ util.Arrays.asList(ConsumerPartitionAssignor.RebalanceProtocol.COOPERATIVE)
+
+ override def assign(
+ partitionsPerTopic: util.Map[String, Integer],
+ subscriptions: util.Map[String, ConsumerPartitionAssignor.Subscription])
+ : util.Map[String, util.List[TopicPartition]] = {
+ val desired = super.assign(partitionsPerTopic, subscriptions).asScala
+ val currentOwner: Map[TopicPartition, String] = (for {
+ (memberId, subscription) <- subscriptions.asScala.toSeq
+ tp <- subscription.ownedPartitions().asScala
+ } yield tp -> memberId).toMap
+
+ desired.map {
+ case (memberId, tps) =>
+ val withoutMovingPartitions = tps.asScala.filter { tp =>
+ currentOwner.get(tp).forall(_ == memberId)
+ }
+ memberId -> withoutMovingPartitions.asJava
+ }.asJava
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
