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

weed/mq/sub_coordinator/consumer_group.go

@@ -0,0 +1,113 @@
+package sub_coordinator
+
+import (
+	"fmt"
+	"time"
+
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/filer_client"
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+type ConsumerGroup struct {
+	topic topic.Topic
+	// map a consumer group instance id to a consumer group instance
+	ConsumerGroupInstances cmap.ConcurrentMap[string, *ConsumerGroupInstance]
+	Market                 *Market
+	reBalanceTimer         *time.Timer
+	filerClientAccessor    *filer_client.FilerClientAccessor
+	stopCh                 chan struct{}
+}
+
+func NewConsumerGroup(t *schema_pb.Topic, reblanceSeconds int32, filerClientAccessor *filer_client.FilerClientAccessor) *ConsumerGroup {
+	cg := &ConsumerGroup{
+		topic:                  topic.FromPbTopic(t),
+		ConsumerGroupInstances: cmap.New[*ConsumerGroupInstance](),
+		filerClientAccessor:    filerClientAccessor,
+		stopCh:                 make(chan struct{}),
+	}
+	if conf, err := cg.filerClientAccessor.ReadTopicConfFromFiler(cg.topic); err == nil {
+		var partitions []topic.Partition
+		for _, assignment := range conf.BrokerPartitionAssignments {
+			partitions = append(partitions, topic.FromPbPartition(assignment.Partition))
+		}
+		cg.Market = NewMarket(partitions, time.Duration(reblanceSeconds)*time.Second)
+	} else {
+		glog.V(0).Infof("fail to read topic conf from filer: %v", err)
+		return nil
+	}
+
+	go func() {
+		for {
+			select {
+			case adjustment := <-cg.Market.AdjustmentChan:
+				cgi, found := cg.ConsumerGroupInstances.Get(string(adjustment.consumer))
+				if !found {
+					glog.V(0).Infof("consumer group instance %s not found", adjustment.consumer)
+					continue
+				}
+				if adjustment.isAssign {
+					if conf, err := cg.filerClientAccessor.ReadTopicConfFromFiler(cg.topic); err == nil {
+						for _, assignment := range conf.BrokerPartitionAssignments {
+							if adjustment.partition.Equals(topic.FromPbPartition(assignment.Partition)) {
+								cgi.ResponseChan <- &mq_pb.SubscriberToSubCoordinatorResponse{
+									Message: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment_{
+										Assignment: &mq_pb.SubscriberToSubCoordinatorResponse_Assignment{
+											PartitionAssignment: &mq_pb.BrokerPartitionAssignment{
+												Partition:      adjustment.partition.ToPbPartition(),
+												LeaderBroker:   assignment.LeaderBroker,
+												FollowerBroker: assignment.FollowerBroker,
+											},
+										},
+									},
+								}
+								glog.V(0).Infof("send assignment %v to %s", adjustment.partition, adjustment.consumer)
+								break
+							}
+						}
+					}
+				} else {
+					cgi.ResponseChan <- &mq_pb.SubscriberToSubCoordinatorResponse{
+						Message: &mq_pb.SubscriberToSubCoordinatorResponse_UnAssignment_{
+							UnAssignment: &mq_pb.SubscriberToSubCoordinatorResponse_UnAssignment{
+								Partition: adjustment.partition.ToPbPartition(),
+							},
+						},
+					}
+					glog.V(0).Infof("send unassignment %v to %s", adjustment.partition, adjustment.consumer)
+				}
+			case <-cg.stopCh:
+				return
+			}
+		}
+	}()
+
+	return cg
+}
+
+func (cg *ConsumerGroup) AckAssignment(cgi *ConsumerGroupInstance, assignment *mq_pb.SubscriberToSubCoordinatorRequest_AckAssignmentMessage) {
+	fmt.Printf("ack assignment %v\n", assignment)
+	cg.Market.ConfirmAdjustment(&Adjustment{
+		consumer:  cgi.InstanceId,
+		partition: topic.FromPbPartition(assignment.Partition),
+		isAssign:  true,
+	})
+}
+func (cg *ConsumerGroup) AckUnAssignment(cgi *ConsumerGroupInstance, assignment *mq_pb.SubscriberToSubCoordinatorRequest_AckUnAssignmentMessage) {
+	fmt.Printf("ack unassignment %v\n", assignment)
+	cg.Market.ConfirmAdjustment(&Adjustment{
+		consumer:  cgi.InstanceId,
+		partition: topic.FromPbPartition(assignment.Partition),
+		isAssign:  false,
+	})
+}
+
+func (cg *ConsumerGroup) OnPartitionListChange(assignments []*mq_pb.BrokerPartitionAssignment) {
+}
+
+func (cg *ConsumerGroup) Shutdown() {
+	close(cg.stopCh)
+}

weed/mq/sub_coordinator/consumer_group_instance.go

@@ -0,0 +1,23 @@
+package sub_coordinator
+
+import (
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+)
+
+type ConsumerGroupInstanceId string
+
+type ConsumerGroupInstance struct {
+	InstanceId         ConsumerGroupInstanceId
+	AssignedPartitions []topic.Partition
+	ResponseChan       chan *mq_pb.SubscriberToSubCoordinatorResponse
+	MaxPartitionCount  int32
+}
+
+func NewConsumerGroupInstance(instanceId string, maxPartitionCount int32) *ConsumerGroupInstance {
+	return &ConsumerGroupInstance{
+		InstanceId:        ConsumerGroupInstanceId(instanceId),
+		ResponseChan:      make(chan *mq_pb.SubscriberToSubCoordinatorResponse, 1),
+		MaxPartitionCount: maxPartitionCount,
+	}
+}

weed/mq/sub_coordinator/inflight_message_tracker.go

@@ -0,0 +1,144 @@
+package sub_coordinator
+
+import (
+	"sort"
+	"sync"
+)
+
+type InflightMessageTracker struct {
+	messages   map[string]int64
+	mu         sync.Mutex
+	timestamps *RingBuffer
+}
+
+func NewInflightMessageTracker(capacity int) *InflightMessageTracker {
+	return &InflightMessageTracker{
+		messages:   make(map[string]int64),
+		timestamps: NewRingBuffer(capacity),
+	}
+}
+
+// EnflightMessage tracks the message with the key and timestamp.
+// These messages are sent to the consumer group instances and waiting for ack.
+func (imt *InflightMessageTracker) EnflightMessage(key []byte, tsNs int64) {
+	// fmt.Printf("EnflightMessage(%s,%d)\n", string(key), tsNs)
+	imt.mu.Lock()
+	defer imt.mu.Unlock()
+	imt.messages[string(key)] = tsNs
+	imt.timestamps.EnflightTimestamp(tsNs)
+}
+
+// AcknowledgeMessage acknowledges the message with the key and timestamp.
+func (imt *InflightMessageTracker) AcknowledgeMessage(key []byte, tsNs int64) bool {
+	// fmt.Printf("AcknowledgeMessage(%s,%d)\n", string(key), tsNs)
+	imt.mu.Lock()
+	defer imt.mu.Unlock()
+	timestamp, exists := imt.messages[string(key)]
+	if !exists || timestamp != tsNs {
+		return false
+	}
+	delete(imt.messages, string(key))
+	// Remove the specific timestamp from the ring buffer.
+	imt.timestamps.AckTimestamp(tsNs)
+	return true
+}
+
+func (imt *InflightMessageTracker) GetOldestAckedTimestamp() int64 {
+	return imt.timestamps.OldestAckedTimestamp()
+}
+
+// IsInflight returns true if the message with the key is inflight.
+func (imt *InflightMessageTracker) IsInflight(key []byte) bool {
+	imt.mu.Lock()
+	defer imt.mu.Unlock()
+	_, found := imt.messages[string(key)]
+	return found
+}
+
+// Cleanup clears all in-flight messages. This should be called when a subscriber disconnects
+// to prevent messages from being stuck in the in-flight state indefinitely.
+func (imt *InflightMessageTracker) Cleanup() int {
+	imt.mu.Lock()
+	defer imt.mu.Unlock()
+	count := len(imt.messages)
+	// Clear all in-flight messages
+	imt.messages = make(map[string]int64)
+	return count
+}
+
+type TimestampStatus struct {
+	Timestamp int64
+	Acked     bool
+}
+
+// RingBuffer represents a circular buffer to hold timestamps.
+type RingBuffer struct {
+	buffer        []*TimestampStatus
+	head          int
+	size          int
+	maxTimestamp  int64
+	maxAllAckedTs int64
+}
+
+// NewRingBuffer creates a new RingBuffer of the given capacity.
+func NewRingBuffer(capacity int) *RingBuffer {
+	return &RingBuffer{
+		buffer: newBuffer(capacity),
+	}
+}
+
+func newBuffer(capacity int) []*TimestampStatus {
+	buffer := make([]*TimestampStatus, capacity)
+	for i := range buffer {
+		buffer[i] = &TimestampStatus{}
+	}
+	return buffer
+}
+
+// EnflightTimestamp adds a new timestamp to the ring buffer.
+func (rb *RingBuffer) EnflightTimestamp(timestamp int64) {
+	if rb.size < len(rb.buffer) {
+		rb.size++
+	} else {
+		newBuf := newBuffer(2 * len(rb.buffer))
+		for i := 0; i < rb.size; i++ {
+			newBuf[i] = rb.buffer[(rb.head+len(rb.buffer)-rb.size+i)%len(rb.buffer)]
+		}
+		rb.buffer = newBuf
+		rb.head = rb.size
+		rb.size++
+	}
+	head := rb.buffer[rb.head]
+	head.Timestamp = timestamp
+	head.Acked = false
+	rb.head = (rb.head + 1) % len(rb.buffer)
+	if timestamp > rb.maxTimestamp {
+		rb.maxTimestamp = timestamp
+	}
+}
+
+// AckTimestamp removes the specified timestamp from the ring buffer.
+func (rb *RingBuffer) AckTimestamp(timestamp int64) {
+	// Perform binary search
+	index := sort.Search(rb.size, func(i int) bool {
+		return rb.buffer[(rb.head+len(rb.buffer)-rb.size+i)%len(rb.buffer)].Timestamp >= timestamp
+	})
+	actualIndex := (rb.head + len(rb.buffer) - rb.size + index) % len(rb.buffer)
+
+	rb.buffer[actualIndex].Acked = true
+
+	// Remove all the continuously acknowledged timestamps from the buffer
+	startPos := (rb.head + len(rb.buffer) - rb.size) % len(rb.buffer)
+	for i := 0; i < len(rb.buffer) && rb.buffer[(startPos+i)%len(rb.buffer)].Acked; i++ {
+		t := rb.buffer[(startPos+i)%len(rb.buffer)]
+		if rb.maxAllAckedTs < t.Timestamp {
+			rb.size--
+			rb.maxAllAckedTs = t.Timestamp
+		}
+	}
+}
+
+// OldestAckedTimestamp returns the oldest that is already acked timestamp in the ring buffer.
+func (rb *RingBuffer) OldestAckedTimestamp() int64 {
+	return rb.maxAllAckedTs
+}

weed/mq/sub_coordinator/market.go

@@ -0,0 +1,368 @@
+package sub_coordinator
+
+import (
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+)
+
+/*
+Market is a data structure that keeps track of the state of the consumer group instances and the partitions.
+
+When rebalancing, the market will try to balance the load of the partitions among the consumer group instances.
+For each loop, the market will:
+* If a consumer group instance has more partitions than the average, it will unassign some partitions.
+* If a consumer group instance has less partitions than the average, it will assign some partitions.
+
+Trigger rebalance when:
+* A new consumer group instance is added
+* Some partitions are unassigned from a consumer group instance.
+
+If multiple rebalance requests are received, after a certain period, the market will only process the latest request.
+
+However, if the number of unassigned partition is increased to exactly the total number of partitions,
+and total partitions are less than or equal to the sum of the max partition count of all consumer group instances,
+the market will process the request immediately.
+This is to ensure a partition can be migrated to another consumer group instance as soon as possible.
+
+Emit these adjustments to the subscriber coordinator:
+* Assign a partition to a consumer group instance
+* Unassign a partition from a consumer group instance
+
+Because the adjustment is sent to the subscriber coordinator, the market will keep track of the inflight adjustments.
+The subscriber coordinator will send back the response to the market when the adjustment is processed.
+If the adjustment is older than a certain time(inflightAdjustmentTTL), it would be considered expired.
+Otherwise, the adjustment is considered inflight, so it would be used when calculating the load.
+
+Later features:
+* A consumer group instance is not keeping up with the load.
+
+Since a coordinator, and thus the market, may be restarted or moved to another node, the market should be able to recover the state from the subscriber coordinator.
+The subscriber coordinator should be able to send the current state of the consumer group instances and the partitions to the market.
+
+*/
+
+type PartitionSlot struct {
+	Partition  topic.Partition
+	AssignedTo *ConsumerGroupInstance // Track the consumer assigned to this partition slot
+}
+
+type Adjustment struct {
+	isAssign  bool
+	partition topic.Partition
+	consumer  ConsumerGroupInstanceId
+	ts        time.Time
+}
+
+type Market struct {
+	mu                    sync.Mutex
+	partitions            map[topic.Partition]*PartitionSlot
+	consumerInstances     map[ConsumerGroupInstanceId]*ConsumerGroupInstance
+	AdjustmentChan        chan *Adjustment
+	inflightAdjustments   []*Adjustment
+	inflightAdjustmentTTL time.Duration
+	lastBalancedTime      time.Time
+	stopChan              chan struct{}
+	balanceRequestChan    chan struct{}
+	hasBalanceRequest     bool
+}
+
+func NewMarket(partitions []topic.Partition, inflightAdjustmentTTL time.Duration) *Market {
+	partitionMap := make(map[topic.Partition]*PartitionSlot)
+	for _, partition := range partitions {
+		partitionMap[partition] = &PartitionSlot{
+			Partition: partition,
+		}
+	}
+	m := &Market{
+		partitions:            partitionMap,
+		consumerInstances:     make(map[ConsumerGroupInstanceId]*ConsumerGroupInstance),
+		AdjustmentChan:        make(chan *Adjustment, 100),
+		inflightAdjustmentTTL: inflightAdjustmentTTL,
+		stopChan:              make(chan struct{}),
+		balanceRequestChan:    make(chan struct{}),
+	}
+	m.lastBalancedTime = time.Now()
+	go m.loopBalanceLoad()
+
+	return m
+}
+
+func (m *Market) ShutdownMarket() {
+	close(m.stopChan)
+	close(m.AdjustmentChan)
+}
+
+func (m *Market) AddConsumerInstance(consumer *ConsumerGroupInstance) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if _, exists := m.consumerInstances[consumer.InstanceId]; exists {
+		return errors.New("consumer instance already exists")
+	}
+
+	m.consumerInstances[consumer.InstanceId] = consumer
+	m.balanceRequestChan <- struct{}{}
+
+	return nil
+}
+
+func (m *Market) RemoveConsumerInstance(consumerId ConsumerGroupInstanceId) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	consumer, exists := m.consumerInstances[consumerId]
+	if !exists {
+		return nil
+	}
+	delete(m.consumerInstances, consumerId)
+
+	for _, partition := range consumer.AssignedPartitions {
+		if partitionSlot, exists := m.partitions[partition]; exists {
+			partitionSlot.AssignedTo = nil
+		}
+	}
+	m.balanceRequestChan <- struct{}{}
+
+	return nil
+}
+
+func (m *Market) assignPartitionToConsumer(partition *PartitionSlot) {
+	var bestConsumer *ConsumerGroupInstance
+	var minLoad = int(^uint(0) >> 1) // Max int value
+
+	inflightConsumerAdjustments := make(map[ConsumerGroupInstanceId]int)
+	for _, adjustment := range m.inflightAdjustments {
+		if adjustment.isAssign {
+			inflightConsumerAdjustments[adjustment.consumer]++
+		} else {
+			inflightConsumerAdjustments[adjustment.consumer]--
+		}
+	}
+	for _, consumer := range m.consumerInstances {
+		consumerLoad := len(consumer.AssignedPartitions)
+		if inflightAdjustments, exists := inflightConsumerAdjustments[consumer.InstanceId]; exists {
+			consumerLoad += inflightAdjustments
+		}
+		// fmt.Printf("Consumer %+v has load %d, max %d, min %d\n", consumer.InstanceId, consumerLoad, consumer.MaxPartitionCount, minLoad)
+		if consumerLoad < int(consumer.MaxPartitionCount) {
+			if consumerLoad < minLoad {
+				bestConsumer = consumer
+				minLoad = consumerLoad
+				// fmt.Printf("picked: Consumer %+v has load %d, max %d, min %d\n", consumer.InstanceId, consumerLoad, consumer.MaxPartitionCount, minLoad)
+			}
+		}
+	}
+
+	if bestConsumer != nil {
+		// change consumer assigned partitions later when the adjustment is confirmed
+		adjustment := &Adjustment{
+			isAssign:  true,
+			partition: partition.Partition,
+			consumer:  bestConsumer.InstanceId,
+			ts:        time.Now(),
+		}
+		m.AdjustmentChan <- adjustment
+		m.inflightAdjustments = append(m.inflightAdjustments, adjustment)
+		m.lastBalancedTime = adjustment.ts
+	}
+}
+
+func (m *Market) loopBalanceLoad() {
+	ticker := time.NewTicker(500 * time.Millisecond)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			if m.hasBalanceRequest {
+				m.hasBalanceRequest = false
+				inflightAdjustments := make([]*Adjustment, 0, len(m.inflightAdjustments))
+				for _, adjustment := range m.inflightAdjustments {
+					if adjustment.ts.Add(m.inflightAdjustmentTTL).After(time.Now()) {
+						inflightAdjustments = append(inflightAdjustments, adjustment)
+					}
+				}
+				m.inflightAdjustments = inflightAdjustments
+
+				m.doBalanceLoad()
+				// println("Balance load completed.")
+				m.Status()
+			}
+		case <-m.balanceRequestChan:
+			m.hasBalanceRequest = true
+		case <-m.stopChan:
+			return
+		}
+	}
+}
+
+// doBalanceLoad will balance the load of the partitions among the consumer group instances.
+// It will try to unassign partitions from the consumer group instances that have more partitions than the average.
+// It will try to assign partitions to the consumer group instances that have less partitions than the average.
+func (m *Market) doBalanceLoad() {
+	if len(m.consumerInstances) == 0 {
+		return
+	}
+
+	// find the average load for all consumers
+	averageLoad := m.findAverageLoad()
+
+	// find the consumers with the higher load than average
+	if m.adjustBusyConsumers(averageLoad) {
+		return
+	}
+
+	// find partitions with no consumer assigned
+	m.adjustUnassignedPartitions()
+}
+func (m *Market) findAverageLoad() (averageLoad float32) {
+	var totalLoad int
+	for _, consumer := range m.consumerInstances {
+		totalLoad += len(consumer.AssignedPartitions)
+	}
+	for _, adjustment := range m.inflightAdjustments {
+		if adjustment.isAssign {
+			totalLoad++
+		} else {
+			totalLoad--
+		}
+	}
+	averageLoad = float32(totalLoad) / float32(len(m.consumerInstances))
+	return
+}
+
+func (m *Market) adjustBusyConsumers(averageLoad float32) (hasAdjustments bool) {
+	inflightConsumerAdjustments := make(map[ConsumerGroupInstanceId]int)
+	for _, adjustment := range m.inflightAdjustments {
+		if adjustment.isAssign {
+			inflightConsumerAdjustments[adjustment.consumer]++
+		} else {
+			inflightConsumerAdjustments[adjustment.consumer]--
+		}
+	}
+	for _, consumer := range m.consumerInstances {
+		consumerLoad := len(consumer.AssignedPartitions)
+		if inflightAdjustment, exists := inflightConsumerAdjustments[consumer.InstanceId]; exists {
+			consumerLoad += inflightAdjustment
+		}
+		delta := int(float32(consumerLoad) - averageLoad)
+		if delta <= 0 {
+			continue
+		}
+		adjustTime := time.Now()
+		for i := 0; i < delta; i++ {
+			adjustment := &Adjustment{
+				isAssign:  false,
+				partition: consumer.AssignedPartitions[i],
+				consumer:  consumer.InstanceId,
+				ts:        adjustTime,
+			}
+			m.AdjustmentChan <- adjustment
+			m.inflightAdjustments = append(m.inflightAdjustments, adjustment)
+			m.lastBalancedTime = adjustment.ts
+		}
+		hasAdjustments = true
+	}
+	return
+}
+
+func (m *Market) adjustUnassignedPartitions() {
+	inflightPartitionAdjustments := make(map[topic.Partition]bool)
+	for _, adjustment := range m.inflightAdjustments {
+		inflightPartitionAdjustments[adjustment.partition] = true
+	}
+	for _, partitionSlot := range m.partitions {
+		if partitionSlot.AssignedTo == nil {
+			if _, exists := inflightPartitionAdjustments[partitionSlot.Partition]; exists {
+				continue
+			}
+			// fmt.Printf("Assigning partition %+v to consumer\n", partitionSlot.Partition)
+			m.assignPartitionToConsumer(partitionSlot)
+		}
+	}
+}
+
+func (m *Market) ConfirmAdjustment(adjustment *Adjustment) {
+	if adjustment.isAssign {
+		m.confirmAssignPartition(adjustment.partition, adjustment.consumer)
+	} else {
+		m.unassignPartitionSlot(adjustment.partition)
+	}
+	glog.V(1).Infof("ConfirmAdjustment %+v", adjustment)
+	m.Status()
+}
+
+func (m *Market) unassignPartitionSlot(partition topic.Partition) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	partitionSlot, exists := m.partitions[partition]
+	if !exists {
+		glog.V(0).Infof("partition %+v slot is not tracked", partition)
+		return
+	}
+
+	if partitionSlot.AssignedTo == nil {
+		glog.V(0).Infof("partition %+v slot is not assigned to any consumer", partition)
+		return
+	}
+
+	consumer := partitionSlot.AssignedTo
+	for i, p := range consumer.AssignedPartitions {
+		if p == partition {
+			consumer.AssignedPartitions = append(consumer.AssignedPartitions[:i], consumer.AssignedPartitions[i+1:]...)
+			partitionSlot.AssignedTo = nil
+			m.balanceRequestChan <- struct{}{}
+			return
+		}
+	}
+
+	glog.V(0).Infof("partition %+v slot not found in assigned consumer", partition)
+
+}
+
+func (m *Market) confirmAssignPartition(partition topic.Partition, consumerInstanceId ConsumerGroupInstanceId) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	partitionSlot, exists := m.partitions[partition]
+	if !exists {
+		glog.V(0).Infof("partition %+v slot is not tracked", partition)
+		return
+	}
+
+	if partitionSlot.AssignedTo != nil {
+		glog.V(0).Infof("partition %+v slot is already assigned to %+v", partition, partitionSlot.AssignedTo.InstanceId)
+		return
+	}
+
+	consumerInstance, exists := m.consumerInstances[consumerInstanceId]
+	if !exists {
+		glog.V(0).Infof("consumer %+v is not tracked", consumerInstanceId)
+		return
+	}
+
+	partitionSlot.AssignedTo = consumerInstance
+	consumerInstance.AssignedPartitions = append(consumerInstance.AssignedPartitions, partition)
+
+}
+
+func (m *Market) Status() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	glog.V(1).Infof("Market has %d partitions and %d consumer instances", len(m.partitions), len(m.consumerInstances))
+	for partition, slot := range m.partitions {
+		if slot.AssignedTo == nil {
+			glog.V(1).Infof("Partition %+v is not assigned to any consumer", partition)
+		} else {
+			glog.V(1).Infof("Partition %+v is assigned to consumer %+v", partition, slot.AssignedTo.InstanceId)
+		}
+	}
+	for _, consumer := range m.consumerInstances {
+		glog.V(1).Infof("Consumer %+v has %d partitions", consumer.InstanceId, len(consumer.AssignedPartitions))
+	}
+}

