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

weed/mq/client/sub_client/connect_to_sub_coordinator.go

@@ -0,0 +1,106 @@
+package sub_client
+
+import (
+	"time"
+
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+)
+
+func (sub *TopicSubscriber) doKeepConnectedToSubCoordinator() {
+	waitTime := 1 * time.Second
+	for {
+		for _, broker := range sub.bootstrapBrokers {
+
+			select {
+			case <-sub.ctx.Done():
+				return
+			default:
+			}
+
+			// lookup topic brokers
+			var brokerLeader string
+			err := pb.WithBrokerGrpcClient(false, broker, sub.SubscriberConfig.GrpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+				resp, err := client.FindBrokerLeader(sub.ctx, &mq_pb.FindBrokerLeaderRequest{})
+				if err != nil {
+					return err
+				}
+				brokerLeader = resp.Broker
+				return nil
+			})
+			if err != nil {
+				glog.V(0).Infof("broker coordinator on %s: %v", broker, err)
+				continue
+			}
+			glog.V(0).Infof("found broker coordinator: %v", brokerLeader)
+
+			// connect to the balancer
+			pb.WithBrokerGrpcClient(true, brokerLeader, sub.SubscriberConfig.GrpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+
+				stream, err := client.SubscriberToSubCoordinator(sub.ctx)
+				if err != nil {
+					glog.V(0).Infof("subscriber %s: %v", sub.ContentConfig.Topic, err)
+					return err
+				}
+				waitTime = 1 * time.Second
+
+				// Maybe later: subscribe to multiple topics instead of just one
+
+				if err := stream.Send(&mq_pb.SubscriberToSubCoordinatorRequest{
+					Message: &mq_pb.SubscriberToSubCoordinatorRequest_Init{
+						Init: &mq_pb.SubscriberToSubCoordinatorRequest_InitMessage{
+							ConsumerGroup:           sub.SubscriberConfig.ConsumerGroup,
+							ConsumerGroupInstanceId: sub.SubscriberConfig.ConsumerGroupInstanceId,
+							Topic:                   sub.ContentConfig.Topic.ToPbTopic(),
+							MaxPartitionCount:       sub.SubscriberConfig.MaxPartitionCount,
+						},
+					},
+				}); err != nil {
+					glog.V(0).Infof("subscriber %s send init: %v", sub.ContentConfig.Topic, err)
+					return err
+				}
+
+				go func() {
+					for reply := range sub.brokerPartitionAssignmentAckChan {
+
+						select {
+						case <-sub.ctx.Done():
+							return
+						default:
+						}
+
+						glog.V(0).Infof("subscriber instance %s ack %+v", sub.SubscriberConfig.ConsumerGroupInstanceId, reply)
+						if err := stream.Send(reply); err != nil {
+							glog.V(0).Infof("subscriber %s reply: %v", sub.ContentConfig.Topic, err)
+							return
+						}
+					}
+				}()
+
+				// keep receiving messages from the sub coordinator
+				for {
+					resp, err := stream.Recv()
+					if err != nil {
+						glog.V(0).Infof("subscriber %s receive: %v", sub.ContentConfig.Topic, err)
+						return err
+					}
+
+					select {
+					case <-sub.ctx.Done():
+						return nil
+					default:
+					}
+
+					sub.brokerPartitionAssignmentChan <- resp
+					glog.V(0).Infof("Received assignment: %+v", resp)
+				}
+			})
+		}
+		glog.V(0).Infof("subscriber %s/%s waiting for more assignments", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup)
+		if waitTime < 10*time.Second {
+			waitTime += 1 * time.Second
+		}
+		time.Sleep(waitTime)
+	}
+}

weed/mq/client/sub_client/on_each_partition.go

