KAFKA-18538: Add Streams membership manager #18551

cadonna · 2025-01-15T13:51:05Z

The Streams membership manager is used client-side in the
background thread of the async consumer. For each member
/consumer, it is responsible for:

keeping the member state,
keeping assignments for the member,
reconciling the assignments of the member -- for example
when tasks need to be revoked before other tasks are assigned
requesting invocations of assignment and revocation callbacks
by the stream thread.

The Streams membership manager is called by the background thread of
the async consumer, directly in its event loop and from the
Streams group heartbeat request manager. The Streams membership
manager uses the Streams rebalance events processor to request
assignment/revocation callback in the stream thread.

More detailed description of your change,
if necessary. The PR title and PR message become
the squashed commit message, so use a separate
comment to ping reviewers.

Summary of testing strategy (including rationale)
for the feature or bug fix. Unit and/or integration
tests are expected for any behaviour change and
system tests should be considered for larger changes.

Committer Checklist (excluded from commit message)

Verify design and implementation
Verify test coverage and CI build status
Verify documentation (including upgrade notes)

lucasbru

I made a first pass on the membership manager, but this is one massive PR. I am not done yet. But a more general point: In the original PR, we discussed changing what we have to change in the membership manager and then trying to share some code with the AbstractMembershipManager. Is that something we have abandoned / is it too hard to do? This class is a complex beast, I was hoping we can share the state machine, and just plug a different reconciliation logic.

If we are not going to share the code, I think we have to go through the last months of changes on AbstractMembershipManager and see which of the changes apply to our membership manager as well. It looks like a number of changes were made in the state machine since we have forked the membership manager - we need to check whether they need to be adopted by our membership manager.

lucasbru · 2025-01-21T10:44:13Z

.../src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java

+    private static final int PARTITION_0 = 0;
+    private static final int PARTITION_1 = 1;
+
+    private Time time = new MockTime(0);


lucasbru · 2025-01-21T10:44:20Z

.../src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java

+    private static final int PARTITION_1 = 1;
+
+    private Time time = new MockTime(0);
+    private Metrics metrics = new Metrics(time);