weed/mq/sub_coordinator/market_test.go

@@ -0,0 +1,103 @@
+package sub_coordinator
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/stretchr/testify/assert"
+)
+
+var partitions = []topic.Partition{
+	{
+		RangeStart: 0,
+		RangeStop:  1,
+		RingSize:   3,
+		UnixTimeNs: 0,
+	},
+	{
+		RangeStart: 1,
+		RangeStop:  2,
+		RingSize:   3,
+		UnixTimeNs: 0,
+	},
+	{
+		RangeStart: 2,
+		RangeStop:  3,
+		RingSize:   3,
+		UnixTimeNs: 0,
+	},
+}
+
+func TestAddConsumerInstance(t *testing.T) {
+	market := NewMarket(partitions, 10*time.Second)
+
+	consumer := &ConsumerGroupInstance{
+		InstanceId:        "first",
+		MaxPartitionCount: 2,
+	}
+	err := market.AddConsumerInstance(consumer)
+
+	assert.Nil(t, err)
+	time.Sleep(1 * time.Second) // Allow time for background rebalancing
+	market.ShutdownMarket()
+	for adjustment := range market.AdjustmentChan {
+		fmt.Printf("%+v\n", adjustment)
+	}
+}
+
+func TestMultipleConsumerInstances(t *testing.T) {
+	market := NewMarket(partitions, 10*time.Second)
+
+	market.AddConsumerInstance(&ConsumerGroupInstance{
+		InstanceId:        "first",
+		MaxPartitionCount: 2,
+	})
+	market.AddConsumerInstance(&ConsumerGroupInstance{
+		InstanceId:        "second",
+		MaxPartitionCount: 2,
+	})
+	market.AddConsumerInstance(&ConsumerGroupInstance{
+		InstanceId:        "third",
+		MaxPartitionCount: 2,
+	})
+
+	time.Sleep(1 * time.Second) // Allow time for background rebalancing
+	market.ShutdownMarket()
+	for adjustment := range market.AdjustmentChan {
+		fmt.Printf("%+v\n", adjustment)
+	}
+}
+
+func TestConfirmAdjustment(t *testing.T) {
+	market := NewMarket(partitions, 1*time.Second)
+
+	market.AddConsumerInstance(&ConsumerGroupInstance{
+		InstanceId:        "first",
+		MaxPartitionCount: 2,
+	})
+	market.AddConsumerInstance(&ConsumerGroupInstance{
+		InstanceId:        "second",
+		MaxPartitionCount: 2,
+	})
+	market.AddConsumerInstance(&ConsumerGroupInstance{
+		InstanceId:        "third",
+		MaxPartitionCount: 2,
+	})
+
+	go func() {
+		time.Sleep(5 * time.Second) // Allow time for background rebalancing
+		market.ShutdownMarket()
+	}()
+	go func() {
+		time.Sleep(2 * time.Second)
+		market.RemoveConsumerInstance("third")
+	}()
+
+	for adjustment := range market.AdjustmentChan {
+		fmt.Printf("%+v\n", adjustment)
+		market.ConfirmAdjustment(adjustment)
+	}
+
+}