@@ -0,0 +1,142 @@
+package sub_client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+type KeyedTimestamp struct {
+	Key  []byte
+	TsNs int64 // Timestamp in nanoseconds for acknowledgment
+}
+
+func (sub *TopicSubscriber) onEachPartition(assigned *mq_pb.BrokerPartitionAssignment, stopCh chan struct{}, onDataMessageFn OnDataMessageFn) error {
+	// connect to the partition broker
+	return pb.WithBrokerGrpcClient(true, assigned.LeaderBroker, sub.SubscriberConfig.GrpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+
+		subscribeClient, err := client.SubscribeMessage(context.Background())
+		if err != nil {
+			return fmt.Errorf("create subscribe client: %w", err)
+		}
+
+		slidingWindowSize := sub.SubscriberConfig.SlidingWindowSize
+		if slidingWindowSize <= 0 {
+			slidingWindowSize = 1
+		}
+
+		po := findPartitionOffset(sub.ContentConfig.PartitionOffsets, assigned.Partition)
+		if po == nil {
+			po = &schema_pb.PartitionOffset{
+				Partition: assigned.Partition,
+				StartTsNs: sub.ContentConfig.OffsetTsNs,
+			}
+		}
+
+		if err = subscribeClient.Send(&mq_pb.SubscribeMessageRequest{
+			Message: &mq_pb.SubscribeMessageRequest_Init{
+				Init: &mq_pb.SubscribeMessageRequest_InitMessage{
+					ConsumerGroup:     sub.SubscriberConfig.ConsumerGroup,
+					ConsumerId:        sub.SubscriberConfig.ConsumerGroupInstanceId,
+					Topic:             sub.ContentConfig.Topic.ToPbTopic(),
+					PartitionOffset:   po,
+					OffsetType:        sub.ContentConfig.OffsetType,
+					Filter:            sub.ContentConfig.Filter,
+					FollowerBroker:    assigned.FollowerBroker,
+					SlidingWindowSize: slidingWindowSize,
+				},
+			},
+		}); err != nil {
+			glog.V(0).Infof("subscriber %s connected to partition %+v at %v: %v", sub.ContentConfig.Topic, assigned.Partition, assigned.LeaderBroker, err)
+		}
+
+		glog.V(0).Infof("subscriber %s connected to partition %+v at %v", sub.ContentConfig.Topic, assigned.Partition, assigned.LeaderBroker)
+
+		if sub.OnCompletionFunc != nil {
+			defer sub.OnCompletionFunc()
+		}
+
+		go func() {
+			for {
+				select {
+				case <-sub.ctx.Done():
+					subscribeClient.CloseSend()
+					return
+				case <-stopCh:
+					subscribeClient.CloseSend()
+					return
+				case ack, ok := <-sub.PartitionOffsetChan:
+					if !ok {
+						subscribeClient.CloseSend()
+						return
+					}
+					subscribeClient.SendMsg(&mq_pb.SubscribeMessageRequest{
+						Message: &mq_pb.SubscribeMessageRequest_Ack{
+							Ack: &mq_pb.SubscribeMessageRequest_AckMessage{
+								Key:  ack.Key,
+								TsNs: ack.TsNs,
+							},
+						},
+					})
+				}
+			}
+		}()
+
+		for {
+			// glog.V(0).Infof("subscriber %s/%s waiting for message", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup)
+			resp, err := subscribeClient.Recv()
+			if err != nil {
+				if errors.Is(err, io.EOF) {
+					return nil
+				}
+				return fmt.Errorf("subscribe recv: %w", err)
+			}
+			if resp.Message == nil {
+				glog.V(0).Infof("subscriber %s/%s received nil message", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup)
+				continue
+			}
+
+			select {
+			case <-sub.ctx.Done():
+				return nil
+			case <-stopCh:
+				return nil
+			default:
+			}
+
+			switch m := resp.Message.(type) {
+			case *mq_pb.SubscribeMessageResponse_Data:
+				if m.Data.Ctrl != nil {
+					glog.V(2).Infof("subscriber %s received control from producer:%s isClose:%v", sub.SubscriberConfig.ConsumerGroup, m.Data.Ctrl.PublisherName, m.Data.Ctrl.IsClose)
+					continue
+				}
+				if len(m.Data.Key) == 0 {
+					// fmt.Printf("empty key %+v, type %v\n", m, reflect.TypeOf(m))
+					continue
+				}
+				onDataMessageFn(m)
+			case *mq_pb.SubscribeMessageResponse_Ctrl:
+				// glog.V(0).Infof("subscriber %s/%s/%s received control %+v", sub.ContentConfig.Namespace, sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup, m.Ctrl)
+				if m.Ctrl.IsEndOfStream || m.Ctrl.IsEndOfTopic {
+					return io.EOF
+				}
+			}
+		}
+
+	})
+}
+
+func findPartitionOffset(partitionOffsets []*schema_pb.PartitionOffset, partition *schema_pb.Partition) *schema_pb.PartitionOffset {
+	for _, po := range partitionOffsets {
+		if po.Partition == partition {
+			return po
+		}
+	}
+	return nil
+}

