clean up consumer protocols

weed/mq/kafka/consumer/assignment.go

@@ -4,6 +4,14 @@ import (
 	"sort"
 )
 
+// Assignment strategy protocol names
+const (
+	ProtocolNameRange             = "range"
+	ProtocolNameRoundRobin        = "roundrobin"
+	ProtocolNameSticky            = "sticky"
+	ProtocolNameCooperativeSticky = "cooperative-sticky"
+)
+
 // AssignmentStrategy defines how partitions are assigned to consumers
 type AssignmentStrategy interface {
 	Name() string
@@ -15,7 +23,7 @@ type AssignmentStrategy interface {
 type RangeAssignmentStrategy struct{}
 
 func (r *RangeAssignmentStrategy) Name() string {
-	return "range"
+	return ProtocolNameRange
 }
 
 func (r *RangeAssignmentStrategy) Assign(members []*GroupMember, topicPartitions map[string][]int32) map[string][]PartitionAssignment {
@@ -104,7 +112,7 @@ func (r *RangeAssignmentStrategy) Assign(members []*GroupMember, topicPartitions
 type RoundRobinAssignmentStrategy struct{}
 
 func (rr *RoundRobinAssignmentStrategy) Name() string {
-	return "roundrobin"
+	return ProtocolNameRoundRobin
 }
 
 func (rr *RoundRobinAssignmentStrategy) Assign(members []*GroupMember, topicPartitions map[string][]int32) map[string][]PartitionAssignment {
@@ -194,191 +202,14 @@ func (rr *RoundRobinAssignmentStrategy) Assign(members []*GroupMember, topicPart
 	return assignments
 }
 
-// CooperativeStickyAssignmentStrategy implements the cooperative-sticky assignment strategy
-// This strategy tries to minimize partition movement during rebalancing while ensuring fairness
-type CooperativeStickyAssignmentStrategy struct{}
-
-func (cs *CooperativeStickyAssignmentStrategy) Name() string {
-	return "cooperative-sticky"
-}
-
-func (cs *CooperativeStickyAssignmentStrategy) Assign(members []*GroupMember, topicPartitions map[string][]int32) map[string][]PartitionAssignment {
-	if len(members) == 0 {
-		return make(map[string][]PartitionAssignment)
-	}
-
-	assignments := make(map[string][]PartitionAssignment)
-	for _, member := range members {
-		assignments[member.ID] = make([]PartitionAssignment, 0)
-	}
-
-	// Sort members for consistent assignment
-	sortedMembers := make([]*GroupMember, len(members))
-	copy(sortedMembers, members)
-	sort.Slice(sortedMembers, func(i, j int) bool {
-		return sortedMembers[i].ID < sortedMembers[j].ID
-	})
-
-	// Get all subscribed topics
-	subscribedTopics := make(map[string]bool)
-	for _, member := range members {
-		for _, topic := range member.Subscription {
-			subscribedTopics[topic] = true
-		}
-	}
-
-	// Collect all partitions that need assignment
-	allPartitions := make([]PartitionAssignment, 0)
-	for topic := range subscribedTopics {
-		partitions, exists := topicPartitions[topic]
-		if !exists {
-			continue
-		}
-
-		for _, partition := range partitions {
-			allPartitions = append(allPartitions, PartitionAssignment{
-				Topic:     topic,
-				Partition: partition,
-			})
-		}
-	}
-
-	// Sort partitions for consistent assignment
-	sort.Slice(allPartitions, func(i, j int) bool {
-		if allPartitions[i].Topic != allPartitions[j].Topic {
-			return allPartitions[i].Topic < allPartitions[j].Topic
-		}
-		return allPartitions[i].Partition < allPartitions[j].Partition
-	})
-
-	// Calculate target assignment counts for fairness
-	totalPartitions := len(allPartitions)
-	numMembers := len(sortedMembers)
-	baseAssignments := totalPartitions / numMembers
-	extraAssignments := totalPartitions % numMembers
-
-	// Phase 1: Try to preserve existing assignments (sticky behavior) but respect fairness
-	currentAssignments := make(map[string]map[PartitionAssignment]bool)
-	for _, member := range sortedMembers {
-		currentAssignments[member.ID] = make(map[PartitionAssignment]bool)
-		for _, assignment := range member.Assignment {
-			currentAssignments[member.ID][assignment] = true
-		}
-	}
-
-	// Track which partitions are already assigned
-	assignedPartitions := make(map[PartitionAssignment]bool)
-
-	// Preserve existing assignments where possible, but respect target counts
-	for i, member := range sortedMembers {
-		// Calculate target count for this member
-		targetCount := baseAssignments
-		if i < extraAssignments {
-			targetCount++
-		}
-
-		assignedCount := 0
-		for assignment := range currentAssignments[member.ID] {
-			// Stop if we've reached the target count for this member
-			if assignedCount >= targetCount {
-				break
-			}
-
-			// Check if member is still subscribed to this topic
-			subscribed := false
-			for _, topic := range member.Subscription {
-				if topic == assignment.Topic {
-					subscribed = true
-					break
-				}
-			}
-
-			if subscribed && !assignedPartitions[assignment] {
-				assignments[member.ID] = append(assignments[member.ID], assignment)
-				assignedPartitions[assignment] = true
-				assignedCount++
-			}
-		}
-	}
-
-	// Phase 2: Assign remaining partitions using round-robin for fairness
-	unassignedPartitions := make([]PartitionAssignment, 0)
-	for _, partition := range allPartitions {
-		if !assignedPartitions[partition] {
-			unassignedPartitions = append(unassignedPartitions, partition)
-		}
-	}
-
-	// Assign remaining partitions to achieve fairness
-	memberIndex := 0
-	for _, partition := range unassignedPartitions {
-		// Find a member that needs more partitions and is subscribed to this topic
-		assigned := false
-		startIndex := memberIndex
-
-		for !assigned {
-			member := sortedMembers[memberIndex]
-
-			// Check if this member is subscribed to the topic
-			subscribed := false
-			for _, topic := range member.Subscription {
-				if topic == partition.Topic {
-					subscribed = true
-					break
-				}
-			}
-
-			if subscribed {
-				// Calculate target count for this member
-				targetCount := baseAssignments
-				if memberIndex < extraAssignments {
-					targetCount++
-				}
-
-				// Assign if member needs more partitions
-				if len(assignments[member.ID]) < targetCount {
-					assignments[member.ID] = append(assignments[member.ID], partition)
-					assigned = true
-				}
-			}
-
-			memberIndex = (memberIndex + 1) % numMembers
-
-			// Prevent infinite loop
-			if memberIndex == startIndex && !assigned {
-				// Force assign to any subscribed member
-				for _, member := range sortedMembers {
-					subscribed := false
-					for _, topic := range member.Subscription {
-						if topic == partition.Topic {
-							subscribed = true
-							break
-						}
-					}
-					if subscribed {
-						assignments[member.ID] = append(assignments[member.ID], partition)
-						assigned = true
-						break
-					}
-				}
-				break
-			}
-		}
-	}
-
-	return assignments
-}
-
 // GetAssignmentStrategy returns the appropriate assignment strategy
 func GetAssignmentStrategy(name string) AssignmentStrategy {
 	switch name {
-	case "range":
+	case ProtocolNameRange:
 		return &RangeAssignmentStrategy{}
-	case "roundrobin":
+	case ProtocolNameRoundRobin:
 		return &RoundRobinAssignmentStrategy{}
-	case "cooperative-sticky":
-		return &CooperativeStickyAssignmentStrategy{}
-	case "incremental-cooperative":
+	case ProtocolNameCooperativeSticky:
 		return NewIncrementalCooperativeAssignmentStrategy()
 	default:
 		// Default to range strategy

weed/mq/kafka/consumer/assignment_test.go

@@ -9,8 +9,8 @@ import (
 func TestRangeAssignmentStrategy(t *testing.T) {
 	strategy := &RangeAssignmentStrategy{}
 
-	if strategy.Name() != "range" {
-		t.Errorf("Expected strategy name 'range', got '%s'", strategy.Name())
+	if strategy.Name() != ProtocolNameRange {
+		t.Errorf("Expected strategy name '%s', got '%s'", ProtocolNameRange, strategy.Name())
 	}
 
 	// Test with 2 members, 4 partitions on one topic
@@ -129,8 +129,8 @@ func TestRangeAssignmentStrategy_MultipleTopics(t *testing.T) {
 func TestRoundRobinAssignmentStrategy(t *testing.T) {
 	strategy := &RoundRobinAssignmentStrategy{}
 
-	if strategy.Name() != "roundrobin" {
-		t.Errorf("Expected strategy name 'roundrobin', got '%s'", strategy.Name())
+	if strategy.Name() != ProtocolNameRoundRobin {
+		t.Errorf("Expected strategy name '%s', got '%s'", ProtocolNameRoundRobin, strategy.Name())
 	}
 
 	// Test with 2 members, 4 partitions on one topic
@@ -206,19 +206,19 @@ func TestRoundRobinAssignmentStrategy_MultipleTopics(t *testing.T) {
 }
 
 func TestGetAssignmentStrategy(t *testing.T) {
-	rangeStrategy := GetAssignmentStrategy("range")
-	if rangeStrategy.Name() != "range" {
+	rangeStrategy := GetAssignmentStrategy(ProtocolNameRange)
+	if rangeStrategy.Name() != ProtocolNameRange {
 		t.Errorf("Expected range strategy, got %s", rangeStrategy.Name())
 	}
 
-	rrStrategy := GetAssignmentStrategy("roundrobin")
-	if rrStrategy.Name() != "roundrobin" {
+	rrStrategy := GetAssignmentStrategy(ProtocolNameRoundRobin)
+	if rrStrategy.Name() != ProtocolNameRoundRobin {
 		t.Errorf("Expected roundrobin strategy, got %s", rrStrategy.Name())
 	}
 
 	// Unknown strategy should default to range
 	defaultStrategy := GetAssignmentStrategy("unknown")
-	if defaultStrategy.Name() != "range" {
+	if defaultStrategy.Name() != ProtocolNameRange {
 		t.Errorf("Expected default strategy to be range, got %s", defaultStrategy.Name())
 	}
 }
@@ -226,7 +226,7 @@ func TestGetAssignmentStrategy(t *testing.T) {
 func TestConsumerGroup_AssignPartitions(t *testing.T) {
 	group := &ConsumerGroup{
 		ID:       "test-group",
-		Protocol: "range",
+		Protocol: ProtocolNameRange,
 		Members: map[string]*GroupMember{
 			"member1": {
 				ID:           "member1",

weed/mq/kafka/consumer/cooperative_sticky_test.go

@@ -5,14 +5,14 @@ import (
 )
 
 func TestCooperativeStickyAssignmentStrategy_Name(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
-	if strategy.Name() != "cooperative-sticky" {
-		t.Errorf("Expected strategy name 'cooperative-sticky', got '%s'", strategy.Name())
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
+	if strategy.Name() != ProtocolNameCooperativeSticky {
+		t.Errorf("Expected strategy name '%s', got '%s'", ProtocolNameCooperativeSticky, strategy.Name())
 	}
 }
 
 func TestCooperativeStickyAssignmentStrategy_InitialAssignment(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
 
 	members := []*GroupMember{
 		{ID: "member1", Subscription: []string{"topic1"}, Assignment: []PartitionAssignment{}},
@@ -55,12 +55,12 @@ func TestCooperativeStickyAssignmentStrategy_InitialAssignment(t *testing.T) {
 }
 
 func TestCooperativeStickyAssignmentStrategy_StickyBehavior(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
 
 	// Initial state: member1 has partitions 0,1 and member2 has partitions 2,3
 	members := []*GroupMember{
 		{
-			ID: "member1", 
+			ID:           "member1",
 			Subscription: []string{"topic1"},
 			Assignment: []PartitionAssignment{
 				{Topic: "topic1", Partition: 0},
@@ -68,7 +68,7 @@ func TestCooperativeStickyAssignmentStrategy_StickyBehavior(t *testing.T) {
 			},
 		},
 		{
-			ID: "member2", 
+			ID:           "member2",
 			Subscription: []string{"topic1"},
 			Assignment: []PartitionAssignment{
 				{Topic: "topic1", Partition: 2},
@@ -121,12 +121,12 @@ func TestCooperativeStickyAssignmentStrategy_StickyBehavior(t *testing.T) {
 }
 
 func TestCooperativeStickyAssignmentStrategy_NewMemberJoin(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
 
 	// Scenario: member1 has all partitions, member2 joins
 	members := []*GroupMember{
 		{
-			ID: "member1", 
+			ID:           "member1",
 			Subscription: []string{"topic1"},
 			Assignment: []PartitionAssignment{
 				{Topic: "topic1", Partition: 0},
@@ -136,9 +136,9 @@ func TestCooperativeStickyAssignmentStrategy_NewMemberJoin(t *testing.T) {
 			},
 		},
 		{
-			ID: "member2", 
+			ID:           "member2",
 			Subscription: []string{"topic1"},
-			Assignment: []PartitionAssignment{}, // New member, no existing assignment
+			Assignment:   []PartitionAssignment{}, // New member, no existing assignment
 		},
 	}
 
@@ -146,6 +146,17 @@ func TestCooperativeStickyAssignmentStrategy_NewMemberJoin(t *testing.T) {
 		"topic1": {0, 1, 2, 3},
 	}
 
+	// First call: revocation phase
+	assignments1 := strategy.Assign(members, topicPartitions)
+
+	// Update members with revocation results
+	members[0].Assignment = assignments1["member1"]
+	members[1].Assignment = assignments1["member2"]
+
+	// Force completion of revocation timeout
+	strategy.GetRebalanceState().RevocationTimeout = 0
+
+	// Second call: assignment phase
 	assignments := strategy.Assign(members, topicPartitions)
 
 	// Verify fair redistribution (2 partitions each)
@@ -177,12 +188,12 @@ func TestCooperativeStickyAssignmentStrategy_NewMemberJoin(t *testing.T) {
 }
 
 func TestCooperativeStickyAssignmentStrategy_MemberLeave(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
 
 	// Scenario: member2 leaves, member1 should get its partitions
 	members := []*GroupMember{
 		{
-			ID: "member1", 
+			ID:           "member1",
 			Subscription: []string{"topic1"},
 			Assignment: []PartitionAssignment{
 				{Topic: "topic1", Partition: 0},
@@ -223,11 +234,11 @@ func TestCooperativeStickyAssignmentStrategy_MemberLeave(t *testing.T) {
 }
 
 func TestCooperativeStickyAssignmentStrategy_MultipleTopics(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
 
 	members := []*GroupMember{
 		{
-			ID: "member1", 
+			ID:           "member1",
 			Subscription: []string{"topic1", "topic2"},
 			Assignment: []PartitionAssignment{
 				{Topic: "topic1", Partition: 0},
@@ -235,7 +246,7 @@ func TestCooperativeStickyAssignmentStrategy_MultipleTopics(t *testing.T) {
 			},
 		},
 		{
-			ID: "member2", 
+			ID:           "member2",
 			Subscription: []string{"topic1", "topic2"},
 			Assignment: []PartitionAssignment{
 				{Topic: "topic1", Partition: 1},
@@ -299,7 +310,7 @@ func TestCooperativeStickyAssignmentStrategy_MultipleTopics(t *testing.T) {
 }
 
 func TestCooperativeStickyAssignmentStrategy_UnevenPartitions(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
 
 	// 5 partitions, 2 members - should distribute 3:2 or 2:3
 	members := []*GroupMember{
@@ -334,7 +345,7 @@ func TestCooperativeStickyAssignmentStrategy_UnevenPartitions(t *testing.T) {
 }
 
 func TestCooperativeStickyAssignmentStrategy_PartialSubscription(t *testing.T) {
-	strategy := &CooperativeStickyAssignmentStrategy{}
+	strategy := NewIncrementalCooperativeAssignmentStrategy()
 
 	// member1 subscribes to both topics, member2 only to topic1
 	members := []*GroupMember{
@@ -393,20 +404,20 @@ func TestCooperativeStickyAssignmentStrategy_PartialSubscription(t *testing.T) {
 		}
 	}
 
-	if member1Topic1Count + member2Topic1Count != 2 {
+	if member1Topic1Count+member2Topic1Count != 2 {
 		t.Errorf("Expected all topic1 partitions to be assigned, got %d + %d = %d",
-			member1Topic1Count, member2Topic1Count, member1Topic1Count + member2Topic1Count)
+			member1Topic1Count, member2Topic1Count, member1Topic1Count+member2Topic1Count)
 	}
 }
 
 func TestGetAssignmentStrategy_CooperativeSticky(t *testing.T) {
-	strategy := GetAssignmentStrategy("cooperative-sticky")
-	if strategy.Name() != "cooperative-sticky" {
+	strategy := GetAssignmentStrategy(ProtocolNameCooperativeSticky)
+	if strategy.Name() != ProtocolNameCooperativeSticky {
 		t.Errorf("Expected cooperative-sticky strategy, got %s", strategy.Name())
 	}
 
 	// Verify it's the correct type
-	if _, ok := strategy.(*CooperativeStickyAssignmentStrategy); !ok {
-		t.Errorf("Expected CooperativeStickyAssignmentStrategy, got %T", strategy)
+	if _, ok := strategy.(*IncrementalCooperativeAssignmentStrategy); !ok {
+		t.Errorf("Expected IncrementalCooperativeAssignmentStrategy, got %T", strategy)
 	}
 }

weed/mq/kafka/consumer/incremental_rebalancing.go

@@ -31,8 +31,8 @@ func (rp RebalancePhase) String() string {
 // IncrementalRebalanceState tracks the state of incremental cooperative rebalancing
 type IncrementalRebalanceState struct {
 	Phase                RebalancePhase
-	RevocationGeneration int32 // Generation when revocation started
-	AssignmentGeneration int32 // Generation when assignment started
+	RevocationGeneration int32                            // Generation when revocation started
+	AssignmentGeneration int32                            // Generation when assignment started
 	RevokedPartitions    map[string][]PartitionAssignment // Member ID -> revoked partitions
 	PendingAssignments   map[string][]PartitionAssignment // Member ID -> pending assignments
 	StartTime            time.Time
@@ -64,7 +64,7 @@ func NewIncrementalCooperativeAssignmentStrategy() *IncrementalCooperativeAssign
 }
 
 func (ics *IncrementalCooperativeAssignmentStrategy) Name() string {
-	return "cooperative-sticky"
+	return ProtocolNameCooperativeSticky
 }
 
 func (ics *IncrementalCooperativeAssignmentStrategy) Assign(
@@ -334,9 +334,8 @@ func (ics *IncrementalCooperativeAssignmentStrategy) performRegularAssignment(
 	// Reset rebalance state
 	ics.rebalanceState = NewIncrementalRebalanceState()
 
-	// Use regular cooperative-sticky logic
-	cooperativeSticky := &CooperativeStickyAssignmentStrategy{}
-	return cooperativeSticky.Assign(members, topicPartitions)
+	// Use ideal assignment calculation (non-incremental cooperative assignment)
+	return ics.calculateIdealAssignment(members, topicPartitions)
 }
 
 // GetRebalanceState returns the current rebalance state (for monitoring/debugging)

weed/mq/kafka/protocol/consumer_group_metadata.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"net"
 	"sync"
+
+	"github.com/seaweedfs/seaweedfs/weed/mq/kafka/consumer"
 )
 
 // ConsumerProtocolMetadata represents parsed consumer protocol metadata
@@ -148,10 +150,10 @@ func ParseConsumerProtocolMetadata(metadata []byte, strategyName string) (*Consu
 // ValidateAssignmentStrategy checks if an assignment strategy is supported
 func ValidateAssignmentStrategy(strategy string) bool {
 	supportedStrategies := map[string]bool{
-		"range":              true,
-		"roundrobin":         true,
-		"sticky":             true,
-		"cooperative-sticky": false, // Not yet implemented
+		consumer.ProtocolNameRange:             true,
+		consumer.ProtocolNameRoundRobin:        true,
+		consumer.ProtocolNameSticky:            true,
+		consumer.ProtocolNameCooperativeSticky: true, // Incremental cooperative rebalancing (Kafka 2.4+)
 	}
 
 	return supportedStrategies[strategy]
@@ -184,7 +186,7 @@ func ExtractTopicsFromMetadata(protocols []GroupProtocol, fallbackTopics []strin
 // SelectBestProtocol chooses the best assignment protocol from available options
 func SelectBestProtocol(protocols []GroupProtocol, groupProtocols []string) string {
 	// Priority order: sticky > roundrobin > range
-	protocolPriority := []string{"sticky", "roundrobin", "range"}
+	protocolPriority := []string{consumer.ProtocolNameSticky, consumer.ProtocolNameRoundRobin, consumer.ProtocolNameRange}
 
 	// Find supported protocols in client's list
 	clientProtocols := make(map[string]bool)
@@ -218,8 +220,8 @@ func SelectBestProtocol(protocols []GroupProtocol, groupProtocols []string) stri
 
 		// No common protocol found - handle special fallback case
 		// If client supports nothing we validate, but group supports "range", use "range"
-		if len(clientProtocols) == 0 && groupProtocolSet["range"] {
-			return "range"
+		if len(clientProtocols) == 0 && groupProtocolSet[consumer.ProtocolNameRange] {
+			return consumer.ProtocolNameRange
 		}
 
 		// Return empty string to indicate no compatible protocol found
@@ -234,7 +236,7 @@ func SelectBestProtocol(protocols []GroupProtocol, groupProtocols []string) stri
 	}
 
 	// Last resort
-	return "range"
+	return consumer.ProtocolNameRange
 }
 
 // ProtocolMetadataDebugInfo returns debug information about protocol metadata

weed/mq/kafka/protocol/joingroup.go

@@ -232,7 +232,7 @@ func (h *Handler) handleJoinGroup(connContext *ConnectionContext, correlationID
 
 	// Ensure we have a valid protocol - fallback to "range" if empty
 	if groupProtocol == "" {
-		groupProtocol = "range"
+		groupProtocol = consumer.ProtocolNameRange
 	}
 
 	// If a protocol is already selected for the group, reject joins that do not support it.
@@ -615,7 +615,7 @@ func (h *Handler) buildJoinGroupResponse(response JoinGroupResponse) []byte {
 		} else {
 			// NON-nullable compact string in v6 - must not be empty!
 			if response.ProtocolName == "" {
-				response.ProtocolName = "range" // fallback to default
+				response.ProtocolName = consumer.ProtocolNameRange // fallback to default
 			}
 			out = append(out, FlexibleString(response.ProtocolName)...)
 		}
@@ -762,9 +762,9 @@ func (h *Handler) buildJoinGroupErrorResponse(correlationID uint32, errorCode in
 		ThrottleTimeMs: 0,
 		ErrorCode:      errorCode,
 		GenerationID:   -1,
-		ProtocolName:   "range",   // Use "range" as default protocol instead of empty string
-		Leader:         "unknown", // Use "unknown" instead of empty string for non-nullable field
-		MemberID:       "unknown", // Use "unknown" instead of empty string for non-nullable field
+		ProtocolName:   consumer.ProtocolNameRange, // Use "range" as default protocol instead of empty string
+		Leader:         "unknown",                  // Use "unknown" instead of empty string for non-nullable field
+		MemberID:       "unknown",                  // Use "unknown" instead of empty string for non-nullable field
 		Version:        apiVersion,
 		Members:        []JoinGroupMember{},
 	}