feat: Initial support for cooperative-sticky rebalancing #407
Conversation
Fix one bug in StreamProcessor where it assumed the passed assignments are replacing the old ones. Our consumer backends mostly work as-is, and are already passing the right values in callbacks.
There was a problem hiding this comment.
Looks good!
Were you wanting to merge + publish as-is or do more testing against this branch?
Is it worth including a test that fails against the current cluster but works with newer versions? Attempting to commit on an existing partition during a rebalance might do it.
| @@ -245,6 +256,10 @@ def test_consumer_polls_when_paused(self) -> None: | |||
| assert consumer.paused() == [] | |||
|
|
|||
|
|
|||
| class TestKafkaStreamsIncrementalRebalancing(TestKafkaStreams): | |||
There was a problem hiding this comment.
no, this actually re-declares all the tests in TestKafkaStreams, just re-running them with cooperative-sticky rebalancing
| # Second partition assigned | ||
| offsets_p1 = {Partition(topic, 1): 0} | ||
| assignment_callback(offsets_p1) | ||
|
|
||
| create_args, _ = factory.create_with_partitions.call_args | ||
| assert factory.create_with_partitions.call_count == 2 | ||
| assert create_args[1] == offsets_p1 | ||
| assert create_args[1] == {**offsets_p1, **offsets_p0} |
There was a problem hiding this comment.
was this test change related to your other changes? since there's no cooperative rebalancing here, seems like the assertions should stay the same?
There was a problem hiding this comment.
the mocked return value for consumer.tell was wrong, so this had the wrong value. the assignments in this test are actually incremental: first p1 is assigned, then p0, and there's no revocation.
arroyo/backends/kafka/consumer.py
Outdated
| @@ -161,6 +161,7 @@ def __init__( | |||
| ) | |||
|
|
|||
| configuration = dict(configuration) | |||
| self.__assignment_strategy = configuration.get("partition.assignment.strategy") | |||
There was a problem hiding this comment.
sorry i said the wrong thing earlier, this should be group.protocol
There was a problem hiding this comment.
after discussing offline i think we can support KIP-848 (group.protocol) as well as cooperative-sticky rebalancing. they're the same as far as rdkafka API is concerned. i just can't get it to work right now and might scope it out of this PR if it takes too much time.
can you elaborate on this? |
| logger.info("skipping empty assignment") | ||
| return |
There was a problem hiding this comment.
Why do you need a different logic for partition assignment between cooperative and standard in case of empty assignment?
I assume you can get an empty assignment in the cooperative rebalancing when, after a rebalancing, your assignments do not change. Is that the scenario where you do not want to touch the existing assignments ?
There was a problem hiding this comment.
after sleeping on it i agree. i only added this because it made cooperative rebalancing more comprehensible, and wasn't sure of the implications on regular rebalancing. i think we can skip empty assignments regardless of the assignment strategy.
| @@ -107,7 +107,7 @@ def test_dlq_policy_wrapper() -> None: | |||
| ) | |||
| partition = Partition(topic, 0) | |||
| wrapper = DlqPolicyWrapper(dlq_policy) | |||
| wrapper.reset_offsets({partition: 0}) | |||
| wrapper.reset_dlq_limits({partition: 0}) | |||
There was a problem hiding this comment.
this rename is just to align with rust btw
Fix one bug in StreamProcessor where it assumed the passed assignments
are replacing the old ones.
Our consumer backends mostly work as-is, and are already passing the
right values in callbacks.