weed/mq/client/sub_client/subscribe.go

@@ -0,0 +1,138 @@
+package sub_client
+
+import (
+	"sync"
+	"time"
+
+	"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/util"
+)
+
+type ProcessorState struct {
+	stopCh chan struct{}
+}
+
+// Subscribe subscribes to a topic's specified partitions.
+// If a partition is moved to another broker, the subscriber will automatically reconnect to the new broker.
+
+func (sub *TopicSubscriber) Subscribe() error {
+
+	go sub.startProcessors()
+
+	// loop forever
+	// TODO shutdown the subscriber when not needed anymore
+	sub.doKeepConnectedToSubCoordinator()
+
+	return nil
+}
+
+func (sub *TopicSubscriber) startProcessors() {
+	// listen to the messages from the sub coordinator
+	// start one processor per partition
+	var wg sync.WaitGroup
+	semaphore := make(chan struct{}, sub.SubscriberConfig.MaxPartitionCount)
+
+	for message := range sub.brokerPartitionAssignmentChan {
+		if assigned := message.GetAssignment(); assigned != nil {
+			wg.Add(1)
+			semaphore <- struct{}{}
+
+			topicPartition := topic.FromPbPartition(assigned.PartitionAssignment.Partition)
+
+			// wait until no covering partition is still in progress
+			sub.waitUntilNoOverlappingPartitionInFlight(topicPartition)
+
+			// start a processors
+			stopChan := make(chan struct{})
+			sub.activeProcessorsLock.Lock()
+			sub.activeProcessors[topicPartition] = &ProcessorState{
+				stopCh: stopChan,
+			}
+			sub.activeProcessorsLock.Unlock()
+
+			go func(assigned *mq_pb.BrokerPartitionAssignment, topicPartition topic.Partition) {
+				defer func() {
+					sub.activeProcessorsLock.Lock()
+					delete(sub.activeProcessors, topicPartition)
+					sub.activeProcessorsLock.Unlock()
+
+					<-semaphore
+					wg.Done()
+				}()
+				glog.V(0).Infof("subscriber %s/%s assigned partition %+v at %v", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup, assigned.Partition, assigned.LeaderBroker)
+				sub.brokerPartitionAssignmentAckChan <- &mq_pb.SubscriberToSubCoordinatorRequest{
+					Message: &mq_pb.SubscriberToSubCoordinatorRequest_AckAssignment{
+						AckAssignment: &mq_pb.SubscriberToSubCoordinatorRequest_AckAssignmentMessage{
+							Partition: assigned.Partition,
+						},
+					},
+				}
+
+				executors := util.NewLimitedConcurrentExecutor(int(sub.SubscriberConfig.SlidingWindowSize))
+				onDataMessageFn := func(m *mq_pb.SubscribeMessageResponse_Data) {
+					executors.Execute(func() {
+						if sub.OnDataMessageFunc != nil {
+							sub.OnDataMessageFunc(m)
+						}
+						sub.PartitionOffsetChan <- KeyedTimestamp{
+							Key:  m.Data.Key,
+							TsNs: m.Data.TsNs,
+						}
+					})
+				}
+
+				err := sub.onEachPartition(assigned, stopChan, onDataMessageFn)
+				if err != nil {
+					glog.V(0).Infof("subscriber %s/%s partition %+v at %v: %v", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup, assigned.Partition, assigned.LeaderBroker, err)
+				} else {
+					glog.V(0).Infof("subscriber %s/%s partition %+v at %v completed", sub.ContentConfig.Topic, sub.SubscriberConfig.ConsumerGroup, assigned.Partition, assigned.LeaderBroker)
+				}
+				sub.brokerPartitionAssignmentAckChan <- &mq_pb.SubscriberToSubCoordinatorRequest{
+					Message: &mq_pb.SubscriberToSubCoordinatorRequest_AckUnAssignment{
+						AckUnAssignment: &mq_pb.SubscriberToSubCoordinatorRequest_AckUnAssignmentMessage{
+							Partition: assigned.Partition,
+						},
+					},
+				}
+			}(assigned.PartitionAssignment, topicPartition)
+		}
+		if unAssignment := message.GetUnAssignment(); unAssignment != nil {
+			topicPartition := topic.FromPbPartition(unAssignment.Partition)
+			sub.activeProcessorsLock.Lock()
+			if processor, found := sub.activeProcessors[topicPartition]; found {
+				close(processor.stopCh)
+				delete(sub.activeProcessors, topicPartition)
+			}
+			sub.activeProcessorsLock.Unlock()
+		}
+	}
+
+	wg.Wait()
+
+}
+
+func (sub *TopicSubscriber) waitUntilNoOverlappingPartitionInFlight(topicPartition topic.Partition) {
+	foundOverlapping := true
+	for foundOverlapping {
+		sub.activeProcessorsLock.Lock()
+		foundOverlapping = false
+		var overlappedPartition topic.Partition
+		for partition, _ := range sub.activeProcessors {
+			if partition.Overlaps(topicPartition) {
+				if partition.Equals(topicPartition) {
+					continue
+				}
+				foundOverlapping = true
+				overlappedPartition = partition
+				break
+			}
+		}
+		sub.activeProcessorsLock.Unlock()
+		if foundOverlapping {
+			glog.V(0).Infof("subscriber %s new partition %v waiting for partition %+v to complete", sub.ContentConfig.Topic, topicPartition, overlappedPartition)
+			time.Sleep(1 * time.Second)
+		}
+	}
+}

