snapshot of https://github.com/seaweedfs/seaweedfs.git (master, 2026-09-16; blobs > 131072B omitted)

weed/mq/kafka/consumer/incremental_rebalancing_test.go

@@ -0,0 +1,399 @@
+package consumer
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestIncrementalCooperativeAssignmentStrategy_BasicAssignment(t *testing.T) {
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+
+	// Create members
+	members := []*GroupMember{
+		{
+			ID:           "member-1",
+			Subscription: []string{"topic-1"},
+			Assignment:   []PartitionAssignment{}, // No existing assignment
+		},
+		{
+			ID:           "member-2",
+			Subscription: []string{"topic-1"},
+			Assignment:   []PartitionAssignment{}, // No existing assignment
+		},
+	}
+
+	// Topic partitions
+	topicPartitions := map[string][]int32{
+		"topic-1": {0, 1, 2, 3},
+	}
+
+	// First assignment (no existing assignments, should be direct)
+	assignments := strategy.Assign(members, topicPartitions)
+
+	// Verify assignments
+	if len(assignments) != 2 {
+		t.Errorf("Expected 2 member assignments, got %d", len(assignments))
+	}
+
+	totalPartitions := 0
+	for memberID, partitions := range assignments {
+		t.Logf("Member %s assigned %d partitions: %v", memberID, len(partitions), partitions)
+		totalPartitions += len(partitions)
+	}
+
+	if totalPartitions != 4 {
+		t.Errorf("Expected 4 total partitions assigned, got %d", totalPartitions)
+	}
+
+	// Should not be in rebalance state for initial assignment
+	if strategy.IsRebalanceInProgress() {
+		t.Error("Expected no rebalance in progress for initial assignment")
+	}
+}
+
+func TestIncrementalCooperativeAssignmentStrategy_RebalanceWithRevocation(t *testing.T) {
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+
+	// Create members with existing assignments
+	members := []*GroupMember{
+		{
+			ID:           "member-1",
+			Subscription: []string{"topic-1"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 0},
+				{Topic: "topic-1", Partition: 1},
+				{Topic: "topic-1", Partition: 2},
+				{Topic: "topic-1", Partition: 3}, // This member has all partitions
+			},
+		},
+		{
+			ID:           "member-2",
+			Subscription: []string{"topic-1"},
+			Assignment:   []PartitionAssignment{}, // New member with no assignments
+		},
+	}
+
+	topicPartitions := map[string][]int32{
+		"topic-1": {0, 1, 2, 3},
+	}
+
+	// First call should start revocation phase
+	assignments1 := strategy.Assign(members, topicPartitions)
+
+	// Should be in revocation phase
+	if !strategy.IsRebalanceInProgress() {
+		t.Error("Expected rebalance to be in progress")
+	}
+
+	state := strategy.GetRebalanceState()
+	if state.Phase != RebalancePhaseRevocation {
+		t.Errorf("Expected revocation phase, got %s", state.Phase)
+	}
+
+	// Member-1 should have some partitions revoked
+	member1Assignments := assignments1["member-1"]
+	if len(member1Assignments) >= 4 {
+		t.Errorf("Expected member-1 to have fewer than 4 partitions after revocation, got %d", len(member1Assignments))
+	}
+
+	// Member-2 should still have no assignments during revocation
+	member2Assignments := assignments1["member-2"]
+	if len(member2Assignments) != 0 {
+		t.Errorf("Expected member-2 to have 0 partitions during revocation, got %d", len(member2Assignments))
+	}
+
+	t.Logf("Revocation phase - Member-1: %d partitions, Member-2: %d partitions",
+		len(member1Assignments), len(member2Assignments))
+
+	// Simulate time passing and second call (should move to assignment phase)
+	time.Sleep(10 * time.Millisecond)
+
+	// Force move to assignment phase by setting timeout to 0
+	state.RevocationTimeout = 0
+
+	assignments2 := strategy.Assign(members, topicPartitions)
+
+	// Should complete rebalance
+	if strategy.IsRebalanceInProgress() {
+		t.Error("Expected rebalance to be completed")
+	}
+
+	// Both members should have partitions now
+	member1FinalAssignments := assignments2["member-1"]
+	member2FinalAssignments := assignments2["member-2"]
+
+	if len(member1FinalAssignments) == 0 {
+		t.Error("Expected member-1 to have some partitions after rebalance")
+	}
+
+	if len(member2FinalAssignments) == 0 {
+		t.Error("Expected member-2 to have some partitions after rebalance")
+	}
+
+	totalFinalPartitions := len(member1FinalAssignments) + len(member2FinalAssignments)
+	if totalFinalPartitions != 4 {
+		t.Errorf("Expected 4 total partitions after rebalance, got %d", totalFinalPartitions)
+	}
+
+	t.Logf("Final assignment - Member-1: %d partitions, Member-2: %d partitions",
+		len(member1FinalAssignments), len(member2FinalAssignments))
+}
+
+func TestIncrementalCooperativeAssignmentStrategy_NoRevocationNeeded(t *testing.T) {
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+
+	// Create members with already balanced assignments
+	members := []*GroupMember{
+		{
+			ID:           "member-1",
+			Subscription: []string{"topic-1"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 0},
+				{Topic: "topic-1", Partition: 1},
+			},
+		},
+		{
+			ID:           "member-2",
+			Subscription: []string{"topic-1"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 2},
+				{Topic: "topic-1", Partition: 3},
+			},
+		},
+	}
+
+	topicPartitions := map[string][]int32{
+		"topic-1": {0, 1, 2, 3},
+	}
+
+	// Assignment should not trigger rebalance
+	assignments := strategy.Assign(members, topicPartitions)
+
+	// Should not be in rebalance state
+	if strategy.IsRebalanceInProgress() {
+		t.Error("Expected no rebalance in progress when assignments are already balanced")
+	}
+
+	// Assignments should remain the same
+	member1Assignments := assignments["member-1"]
+	member2Assignments := assignments["member-2"]
+
+	if len(member1Assignments) != 2 {
+		t.Errorf("Expected member-1 to keep 2 partitions, got %d", len(member1Assignments))
+	}
+
+	if len(member2Assignments) != 2 {
+		t.Errorf("Expected member-2 to keep 2 partitions, got %d", len(member2Assignments))
+	}
+}
+
+func TestIncrementalCooperativeAssignmentStrategy_MultipleTopics(t *testing.T) {
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+
+	// Create members with mixed topic subscriptions
+	members := []*GroupMember{
+		{
+			ID:           "member-1",
+			Subscription: []string{"topic-1", "topic-2"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 0},
+				{Topic: "topic-1", Partition: 1},
+				{Topic: "topic-2", Partition: 0},
+			},
+		},
+		{
+			ID:           "member-2",
+			Subscription: []string{"topic-1"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 2},
+			},
+		},
+		{
+			ID:           "member-3",
+			Subscription: []string{"topic-2"},
+			Assignment:   []PartitionAssignment{}, // New member
+		},
+	}
+
+	topicPartitions := map[string][]int32{
+		"topic-1": {0, 1, 2},
+		"topic-2": {0, 1},
+	}
+
+	// Should trigger rebalance to distribute topic-2 partitions
+	assignments := strategy.Assign(members, topicPartitions)
+
+	// Verify all partitions are assigned
+	allAssignedPartitions := make(map[string]bool)
+	for _, memberAssignments := range assignments {
+		for _, assignment := range memberAssignments {
+			key := fmt.Sprintf("%s:%d", assignment.Topic, assignment.Partition)
+			allAssignedPartitions[key] = true
+		}
+	}
+
+	expectedPartitions := []string{"topic-1:0", "topic-1:1", "topic-1:2", "topic-2:0", "topic-2:1"}
+	for _, expected := range expectedPartitions {
+		if !allAssignedPartitions[expected] {
+			t.Errorf("Expected partition %s to be assigned", expected)
+		}
+	}
+
+	// Debug: Print all assigned partitions
+	t.Logf("All assigned partitions: %v", allAssignedPartitions)
+}
+
+func TestIncrementalCooperativeAssignmentStrategy_ForceComplete(t *testing.T) {
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+
+	// Start a rebalance - create scenario where member-1 has all partitions but member-2 joins
+	members := []*GroupMember{
+		{
+			ID:           "member-1",
+			Subscription: []string{"topic-1"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 0},
+				{Topic: "topic-1", Partition: 1},
+				{Topic: "topic-1", Partition: 2},
+				{Topic: "topic-1", Partition: 3},
+			},
+		},
+		{
+			ID:           "member-2",
+			Subscription: []string{"topic-1"},
+			Assignment:   []PartitionAssignment{}, // New member
+		},
+	}
+
+	topicPartitions := map[string][]int32{
+		"topic-1": {0, 1, 2, 3},
+	}
+
+	// This should start a rebalance (member-2 needs partitions)
+	strategy.Assign(members, topicPartitions)
+
+	if !strategy.IsRebalanceInProgress() {
+		t.Error("Expected rebalance to be in progress")
+	}
+
+	// Force complete the rebalance
+	strategy.ForceCompleteRebalance()
+
+	if strategy.IsRebalanceInProgress() {
+		t.Error("Expected rebalance to be completed after force complete")
+	}
+
+	state := strategy.GetRebalanceState()
+	if state.Phase != RebalancePhaseNone {
+		t.Errorf("Expected phase to be None after force complete, got %s", state.Phase)
+	}
+}
+
+func TestIncrementalCooperativeAssignmentStrategy_RevocationTimeout(t *testing.T) {
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+
+	// Set a very short revocation timeout for testing
+	strategy.rebalanceState.RevocationTimeout = 1 * time.Millisecond
+
+	members := []*GroupMember{
+		{
+			ID:           "member-1",
+			Subscription: []string{"topic-1"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 0},
+				{Topic: "topic-1", Partition: 1},
+				{Topic: "topic-1", Partition: 2},
+				{Topic: "topic-1", Partition: 3},
+			},
+		},
+		{
+			ID:           "member-2",
+			Subscription: []string{"topic-1"},
+			Assignment:   []PartitionAssignment{},
+		},
+	}
+
+	topicPartitions := map[string][]int32{
+		"topic-1": {0, 1, 2, 3},
+	}
+
+	// First call starts revocation
+	strategy.Assign(members, topicPartitions)
+
+	if !strategy.IsRebalanceInProgress() {
+		t.Error("Expected rebalance to be in progress")
+	}
+
+	// Wait for timeout
+	time.Sleep(5 * time.Millisecond)
+
+	// Second call should complete due to timeout
+	assignments := strategy.Assign(members, topicPartitions)
+
+	if strategy.IsRebalanceInProgress() {
+		t.Error("Expected rebalance to be completed after timeout")
+	}
+
+	// Both members should have partitions
+	member1Assignments := assignments["member-1"]
+	member2Assignments := assignments["member-2"]
+
+	if len(member1Assignments) == 0 {
+		t.Error("Expected member-1 to have partitions after timeout")
+	}
+
+	if len(member2Assignments) == 0 {
+		t.Error("Expected member-2 to have partitions after timeout")
+	}
+}
+
+func TestIncrementalCooperativeAssignmentStrategy_StateTransitions(t *testing.T) {
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+
+	// Initial state should be None
+	state := strategy.GetRebalanceState()
+	if state.Phase != RebalancePhaseNone {
+		t.Errorf("Expected initial phase to be None, got %s", state.Phase)
+	}
+
+	// Create scenario that requires rebalancing
+	members := []*GroupMember{
+		{
+			ID:           "member-1",
+			Subscription: []string{"topic-1"},
+			Assignment: []PartitionAssignment{
+				{Topic: "topic-1", Partition: 0},
+				{Topic: "topic-1", Partition: 1},
+				{Topic: "topic-1", Partition: 2},
+				{Topic: "topic-1", Partition: 3},
+			},
+		},
+		{
+			ID:           "member-2",
+			Subscription: []string{"topic-1"},
+			Assignment:   []PartitionAssignment{}, // New member
+		},
+	}
+
+	topicPartitions := map[string][]int32{
+		"topic-1": {0, 1, 2, 3}, // Same partitions, but need rebalancing due to new member
+	}
+
+	// First call should move to revocation phase
+	strategy.Assign(members, topicPartitions)
+	state = strategy.GetRebalanceState()
+	if state.Phase != RebalancePhaseRevocation {
+		t.Errorf("Expected phase to be Revocation, got %s", state.Phase)
+	}
+
+	// Force timeout to move to assignment phase
+	state.RevocationTimeout = 0
+	strategy.Assign(members, topicPartitions)
+
+	// Should complete and return to None
+	state = strategy.GetRebalanceState()
+	if state.Phase != RebalancePhaseNone {
+		t.Errorf("Expected phase to be None after completion, got %s", state.Phase)
+	}
+}