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

test/kafka/integration/rebalancing_test.go

@@ -0,0 +1,474 @@
+package integration
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/IBM/sarama"
+	"github.com/seaweedfs/seaweedfs/test/kafka/internal/testutil"
+)
+
+// runConsumeLoop drives consumerGroup.Consume in a for-loop until ctx is cancelled
+// or the group is closed. This is the idiomatic Sarama pattern: Consume returns
+// when a rebalance cancels the session, and the caller is expected to re-enter.
+// Tests that observe rebalance behaviour MUST loop — otherwise the consumer goroutine
+// exits on the first rebalance and never sees the post-rebalance assignment.
+func runConsumeLoop(t *testing.T, ctx context.Context, name string, cg sarama.ConsumerGroup, topics []string, handler sarama.ConsumerGroupHandler) {
+	for {
+		if err := cg.Consume(ctx, topics, handler); err != nil {
+			if ctx.Err() != nil || errors.Is(err, sarama.ErrClosedConsumerGroup) {
+				return
+			}
+			t.Logf("%s Consume returned: %v (will retry)", name, err)
+		}
+		if ctx.Err() != nil {
+			return
+		}
+	}
+}
+
+func testSingleConsumerAllPartitions(t *testing.T, addr, topicName, groupID string) {
+	config := sarama.NewConfig()
+	config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRange
+	config.Consumer.Offsets.Initial = sarama.OffsetOldest
+	config.Consumer.Return.Errors = true
+
+	client, err := sarama.NewClient([]string{addr}, config)
+	testutil.AssertNoError(t, err, "Failed to create client")
+	defer client.Close()
+
+	consumerGroup, err := sarama.NewConsumerGroupFromClient(groupID, client)
+	testutil.AssertNoError(t, err, "Failed to create consumer group")
+	defer consumerGroup.Close()
+
+	handler := &RebalanceTestHandler{
+		messages:    make(chan *sarama.ConsumerMessage, 20),
+		ready:       make(chan bool),
+		assignments: make(chan []int32, 1),
+		t:           t,
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	// Start consumer
+	go func() {
+		err := consumerGroup.Consume(ctx, []string{topicName}, handler)
+		if err != nil && err != context.DeadlineExceeded {
+			t.Logf("Consumer error: %v", err)
+		}
+	}()
+
+	// Wait for consumer to be ready
+	<-handler.ready
+
+	// Wait for assignment
+	select {
+	case partitions := <-handler.assignments:
+		t.Logf("Single consumer assigned partitions: %v", partitions)
+		if len(partitions) != 4 {
+			t.Errorf("Expected single consumer to get all 4 partitions, got %d", len(partitions))
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("Timeout waiting for partition assignment")
+	}
+
+	// Consume some messages to verify functionality
+	consumedCount := 0
+	for consumedCount < 4 { // At least one from each partition
+		select {
+		case msg := <-handler.messages:
+			t.Logf("Consumed message from partition %d: %s", msg.Partition, string(msg.Value))
+			consumedCount++
+		case <-time.After(5 * time.Second):
+			t.Logf("Consumed %d messages so far", consumedCount)
+			break
+		}
+	}
+
+	if consumedCount == 0 {
+		t.Error("No messages consumed by single consumer")
+	}
+}
+
+func testTwoConsumersRebalance(t *testing.T, addr, topicName, groupID string) {
+	config := sarama.NewConfig()
+	config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRange
+	config.Consumer.Offsets.Initial = sarama.OffsetOldest
+	config.Consumer.Return.Errors = true
+
+	// Start first consumer
+	client1, err := sarama.NewClient([]string{addr}, config)
+	testutil.AssertNoError(t, err, "Failed to create client1")
+	defer client1.Close()
+
+	consumerGroup1, err := sarama.NewConsumerGroupFromClient(groupID, client1)
+	testutil.AssertNoError(t, err, "Failed to create consumer group 1")
+	defer consumerGroup1.Close()
+
+	handler1 := &RebalanceTestHandler{
+		messages:    make(chan *sarama.ConsumerMessage, 20),
+		ready:       make(chan bool),
+		assignments: make(chan []int32, 1),
+		t:           t,
+		name:        "Consumer1",
+	}
+
+	ctx1, cancel1 := context.WithTimeout(context.Background(), 45*time.Second)
+	defer cancel1()
+
+	go runConsumeLoop(t, ctx1, "Consumer1", consumerGroup1, []string{topicName}, handler1)
+
+	// Wait for first consumer to be ready and get initial assignment
+	<-handler1.ready
+	select {
+	case partitions := <-handler1.assignments:
+		t.Logf("Consumer1 initial assignment: %v", partitions)
+		if len(partitions) != 4 {
+			t.Errorf("Expected Consumer1 to initially get all 4 partitions, got %d", len(partitions))
+		}
+	case <-time.After(10 * time.Second):
+		t.Fatal("Timeout waiting for Consumer1 initial assignment")
+	}
+
+	// Start second consumer
+	client2, err := sarama.NewClient([]string{addr}, config)
+	testutil.AssertNoError(t, err, "Failed to create client2")
+	defer client2.Close()
+
+	consumerGroup2, err := sarama.NewConsumerGroupFromClient(groupID, client2)
+	testutil.AssertNoError(t, err, "Failed to create consumer group 2")
+	defer consumerGroup2.Close()
+
+	handler2 := &RebalanceTestHandler{
+		messages:    make(chan *sarama.ConsumerMessage, 20),
+		ready:       make(chan bool),
+		assignments: make(chan []int32, 1),
+		t:           t,
+		name:        "Consumer2",
+	}
+
+	ctx2, cancel2 := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel2()
+
+	go runConsumeLoop(t, ctx2, "Consumer2", consumerGroup2, []string{topicName}, handler2)
+
+	// Wait for second consumer to be ready
+	<-handler2.ready
+
+	// Wait for rebalancing to occur - both consumers should get new assignments
+	var rebalancedAssignment1, rebalancedAssignment2 []int32
+
+	// Consumer1 should get a rebalance assignment
+	select {
+	case partitions := <-handler1.assignments:
+		rebalancedAssignment1 = partitions
+		t.Logf("Consumer1 rebalanced assignment: %v", partitions)
+	case <-time.After(15 * time.Second):
+		t.Error("Timeout waiting for Consumer1 rebalance assignment")
+	}
+
+	// Consumer2 should get its assignment
+	select {
+	case partitions := <-handler2.assignments:
+		rebalancedAssignment2 = partitions
+		t.Logf("Consumer2 assignment: %v", partitions)
+	case <-time.After(15 * time.Second):
+		t.Error("Timeout waiting for Consumer2 assignment")
+	}
+
+	// Verify rebalancing occurred correctly
+	totalPartitions := len(rebalancedAssignment1) + len(rebalancedAssignment2)
+	if totalPartitions != 4 {
+		t.Errorf("Expected total of 4 partitions assigned, got %d", totalPartitions)
+	}
+
+	// Each consumer should have at least 1 partition, and no more than 3
+	if len(rebalancedAssignment1) == 0 || len(rebalancedAssignment1) > 3 {
+		t.Errorf("Consumer1 should have 1-3 partitions, got %d", len(rebalancedAssignment1))
+	}
+	if len(rebalancedAssignment2) == 0 || len(rebalancedAssignment2) > 3 {
+		t.Errorf("Consumer2 should have 1-3 partitions, got %d", len(rebalancedAssignment2))
+	}
+
+	// Verify no partition overlap
+	partitionSet := make(map[int32]bool)
+	for _, p := range rebalancedAssignment1 {
+		if partitionSet[p] {
+			t.Errorf("Partition %d assigned to multiple consumers", p)
+		}
+		partitionSet[p] = true
+	}
+	for _, p := range rebalancedAssignment2 {
+		if partitionSet[p] {
+			t.Errorf("Partition %d assigned to multiple consumers", p)
+		}
+		partitionSet[p] = true
+	}
+
+	t.Logf("Rebalancing test completed successfully")
+}
+
+func testConsumerLeaveRebalance(t *testing.T, addr, topicName, groupID string) {
+	config := sarama.NewConfig()
+	config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRange
+	config.Consumer.Offsets.Initial = sarama.OffsetOldest
+	config.Consumer.Return.Errors = true
+
+	// Start two consumers
+	client1, err := sarama.NewClient([]string{addr}, config)
+	testutil.AssertNoError(t, err, "Failed to create client1")
+	defer client1.Close()
+
+	client2, err := sarama.NewClient([]string{addr}, config)
+	testutil.AssertNoError(t, err, "Failed to create client2")
+	defer client2.Close()
+
+	consumerGroup1, err := sarama.NewConsumerGroupFromClient(groupID, client1)
+	testutil.AssertNoError(t, err, "Failed to create consumer group 1")
+	defer consumerGroup1.Close()
+
+	consumerGroup2, err := sarama.NewConsumerGroupFromClient(groupID, client2)
+	testutil.AssertNoError(t, err, "Failed to create consumer group 2")
+
+	handler1 := &RebalanceTestHandler{
+		messages:    make(chan *sarama.ConsumerMessage, 20),
+		ready:       make(chan bool),
+		assignments: make(chan []int32, 1),
+		t:           t,
+		name:        "Consumer1",
+	}
+
+	handler2 := &RebalanceTestHandler{
+		messages:    make(chan *sarama.ConsumerMessage, 20),
+		ready:       make(chan bool),
+		assignments: make(chan []int32, 1),
+		t:           t,
+		name:        "Consumer2",
+	}
+
+	ctx1, cancel1 := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel1()
+
+	ctx2, cancel2 := context.WithTimeout(context.Background(), 30*time.Second)
+
+	// Start both consumers with loops so they survive rebalance session exits.
+	go runConsumeLoop(t, ctx1, "Consumer1", consumerGroup1, []string{topicName}, handler1)
+	go runConsumeLoop(t, ctx2, "Consumer2", consumerGroup2, []string{topicName}, handler2)
+
+	// Wait for both consumers to be ready
+	<-handler1.ready
+	<-handler2.ready
+
+	// Wait for initial assignments
+	<-handler1.assignments
+	<-handler2.assignments
+
+	t.Logf("Both consumers started, now stopping Consumer2")
+
+	// Stop second consumer (simulate leave)
+	cancel2()
+	consumerGroup2.Close()
+
+	// Wait for Consumer1 to get rebalanced assignment (should get all partitions)
+	select {
+	case partitions := <-handler1.assignments:
+		t.Logf("Consumer1 rebalanced assignment after Consumer2 left: %v", partitions)
+		if len(partitions) != 4 {
+			t.Errorf("Expected Consumer1 to get all 4 partitions after Consumer2 left, got %d", len(partitions))
+		}
+	case <-time.After(20 * time.Second):
+		t.Error("Timeout waiting for Consumer1 rebalance after Consumer2 left")
+	}
+
+	t.Logf("Consumer leave rebalancing test completed successfully")
+}
+
+func testMultipleConsumersJoin(t *testing.T, addr, topicName, groupID string) {
+	config := sarama.NewConfig()
+	config.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRange
+	config.Consumer.Offsets.Initial = sarama.OffsetOldest
+	config.Consumer.Return.Errors = true
+
+	numConsumers := 4
+	consumers := make([]sarama.ConsumerGroup, numConsumers)
+	clients := make([]sarama.Client, numConsumers)
+	handlers := make([]*RebalanceTestHandler, numConsumers)
+	contexts := make([]context.Context, numConsumers)
+	cancels := make([]context.CancelFunc, numConsumers)
+
+	// Start all consumers simultaneously
+	for i := 0; i < numConsumers; i++ {
+		client, err := sarama.NewClient([]string{addr}, config)
+		testutil.AssertNoError(t, err, fmt.Sprintf("Failed to create client%d", i))
+		clients[i] = client
+
+		consumerGroup, err := sarama.NewConsumerGroupFromClient(groupID, client)
+		testutil.AssertNoError(t, err, fmt.Sprintf("Failed to create consumer group %d", i))
+		consumers[i] = consumerGroup
+
+		handlers[i] = &RebalanceTestHandler{
+			messages:    make(chan *sarama.ConsumerMessage, 20),
+			ready:       make(chan bool),
+			assignments: make(chan []int32, 1),
+			t:           t,
+			name:        fmt.Sprintf("Consumer%d", i),
+		}
+
+		contexts[i], cancels[i] = context.WithTimeout(context.Background(), 45*time.Second)
+
+		go func(idx int) {
+			runConsumeLoop(t, contexts[idx], fmt.Sprintf("Consumer%d", idx),
+				consumers[idx], []string{topicName}, handlers[idx])
+		}(i)
+	}
+
+	// Cleanup
+	defer func() {
+		for i := 0; i < numConsumers; i++ {
+			cancels[i]()
+			consumers[i].Close()
+			clients[i].Close()
+		}
+	}()
+
+	// Wait for all consumers to be ready. Four consumers racing in produces
+	// several rebalance rounds; each round costs ~HeartbeatInterval (3s) before
+	// the prior leader sees REBALANCE_IN_PROGRESS and rejoins. Give enough
+	// headroom for the group to converge.
+	for i := 0; i < numConsumers; i++ {
+		select {
+		case <-handlers[i].ready:
+			t.Logf("Consumer%d ready", i)
+		case <-time.After(30 * time.Second):
+			t.Fatalf("Timeout waiting for Consumer%d to be ready", i)
+		}
+	}
+
+	// Collect the settled assignment for each consumer. Because each consumer's
+	// Consume() returns and re-enters on every rebalance, Setup fires multiple
+	// times while the group converges and the handler's assignments channel
+	// accumulates several snapshots. Drain it and keep the most recent value.
+	assignments := make([][]int32, numConsumers)
+	deadline := time.Now().Add(20 * time.Second)
+	for i := 0; i < numConsumers; i++ {
+		select {
+		case partitions := <-handlers[i].assignments:
+			assignments[i] = partitions
+		case <-time.After(time.Until(deadline)):
+			t.Errorf("Timeout waiting for Consumer%d assignment", i)
+			continue
+		}
+	drain:
+		for {
+			select {
+			case partitions := <-handlers[i].assignments:
+				assignments[i] = partitions
+			case <-time.After(3 * time.Second):
+				break drain
+			}
+		}
+		t.Logf("Consumer%d final assignment: %v", i, assignments[i])
+	}
+
+	// Verify all partitions are assigned exactly once
+	assignedPartitions := make(map[int32]int)
+	totalAssigned := 0
+	for i, assignment := range assignments {
+		totalAssigned += len(assignment)
+		for _, partition := range assignment {
+			assignedPartitions[partition]++
+			if assignedPartitions[partition] > 1 {
+				t.Errorf("Partition %d assigned to multiple consumers", partition)
+			}
+		}
+
+		// Each consumer should get exactly 1 partition (4 partitions / 4 consumers)
+		if len(assignment) != 1 {
+			t.Errorf("Consumer%d should get exactly 1 partition, got %d", i, len(assignment))
+		}
+	}
+
+	if totalAssigned != 4 {
+		t.Errorf("Expected 4 total partitions assigned, got %d", totalAssigned)
+	}
+
+	// Verify all partitions 0-3 are assigned
+	for i := int32(0); i < 4; i++ {
+		if assignedPartitions[i] != 1 {
+			t.Errorf("Partition %d assigned %d times, expected 1", i, assignedPartitions[i])
+		}
+	}
+
+	t.Logf("Multiple consumers join test completed successfully")
+}
+
+// RebalanceTestHandler implements sarama.ConsumerGroupHandler with rebalancing awareness
+type RebalanceTestHandler struct {
+	messages    chan *sarama.ConsumerMessage
+	ready       chan bool
+	assignments chan []int32
+	readyOnce   sync.Once
+	t           *testing.T
+	name        string
+}
+
+func (h *RebalanceTestHandler) Setup(session sarama.ConsumerGroupSession) error {
+	h.t.Logf("%s: Consumer group session setup", h.name)
+	h.readyOnce.Do(func() {
+		close(h.ready)
+	})
+
+	// Send partition assignment
+	partitions := make([]int32, 0)
+	for topic, partitionList := range session.Claims() {
+		h.t.Logf("%s: Assigned topic %s with partitions %v", h.name, topic, partitionList)
+		for _, partition := range partitionList {
+			partitions = append(partitions, partition)
+		}
+	}
+
+	// Always publish the latest assignment. If the channel is full, drain one
+	// stale value first so readers see the most recent rebalance outcome (not
+	// whatever arrived first while the group was still converging).
+	for {
+		select {
+		case h.assignments <- partitions:
+			return nil
+		default:
+			select {
+			case <-h.assignments:
+			default:
+			}
+		}
+	}
+}
+
+func (h *RebalanceTestHandler) Cleanup(sarama.ConsumerGroupSession) error {
+	h.t.Logf("%s: Consumer group session cleanup", h.name)
+	return nil
+}
+
+func (h *RebalanceTestHandler) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
+	for {
+		select {
+		case message := <-claim.Messages():
+			if message == nil {
+				return nil
+			}
+			h.t.Logf("%s: Received message from partition %d: %s", h.name, message.Partition, string(message.Value))
+			select {
+			case h.messages <- message:
+			default:
+				// Channel full, drop message for test
+			}
+			session.MarkMessage(message, "")
+		case <-session.Context().Done():
+			return nil
+		}
+	}
+}