weed/mq/sub_coordinator/partition_consumer_mapping.go

@@ -0,0 +1,6 @@
+package sub_coordinator
+
+type PartitionConsumerMapping struct {
+	currentMapping *PartitionSlotToConsumerInstanceList
+	prevMappings   []*PartitionSlotToConsumerInstanceList
+}

weed/mq/sub_coordinator/partition_list.go

@@ -0,0 +1,16 @@
+package sub_coordinator
+
+type PartitionSlotToConsumerInstance struct {
+	RangeStart         int32
+	RangeStop          int32
+	UnixTimeNs         int64
+	Broker             string
+	AssignedInstanceId ConsumerGroupInstanceId
+	FollowerBroker     string
+}
+
+type PartitionSlotToConsumerInstanceList struct {
+	PartitionSlots []*PartitionSlotToConsumerInstance
+	RingSize       int32
+	Version        int64
+}

weed/mq/sub_coordinator/sub_coordinator.go

@@ -0,0 +1,109 @@
+package sub_coordinator
+
+import (
+	"fmt"
+
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/filer_client"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+type TopicConsumerGroups struct {
+	// map a consumer group name to a consumer group
+	ConsumerGroups cmap.ConcurrentMap[string, *ConsumerGroup]
+}
+
+// SubCoordinator coordinates the instances in the consumer group for one topic.
+// It is responsible for:
+// 1. (Maybe) assigning partitions when a consumer instance is up/down.
+
+type SubCoordinator struct {
+	// map topic name to consumer groups
+	TopicSubscribers    cmap.ConcurrentMap[string, *TopicConsumerGroups]
+	FilerClientAccessor *filer_client.FilerClientAccessor
+}
+
+func NewSubCoordinator() *SubCoordinator {
+	return &SubCoordinator{
+		TopicSubscribers: cmap.New[*TopicConsumerGroups](),
+	}
+}
+
+func (c *SubCoordinator) GetTopicConsumerGroups(topic *schema_pb.Topic, createIfMissing bool) *TopicConsumerGroups {
+	topicName := toTopicName(topic)
+	tcg, _ := c.TopicSubscribers.Get(topicName)
+	if tcg == nil && createIfMissing {
+		tcg = &TopicConsumerGroups{
+			ConsumerGroups: cmap.New[*ConsumerGroup](),
+		}
+		if !c.TopicSubscribers.SetIfAbsent(topicName, tcg) {
+			tcg, _ = c.TopicSubscribers.Get(topicName)
+		}
+	}
+	return tcg
+}
+func (c *SubCoordinator) RemoveTopic(topic *schema_pb.Topic) {
+	topicName := toTopicName(topic)
+	c.TopicSubscribers.Remove(topicName)
+}
+
+func toTopicName(topic *schema_pb.Topic) string {
+	topicName := topic.Namespace + "." + topic.Name
+	return topicName
+}
+
+func (c *SubCoordinator) AddSubscriber(initMessage *mq_pb.SubscriberToSubCoordinatorRequest_InitMessage) (*ConsumerGroup, *ConsumerGroupInstance, error) {
+	tcg := c.GetTopicConsumerGroups(initMessage.Topic, true)
+	cg, _ := tcg.ConsumerGroups.Get(initMessage.ConsumerGroup)
+	if cg == nil {
+		cg = NewConsumerGroup(initMessage.Topic, initMessage.RebalanceSeconds, c.FilerClientAccessor)
+		if cg != nil {
+			tcg.ConsumerGroups.SetIfAbsent(initMessage.ConsumerGroup, cg)
+		}
+		cg, _ = tcg.ConsumerGroups.Get(initMessage.ConsumerGroup)
+	}
+	if cg == nil {
+		return nil, nil, fmt.Errorf("fail to create consumer group %s: topic %s not found", initMessage.ConsumerGroup, initMessage.Topic)
+	}
+	cgi, _ := cg.ConsumerGroupInstances.Get(initMessage.ConsumerGroupInstanceId)
+	if cgi == nil {
+		cgi = NewConsumerGroupInstance(initMessage.ConsumerGroupInstanceId, initMessage.MaxPartitionCount)
+		if !cg.ConsumerGroupInstances.SetIfAbsent(initMessage.ConsumerGroupInstanceId, cgi) {
+			cgi, _ = cg.ConsumerGroupInstances.Get(initMessage.ConsumerGroupInstanceId)
+		}
+	}
+	cgi.MaxPartitionCount = initMessage.MaxPartitionCount
+	cg.Market.AddConsumerInstance(cgi)
+	return cg, cgi, nil
+}
+
+func (c *SubCoordinator) RemoveSubscriber(initMessage *mq_pb.SubscriberToSubCoordinatorRequest_InitMessage) {
+	tcg := c.GetTopicConsumerGroups(initMessage.Topic, false)
+	if tcg == nil {
+		return
+	}
+	cg, _ := tcg.ConsumerGroups.Get(initMessage.ConsumerGroup)
+	if cg == nil {
+		return
+	}
+	cg.ConsumerGroupInstances.Remove(initMessage.ConsumerGroupInstanceId)
+	cg.Market.RemoveConsumerInstance(ConsumerGroupInstanceId(initMessage.ConsumerGroupInstanceId))
+	if cg.ConsumerGroupInstances.Count() == 0 {
+		tcg.ConsumerGroups.Remove(initMessage.ConsumerGroup)
+		cg.Shutdown()
+	}
+	if tcg.ConsumerGroups.Count() == 0 {
+		c.RemoveTopic(initMessage.Topic)
+	}
+}
+
+func (c *SubCoordinator) OnPartitionChange(topic *schema_pb.Topic, assignments []*mq_pb.BrokerPartitionAssignment) {
+	tcg, _ := c.TopicSubscribers.Get(toTopicName(topic))
+	if tcg == nil {
+		return
+	}
+	for _, cg := range tcg.ConsumerGroups.Items() {
+		cg.OnPartitionListChange(assignments)
+	}
+}