KAFKA-21060: Make StickyAssignor aware of stateless vs stateful tasks

[gabriellefu]

The original sticky assignor in streams group didn't aware of
stateless/stateful assignor, which can cause some member assigned a lot
of stateful active task while some assigned a lot of stateless tasks,
cause the later member being stale

This pr change the assign order, first stateful active tasks are
assigned, then stateless active tasks, then standby tasks. Each type of
capacity are calculated to ensure each member can have a "balanced"
number of each type of tasks.

Reviewers: Lucas Brutschy lbrutschy@confluent.io

[Copilot]

Copilot review overview

🟢 Approval recommended

The assignment changes are covered by focused tests and no unresolved correctness issues were found.

Review effort: Lite
Findings: None

What changed in this PR

Updates StickyTaskAssignor to balance stateful and stateless active tasks separately while preserving total task capacity and standby behavior.

Changes:

• Adds stateful/stateless task tracking and load comparison.
• Assigns stateful active tasks before stateless tasks.
• Adds focused balancing and stickiness tests.

File	Description
StickyTaskAssignor.java	Implements task-type-aware assignment and quotas.
ProcessState.java	Tracks active task categories and stateless load.
StickyTaskAssignorTest.java	Adds mixed-task balancing tests.
ProcessStateTest.java	Updates tests for the expanded task API.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[lucasbru]

Mostly looks good to me. I left two comments.

I would still benchmark and random-test this against the current sticky assignor to avoid regressions.

[lucasbru]

nit: This quota check only works because assignActive(stateful) always runs before assignActive(stateless)/assignStandby, so the member's task count still equals its stateful count at this point. Nothing enforces that ordering though - worth a comment calling out the invariant, so a future reorder doesn't silently corrupt the quota check?

[gabriellefu]

I have added the comment, thanks for the suggestion

[lucasbru]

Do we need to copy all three task lists here? assignActive consumes its input, but assignStandby only sorts and iterates its list. Could we preserve statefulActiveTaskIds for standby assignment and pass the stateless list directly, avoiding the extra LinkedList copy?

[gabriellefu]

I have updated my pr, thanks

[gabriellefu]

JMH StreamsStickyAssignorBenchmark.doAssignment, trunk 1f419e29f0 vs branch 83d226b586, full parameter matrix (128 combinations), -w 3 -r 3, 5 warmup + 5 measurement iterations, 1 fork, same machine, runs not overlapping.

group	n	median	mean	min	max	significant slower / faster
all	128	-1.4%	-1.5%	-11.2%	+8.5%	9 / 13
FULL	64	-0.8%	-0.7%	-11.2%	+8.5%	6 / 6
INCREMENTAL	64	-1.8%	-2.4%	-9.9%	+4.3%	3 / 7
FULL, membersPerProcess=1	32	+0.4%	+0.6%	-8.8%	+8.5%	5 / 1
FULL, membersPerProcess=50	32	-2.2%	-2.1%	-11.2%	+4.1%	1 / 5
INCREMENTAL, membersPerProcess=1	32	-0.6%	-0.9%	-7.7%	+4.3%	3 / 3
INCREMENTAL, membersPerProcess=50	32	-4.3%	-3.9%	-9.9%	+3.9%	0 / 4

"Significant" = the difference exceeds the sum of both 99.9% error bars. Rows where that holds:

