subscriber can get assignments

weed/mq/sub_coordinator/consumer_group.go

@@ -2,6 +2,7 @@ package sub_coordinator
 
 import (
 	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"github.com/seaweedfs/seaweedfs/weed/mq/pub_balancer"
 	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
@@ -15,16 +16,20 @@ type ConsumerGroupInstance struct {
 	ResponseChan chan *mq_pb.SubscriberToSubCoordinatorResponse
 }
 type ConsumerGroup struct {
+	topic         topic.Topic
 	// map a consumer group instance id to a consumer group instance
 	ConsumerGroupInstances cmap.ConcurrentMap[string, *ConsumerGroupInstance]
 	mapping        *PartitionConsumerMapping
 	reBalanceTimer *time.Timer
+	pubBalancer    *pub_balancer.Balancer
 }
 
-func NewConsumerGroup() *ConsumerGroup {
+func NewConsumerGroup(t *mq_pb.Topic, pubBalancer *pub_balancer.Balancer) *ConsumerGroup {
 	return &ConsumerGroup{
+		topic:                  topic.FromPbTopic(t),
 		ConsumerGroupInstances: cmap.New[*ConsumerGroupInstance](),
 		mapping:                NewPartitionConsumerMapping(pub_balancer.MaxPartitionCount),
+		pubBalancer:            pubBalancer,
 	}
 }
 
@@ -47,7 +52,7 @@ func (cg *ConsumerGroup) onConsumerGroupInstanceChange(){
 		cg.reBalanceTimer = nil
 	}
 	cg.reBalanceTimer = time.AfterFunc(5*time.Second, func() {
-		cg.Rebalance()
+		cg.RebalanceConsumberGroupInstances()
 		cg.reBalanceTimer = nil
 	})
 }
@@ -56,9 +61,66 @@ func (cg *ConsumerGroup) OnPartitionListChange() {
 		cg.reBalanceTimer.Stop()
 		cg.reBalanceTimer = nil
 	}
-	cg.Rebalance()
+	cg.RebalanceConsumberGroupInstances()
 }
 
-func (cg *ConsumerGroup) Rebalance() {
+func (cg *ConsumerGroup) RebalanceConsumberGroupInstances() {
 	println("rebalance...")
+
+	now := time.Now().UnixNano()
+
+	// collect current topic partitions
+	partitionSlotToBrokerList, found := cg.pubBalancer.TopicToBrokers.Get(cg.topic.String())
+	if !found {
+		glog.V(0).Infof("topic %s not found in balancer", cg.topic.String())
+		return
+	}
+	partitions := make([]*topic.Partition, 0)
+	for _, partitionSlot := range partitionSlotToBrokerList.PartitionSlots {
+		partitions = append(partitions, topic.NewPartition(partitionSlot.RangeStart, partitionSlot.RangeStop, partitionSlotToBrokerList.RingSize, now))
+	}
+
+	// collect current consumer group instance ids
+	consumerInstanceIds := make([]string, 0)
+	for _, consumerGroupInstance := range cg.ConsumerGroupInstances.Items() {
+		consumerInstanceIds = append(consumerInstanceIds, consumerGroupInstance.InstanceId)
+	}
+
+	cg.mapping.BalanceToConsumerInstanceIds(partitions, consumerInstanceIds)
+
+	// convert cg.mapping currentMapping to map of consumer group instance id to partition slots
+	consumerInstanceToPartitionSlots := make(map[string][]*PartitionSlotToConsumerInstance)
+	for _, partitionSlot := range cg.mapping.currentMapping.PartitionSlots {
+		consumerInstanceToPartitionSlots[partitionSlot.AssignedInstanceId] = append(consumerInstanceToPartitionSlots[partitionSlot.AssignedInstanceId], partitionSlot)
+	}
+
+	// notify consumer group instances
+	for _, consumerGroupInstance := range cg.ConsumerGroupInstances.Items() {
+		partitionSlots, found := consumerInstanceToPartitionSlots[consumerGroupInstance.InstanceId]
+		if !found {
+			partitionSlots = make([]*PartitionSlotToConsumerInstance, 0)
+		}
+		consumerGroupInstance.Partitions = ToPartitions(partitionSlotToBrokerList.RingSize, partitionSlots, now)
+		assignedPartitions := make([]*mq_pb.SubscriberToSubCoordinatorResponse_AssignedPartition, len(partitionSlots))
+		for i, partitionSlot := range partitionSlots {
+			assignedPartitions[i] = &mq_pb.SubscriberToSubCoordinatorResponse_AssignedPartition{
+				Partition: &mq_pb.Partition{
+					RangeStop: partitionSlot.RangeStop,
+					RangeStart: partitionSlot.RangeStart,
+					RingSize: partitionSlotToBrokerList.RingSize,
+					UnixTimeNs: now,
+				},
+			}
+		}
+		response := &mq_pb.SubscriberToSubCoordinatorResponse{
+			Message: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment_{
+				Assignment: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment{
+					AssignedPartitions: assignedPartitions,
+				},
+			},
+		}
+		consumerGroupInstance.ResponseChan <- response
+	}
+
+
 }

weed/mq/sub_coordinator/coordinator.go

@@ -53,7 +53,7 @@ func (c *Coordinator) AddSubscriber(consumerGroup, consumerGroupInstance string,
 	tcg := c.GetTopicConsumerGroups(topic)
 	cg, _ := tcg.ConsumerGroups.Get(consumerGroup)
 	if cg == nil {
-		cg = NewConsumerGroup()
+		cg = NewConsumerGroup(topic, c.balancer)
 		tcg.ConsumerGroups.Set(consumerGroup, cg)
 	}
 	cgi, _ := cg.ConsumerGroupInstances.Get(consumerGroupInstance)

weed/mq/sub_coordinator/partition_consumer_mapping.go

@@ -29,9 +29,18 @@ func (pcm *PartitionConsumerMapping) BalanceToConsumerInstanceIds(partitions []*
 	}
 	newVersion := time.Now().UnixNano()
 	newMapping := NewPartitionSlotToConsumerInstanceList(partitions[0].RingSize, newVersion)
-	newMapping.PartitionSlots = doBalanceSticky(partitions, consumerInstanceIds, pcm.prevMappings[0])
+	var prevMapping *PartitionSlotToConsumerInstanceList
+	if len(pcm.prevMappings) > 0 {
+		prevMapping = pcm.prevMappings[len(pcm.prevMappings)-1]
+	} else {
+		prevMapping = nil
+	}
+	newMapping.PartitionSlots = doBalanceSticky(partitions, consumerInstanceIds, prevMapping)
 	if pcm.currentMapping != nil {
 		pcm.prevMappings = append(pcm.prevMappings, pcm.currentMapping)
+		if len(pcm.prevMappings) > 10 {
+			pcm.prevMappings = pcm.prevMappings[1:]
+		}
 	}
 	pcm.currentMapping = newMapping
 }

weed/mq/sub_coordinator/partition_list.go

@@ -30,3 +30,11 @@ func ToPartitionSlots(partitions []*topic.Partition) (partitionSlots []*Partitio
 	}
 	return
 }
+
+func ToPartitions(ringSize int32, slots []*PartitionSlotToConsumerInstance, unixTimeNs int64) []*topic.Partition {
+	partitions := make([]*topic.Partition, 0, len(slots))
+	for _, slot := range slots {
+		partitions = append(partitions, topic.NewPartition(slot.RangeStart, slot.RangeStop, ringSize, unixTimeNs))
+	}
+	return partitions
+}