weed/mq/client/sub_client/subscriber.go

@@ -0,0 +1,70 @@
+package sub_client
+
+import (
+	"context"
+	"sync"
+
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+	"google.golang.org/grpc"
+)
+
+type SubscriberConfiguration struct {
+	ClientId                string
+	ConsumerGroup           string
+	ConsumerGroupInstanceId string
+	GrpcDialOption          grpc.DialOption
+	MaxPartitionCount       int32 // how many partitions to process concurrently
+	SlidingWindowSize       int32 // how many messages to process concurrently per partition
+}
+
+func (s *SubscriberConfiguration) String() string {
+	return "ClientId: " + s.ClientId + ", ConsumerGroup: " + s.ConsumerGroup + ", ConsumerGroupInstanceId: " + s.ConsumerGroupInstanceId
+}
+
+type ContentConfiguration struct {
+	Topic            topic.Topic
+	Filter           string
+	PartitionOffsets []*schema_pb.PartitionOffset
+	OffsetType       schema_pb.OffsetType
+	OffsetTsNs       int64
+}
+
+type OnDataMessageFn func(m *mq_pb.SubscribeMessageResponse_Data)
+type OnCompletionFunc func()
+
+type TopicSubscriber struct {
+	ctx                              context.Context
+	SubscriberConfig                 *SubscriberConfiguration
+	ContentConfig                    *ContentConfiguration
+	brokerPartitionAssignmentChan    chan *mq_pb.SubscriberToSubCoordinatorResponse
+	brokerPartitionAssignmentAckChan chan *mq_pb.SubscriberToSubCoordinatorRequest
+	OnDataMessageFunc                OnDataMessageFn
+	OnCompletionFunc                 OnCompletionFunc
+	bootstrapBrokers                 []string
+	activeProcessors                 map[topic.Partition]*ProcessorState
+	activeProcessorsLock             sync.Mutex
+	PartitionOffsetChan              chan KeyedTimestamp
+}
+
+func NewTopicSubscriber(ctx context.Context, bootstrapBrokers []string, subscriber *SubscriberConfiguration, content *ContentConfiguration, partitionOffsetChan chan KeyedTimestamp) *TopicSubscriber {
+	return &TopicSubscriber{
+		ctx:                              ctx,
+		SubscriberConfig:                 subscriber,
+		ContentConfig:                    content,
+		brokerPartitionAssignmentChan:    make(chan *mq_pb.SubscriberToSubCoordinatorResponse, 1024),
+		brokerPartitionAssignmentAckChan: make(chan *mq_pb.SubscriberToSubCoordinatorRequest, 1024),
+		bootstrapBrokers:                 bootstrapBrokers,
+		activeProcessors:                 make(map[topic.Partition]*ProcessorState),
+		PartitionOffsetChan:              partitionOffsetChan,
+	}
+}
+
+func (sub *TopicSubscriber) SetOnDataMessageFn(fn OnDataMessageFn) {
+	sub.OnDataMessageFunc = fn
+}
+
+func (sub *TopicSubscriber) SetCompletionFunc(onCompletionFn OnCompletionFunc) {
+	sub.OnCompletionFunc = onCompletionFn
+}