assignmentType	memberCount	membersPerProcess	partitionCount	reportedOffsets	standbyReplicas	subtopologyCount	trunk ms/op	branch ms/op	delta
FULL	1000	1	10	NONE	0	100	0.383 ± 0.005	0.416 ± 0.006	+8.5%
FULL	1000	1	10	NONE	0	10	0.117 ± 0.001	0.127 ± 0.002	+8.5%
FULL	1000	1	100	NONE	0	100	2.386 ± 0.070	2.568 ± 0.040	+7.6%
FULL	100	1	10	NONE	0	100	0.193 ± 0.009	0.207 ± 0.003	+6.8%
FULL	1000	1	10	NONE	1	10	0.137 ± 0.003	0.144 ± 0.002	+4.9%
INCREMENTAL	1000	1	100	NONE	0	10	0.313 ± 0.005	0.326 ± 0.006	+4.1%
INCREMENTAL	1000	1	10	NONE	1	10	0.142 ± 0.003	0.148 ± 0.001	+4.1%
FULL	1000	50	10	OWNED_AND_DORMANT	1	10	0.096 ± 0.002	0.100 ± 0.001	+4.1%
INCREMENTAL	100	1	10	OWNED_AND_DORMANT	1	10	0.046 ± 0.000	0.047 ± 0.000	+3.5%
FULL	100	50	10	NONE	1	10	0.035 ± 0.000	0.034 ± 0.000	-2.4%
INCREMENTAL	100	50	10	NONE	1	10	0.033 ± 0.000	0.031 ± 0.000	-3.8%
INCREMENTAL	100	1	10	OWNED_AND_DORMANT	1	100	0.339 ± 0.009	0.324 ± 0.002	-4.4%
FULL	100	1	100	OWNED_AND_DORMANT	1	10	0.269 ± 0.002	0.257 ± 0.006	-4.5%
INCREMENTAL	100	50	10	NONE	1	100	0.270 ± 0.004	0.258 ± 0.003	-4.6%
INCREMENTAL	100	1	100	OWNED_AND_DORMANT	0	10	0.161 ± 0.003	0.153 ± 0.002	-4.8%
INCREMENTAL	100	50	10	OWNED_AND_DORMANT	1	10	0.039 ± 0.000	0.037 ± 0.000	-4.8%
FULL	1000	50	10	NONE	1	10	0.097 ± 0.004	0.092 ± 0.001	-5.2%
FULL	100	50	10	NONE	0	10	0.025 ± 0.000	0.024 ± 0.000	-5.4%
FULL	100	50	100	NONE	0	10	0.183 ± 0.005	0.171 ± 0.002	-6.3%
INCREMENTAL	100	1	100	NONE	0	10	0.133 ± 0.006	0.124 ± 0.002	-6.5%
FULL	100	50	10	OWNED_AND_DORMANT	0	10	0.029 ± 0.001	0.027 ± 0.000	-6.6%
INCREMENTAL	100	50	100	OWNED_AND_DORMANT	0	100	2.394 ± 0.117	2.164 ± 0.085	-9.6%

[gabriellefu]

JMH: no measurable regression. StreamsStickyAssignorBenchmark.doAssignment, full 128-combination matrix, trunk vs this branch, same machine, -w 3 -r 3.

• Median delta -1.4% (mean -1.5%). 22 of 128 rows differ beyond the combined 99.9% error bars: 9 slower, 13 faster.
• The only consistent slowdown is FULL assignment with membersPerProcess=1 and no reported offsets, i.e. a cold-start assignment where every process has a single thread: +7.6% to +8.5% at
  memberCount=1000 (0.12 → 0.13 ms/op, 2.39 → 2.57 ms/op). This is the fixed cost of the second pass; all other slower rows are under 5% and scattered.
• Speedups of 5–11% cluster in membersPerProcess=50 and INCREMENTAL cases, where the stateful-first fill leaves fewer heap operations for the stateless pass.

[lucasbru]

Nice! Benchmarks look good. Since you are anyways at it, you should probably run the testbed on this assignor change before merging. But it looks good to me.

[lucasbru]

LGTM, thanks!

[gabriellefu]

Fuzz testbed comparison: trunk vs this PR

Run with the shared assignor testbed from #23556 (StickyTaskAssignorFuzzTest, default base seed 42, identical generated scenarios on both sides). Every assignment passed the invariants and the
sticky-specific checks on both branches; the numbers below are the graded metrics (max−min rows: lower is better).

Headline

profile	metric	trunk	this PR
LARGE, events	stateful active per member (max−min)	3.65 (max 9)	1.10 (max 3)
LARGE, events	stateless active per member (max−min)	3.24 (max 9)	1.58 (max 2)
LARGE, events	active per member (max−min)	1.04	1.17
LARGE, events	active task moves per rebalance	6.1	7.2
LARGE, events	process stickiness	0.989	0.988
LARGE, events	convergence iterations (avg / max)	2.06 / 6	2.18 / 7
LARGE, fresh	stateful active per member (max−min)	3.10 (max 6)	1.00 (max 1)
SMALL, events	stateful active per member (max−min)	1.03 (max 4)	0.73 (max 2)
SMALL, fresh	stateful active per member (max−min)	0.99 (max 3)	0.75 (max 1)

Stateful imbalance drops ~3× at scale and reaches the floor of 1 on fresh groups; stateless balance improves as a side effect. Cost: total-active spread +0.12 on LARGE (two buckets each rounding up),
~1 extra active move per rebalance, no measurable change in stickiness. Standby and rack-diversity rows are unchanged within noise, as expected since assignStandby is untouched.

[gabriellefu]

in the fuzzy test, fresh means it's the first time tasks being assigned, so no previous holder of any task, and events means we randomly choose event( member join, member leave, process join, process leave etc)

[lucasbru]

Nice work, @gabriellefu !