lucasbru · 2025-01-21T11:01:49Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+        }
+    }
+
+    public void onHeartbeatFailure(boolean retriable) {


This is not unit tested

lucasbru · 2025-01-21T11:02:11Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+        }
+    }
+
+    public void maybeRejoinStaleMember() {


this is not unit tested

lucasbru · 2025-01-21T11:03:08Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+        rejoinedWhileReconciliationInProgress = false;
+    }
+
+    public void onTasksRevokedCallbackCompleted(final StreamsOnTasksRevokedCallbackCompletedEvent event) {


This is not unit tested

lucasbru · 2025-01-21T13:49:11Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+    }
+
+    public boolean isLeavingGroup() {
+        MemberState state = state();


Not sure why we'd use the getter here. Maybe a result of copying?

Me neither, I copied this over. I will change it.

lucasbru · 2025-01-21T13:50:55Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+    }
+
+    public boolean shouldHeartbeatNow() {
+        MemberState state = state();


don't use the getter

lucasbru · 2025-01-21T13:50:59Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+    }
+
+    public void onHeartbeatRequestGenerated() {
+        MemberState state = state();


don't use the getter

lucasbru · 2025-01-21T13:52:44Z

.../src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java

+    private static final String TOPIC_0 = "topic-0";
+    private static final String TOPIC_1 = "topic-1";
+
+    private static final int PARTITION_0 = 0;


I think you asked me to call these tasks instead of partitions

You gave good reasons not to do it under my comment, so I actually thought you refrained from changing it.

lucasbru · 2025-01-21T13:53:56Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+        SortedSet<StreamsRebalanceData.TaskId> taskIdSet = new TreeSet<>();
+        for (final Map.Entry<String, SortedSet<Integer>> task : tasks.entrySet()) {
+            final String subtopologyId = task.getKey();
+            final SortedSet<Integer> partitions = task.getValue();


I think you asked me to call these tasks instead of partitions

cadonna · 2025-01-29T16:52:20Z

@lucasbru I included your feedback. I added javadocs mainly on public methods. Please re-review.

bbejeck

Thanks for PR @cadonna - I've made a pass over the StreamsMembershipManager I'll do the test in another pass

bbejeck · 2025-01-27T19:02:23Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+    }
+
+    public void registerStateListener(MemberStateListener listener) {
+        stateUpdatesListeners.add(Objects.requireNonNull(listener, "State updates listener cannot be null"));


Should stateUpdatesListeners be a Set or should we check for possibly adding a duplicate MemberStateListener?

Not allowing duplicate listener makes sense to me. However, I would prefer to check for reference equality instead of using a Set since the latter requires equals() and hashCode() being implemented which is not needed for listener IMO.

Added a check for duplicates.

bbejeck · 2025-01-28T15:26:20Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+        stateUpdatesListeners.add(Objects.requireNonNull(listener, "State updates listener cannot be null"));
+    }
+
+    void notifyEpochChange(Optional<Integer> epoch) {


Why the package-private modifier here?

bbejeck · 2025-01-29T20:55:55Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+     * the interval.
+     */
+    public boolean shouldHeartbeatNow() {
+        return state == MemberState.ACKNOWLEDGING || state == MemberState.LEAVING || state == MemberState.JOINING;


why heartbeat on leaving? more for my own knowledge

If a member leaves the group, it send a heartbeat with the leave epoch to the group coordinator. The leave epoch is -1 for non-static members and -2 for static members. So the group coordinator knows that the member can be safely removed from the group.
A leave group heartbeat request does not need to wait for the expiration of the heartbeat interval, but should be sent immediately.

bbejeck · 2025-01-29T21:28:13Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+     */
+    public void onFenced() {
+        if (state == MemberState.PREPARE_LEAVING) {
+            log.info("Member {} with epoch {} got fenced but it is already preparing to leave " +


the other log statements in onFenced are at the debug level, is there reason this one is info?

Good question! I actually do not know. I copied this over from the AbstractMembershipManager without noticing this. I am going to change the level to debug.

lucasbru

Getting there! I made a deeper pass on the main class. I did not look into the unit tests too closely, since Bill intended to do this.

lucasbru · 2025-01-30T09:35:30Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+     * If the member is already part of the group, this will only ensure that the updated subscription
+     * is included in the next heartbeat request.
+     * <p/>
+     * Note that list of topics of the subscription is taken from the shared subscription state.


Suggested change

* Note that list of topics of the subscription is taken from the shared subscription state.

* Note that the list of topics in the subscription is taken from the shared subscription state.

lucasbru · 2025-01-30T09:37:00Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+     * <p/>
+     * Note that list of topics of the subscription is taken from the shared subscription state.
+     */
+    public void onSubscriptionUpdated() {


What does SubscriptionUpdated mean in the context of Kafka Streams? It seems like we don't really subscribe to topics, so how does this relate to e.g. the topology?

Actually, we do subscribe:

kafka/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java

Line 1046 in ddfadf7

mainConsumer.unsubscribe();

The async consumer passes the subscription back to the network thread by means of an event. Processing that event calls onSubscriptionUpdated().

In future, we could consider removing the explicit subscription and let the StreamsRebalanceData or StreamsRebalanceEventsProcessor subscribe depending on the topology.

lucasbru · 2025-01-30T09:38:52Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+ * <ul>
+ *   <li>Keeping member state</li>
+ *   <li>Keeping assignment for the member</li>
+ *   <li>Reconciling assignment, for example if tasks need to be revoked before other tasks can be assigned</li>


Suggested change

* <li>Reconciling assignment, for example if tasks need to be revoked before other tasks can be assigned</li>

* <li>Reconciling assignment, for example, if tasks need to be revoked before other tasks can be assigned</li>

lucasbru · 2025-01-30T09:46:38Z

.../src/test/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManagerTest.java

+    private static final String MEMBER_ID = "test-member-1";
+    private static final int MEMBER_EPOCH = 1;
+
+    private static final String SUB_TOPOLOGY_ID_0 = "subtopology-0";


If the capitalization is subtopologyId, then this should be SUBTOPOLOGY_ID.

lucasbru · 2025-01-30T09:48:32Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+    public void onHeartbeatSuccess(StreamsGroupHeartbeatResponse response) {
+        StreamsGroupHeartbeatResponseData responseData = response.data();
+        throwIfUnexpectedError(responseData);
+        if (state == MemberState.LEAVING) {


For discussion: Should we test that the onHeartbeatSuccess is a NOOP in these three cases?

The case with state == MemberState.UNSUBSCRIBED && maybeCompleteLeaveInProgress() is not a noop, because it might complete the leaveGroupInProgress future in maybeCompleteLeaveInProgress(). That case is also covered by testLeaveGroupWhenInGroupWithAssignment().

For the other two cases, I need to see how I can sensibly verify a noop.

I verified the noop by checking if the member state listener was NOT called. See testOnHeartbeatSuccessWhenInFenced() for an example.

lucasbru · 2025-01-30T09:49:46Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+     */
+    private void transitionTo(MemberState nextState) {
+        if (!state.equals(nextState) && !nextState.getPreviousValidStates().contains(state)) {
+            throw new IllegalStateException(String.format("Invalid state transition from %s to %s",


Should we test detection of invalid state transitions?

I do not understand why this method is not part of the MemberState class. Then it would also be easy to test the transitions.

lucasbru · 2025-01-30T09:52:44Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+     * Notify when the poll timer expired.
+     */
+    public void onPollTimerExpired() {
+        transitionToSendingLeaveGroup(true);


Should we test what happens in LEAVING and UNSUBSCRIBED state?

Do you mean FATAL and UNSUBSCRIBED?

lucasbru · 2025-01-30T09:53:23Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+                    "current assignment {}. Nothing to reconcile.",
+                targetAssignment, currentAssignment);
+            // Make sure we transition the member back to STABLE if it was RECONCILING (ex.
+            // member was RECONCILING unresolved assignments that were just removed by the


What does it mean for an assignment to be unresolved if we do not use UUIDs?

copy & paste error?

You may want to replace this comment by a test.

I replaced the comment with a test.
It is not clear to me how a member can be in JOINING with a reconciled assignment. When the member starts up and joins it does not a reconciled assignment with the assignment that comes in with onHeartbeatSuccess(). If the member rejoins after being stale because it exceeded the poll time out the current task assignment is cleared. Thus, the assignment cannot be reconciled with the assignment that comes in with onHeartbeatSuccess().

lucasbru · 2025-01-30T10:01:11Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+public class StreamsMembershipManager implements RequestManager {
+
+    /**
+     * A data structure to represent the current task assignment, and current target task assignment of a member in a


You renamed current target task assignment to just target task assignment, so then you should probably update this comment as well.

lucasbru · 2025-01-30T10:03:26Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsMembershipManager.java

+    @Override
+    public NetworkClientDelegate.PollResult poll(long currentTimeMs) {
+        if (state == MemberState.RECONCILING) {
+            maybeReconcile();


I moved the reconciliation to poll in the consumer membership manager back then, because we may receive new metadata for resolving topic UUIDs. Does it still make sense to reconcile in poll?

I would like to postpone this broader discussion after merging this PR.

The Streams membership manager is used client-side in the background thread of the async consumer. For each member /consumer, it is responsible for: * keeping the member state, * keeping assignments for the member, * reconciling the assignments of the member -- for example when tasks need to be revoked before other tasks are assigned * requesting invocations of assignment and revocation callbacks by the stream thread. The Streams membership manager is called by the background thread of the async consumer, directly in its event loop and from the Streams group heartbeat request manager. The Streams membership manager uses the Streams rebalance events processor to request assignment/revocation callback in the stream thread.

cadonna added streams consumer KIP-1071 PRs related to KIP-1071 labels Jan 15, 2025

cadonna requested review from bbejeck, lucasbru and mjsax January 15, 2025 13:51

github-actions bot added triage PRs from the community clients labels Jan 15, 2025

cadonna removed the triage PRs from the community label Jan 15, 2025

cadonna force-pushed the streams_membership_manager branch from c58c298 to 67dd512 Compare January 15, 2025 21:33

lucasbru reviewed Jan 21, 2025

View reviewed changes

cadonna force-pushed the streams_membership_manager branch 2 times, most recently from 8403d02 to 6c65e3f Compare January 22, 2025 08:36

cadonna force-pushed the streams_membership_manager branch 2 times, most recently from b0eaf13 to 5e6812f Compare January 29, 2025 16:51

bbejeck reviewed Jan 29, 2025

View reviewed changes

cadonna force-pushed the streams_membership_manager branch from 5e6812f to 8dab352 Compare January 30, 2025 07:37

lucasbru reviewed Jan 30, 2025

View reviewed changes

cadonna added 8 commits January 31, 2025 18:03

Improve doc

828c8a8

fix compilation issues

b5e0326

include feedback

ab677ee

add javadocs

628b972

change log level for message about fencing

ef34da3

Check for duplicate state listener

050eb0b

Include feedback

5b2a196

cadonna force-pushed the streams_membership_manager branch from 295686b to 5b2a196 Compare January 31, 2025 17:05

	* Note that list of topics of the subscription is taken from the shared subscription state.
	* Note that the list of topics in the subscription is taken from the shared subscription state.

	* <li>Reconciling assignment, for example if tasks need to be revoked before other tasks can be assigned</li>
	* <li>Reconciling assignment, for example, if tasks need to be revoked before other tasks can be assigned</li>

KAFKA-18538: Add Streams membership manager #18551

Are you sure you want to change the base?

KAFKA-18538: Add Streams membership manager #18551

Conversation

cadonna commented Jan 15, 2025 • edited Loading

Committer Checklist (excluded from commit message)

lucasbru 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cadonna commented Jan 29, 2025

bbejeck left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cadonna Jan 30, 2025 • edited Loading

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

cadonna Jan 30, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lucasbru 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

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

cadonna Jan 31, 2025 • edited Loading

Choose a reason for hiding this comment

cadonna commented Jan 15, 2025 •

edited

Loading

cadonna Jan 30, 2025 •

edited

Loading

cadonna Jan 30, 2025 •

edited

Loading

cadonna Jan 31, 2025 •

edited

Loading