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

weed/mq/client/agent_client/publish_session.go

@@ -0,0 +1,82 @@
+package agent_client
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/seaweedfs/seaweedfs/weed/mq/schema"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+type PublishSession struct {
+	schema         *schema.Schema
+	partitionCount int
+	publisherName  string
+	stream         grpc.BidiStreamingClient[mq_agent_pb.PublishRecordRequest, mq_agent_pb.PublishRecordResponse]
+}
+
+func NewPublishSession(agentAddress string, topicSchema *schema.Schema, partitionCount int, publisherName string) (*PublishSession, error) {
+
+	// call local agent grpc server to create a new session
+	clientConn, err := grpc.NewClient(agentAddress, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		return nil, fmt.Errorf("dial agent server %s: %v", agentAddress, err)
+	}
+	agentClient := mq_agent_pb.NewSeaweedMessagingAgentClient(clientConn)
+
+	resp, err := agentClient.StartPublishSession(context.Background(), &mq_agent_pb.StartPublishSessionRequest{
+		Topic: &schema_pb.Topic{
+			Namespace: topicSchema.Namespace,
+			Name:      topicSchema.Name,
+		},
+		PartitionCount: int32(partitionCount),
+		RecordType:     topicSchema.RecordType,
+		PublisherName:  publisherName,
+	})
+	if err != nil {
+		return nil, err
+	}
+	if resp.Error != "" {
+		return nil, fmt.Errorf("start publish session: %v", resp.Error)
+	}
+
+	stream, err := agentClient.PublishRecord(context.Background())
+	if err != nil {
+		return nil, fmt.Errorf("publish record: %w", err)
+	}
+
+	if err = stream.Send(&mq_agent_pb.PublishRecordRequest{
+		SessionId: resp.SessionId,
+	}); err != nil {
+		return nil, fmt.Errorf("send session id: %w", err)
+	}
+
+	return &PublishSession{
+		schema:         topicSchema,
+		partitionCount: partitionCount,
+		publisherName:  publisherName,
+		stream:         stream,
+	}, nil
+}
+
+func (a *PublishSession) CloseSession() error {
+	if a.schema == nil {
+		return nil
+	}
+	err := a.stream.CloseSend()
+	if err != nil {
+		return fmt.Errorf("close send: %w", err)
+	}
+	a.schema = nil
+	return err
+}
+
+func (a *PublishSession) PublishMessageRecord(key []byte, record *schema_pb.RecordValue) error {
+	return a.stream.Send(&mq_agent_pb.PublishRecordRequest{
+		Key:   key,
+		Value: record,
+	})
+}

weed/mq/client/agent_client/subscribe_session.go

@@ -0,0 +1,87 @@
+package agent_client
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_agent_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+type SubscribeOption struct {
+	ConsumerGroup           string
+	ConsumerGroupInstanceId string
+	Topic                   topic.Topic
+	OffsetType              schema_pb.OffsetType
+	OffsetTsNs              int64
+	Filter                  string
+	MaxSubscribedPartitions int32
+	SlidingWindowSize       int32
+}
+
+type SubscribeSession struct {
+	Option *SubscribeOption
+	stream grpc.BidiStreamingClient[mq_agent_pb.SubscribeRecordRequest, mq_agent_pb.SubscribeRecordResponse]
+}
+
+func NewSubscribeSession(agentAddress string, option *SubscribeOption) (*SubscribeSession, error) {
+	// call local agent grpc server to create a new session
+	clientConn, err := grpc.NewClient(agentAddress, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		return nil, fmt.Errorf("dial agent server %s: %v", agentAddress, err)
+	}
+	agentClient := mq_agent_pb.NewSeaweedMessagingAgentClient(clientConn)
+
+	initRequest := &mq_agent_pb.SubscribeRecordRequest_InitSubscribeRecordRequest{
+		ConsumerGroup:           option.ConsumerGroup,
+		ConsumerGroupInstanceId: option.ConsumerGroupInstanceId,
+		Topic: &schema_pb.Topic{
+			Namespace: option.Topic.Namespace,
+			Name:      option.Topic.Name,
+		},
+		OffsetType:              option.OffsetType,
+		OffsetTsNs:              option.OffsetTsNs,
+		MaxSubscribedPartitions: option.MaxSubscribedPartitions,
+		Filter:                  option.Filter,
+		SlidingWindowSize:       option.SlidingWindowSize,
+	}
+
+	stream, err := agentClient.SubscribeRecord(context.Background())
+	if err != nil {
+		return nil, fmt.Errorf("subscribe record: %w", err)
+	}
+
+	if err = stream.Send(&mq_agent_pb.SubscribeRecordRequest{
+		Init: initRequest,
+	}); err != nil {
+		return nil, fmt.Errorf("send session id: %w", err)
+	}
+
+	return &SubscribeSession{
+		Option: option,
+		stream: stream,
+	}, nil
+}
+
+func (s *SubscribeSession) CloseSession() error {
+	err := s.stream.CloseSend()
+	return err
+}
+
+func (a *SubscribeSession) SubscribeMessageRecord(
+	onEachMessageFn func(key []byte, record *schema_pb.RecordValue),
+	onCompletionFn func()) error {
+	for {
+		resp, err := a.stream.Recv()
+		if err != nil {
+			if onCompletionFn != nil {
+				onCompletionFn()
+			}
+			return err
+		}
+		onEachMessageFn(resp.Key, resp.Value)
+	}
+}

weed/mq/client/pub_client/publish.go

@@ -0,0 +1,62 @@
+package pub_client
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/seaweedfs/seaweedfs/weed/mq/pub_balancer"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+	"github.com/seaweedfs/seaweedfs/weed/util"
+)
+
+func (p *TopicPublisher) Publish(key, value []byte) error {
+	if p.config.RecordType != nil {
+		return fmt.Errorf("record type is set, use PublishRecord instead")
+	}
+	return p.doPublish(key, value)
+}
+
+func (p *TopicPublisher) doPublish(key, value []byte) error {
+	hashKey := util.HashToInt32(key) % pub_balancer.MaxPartitionCount
+	if hashKey < 0 {
+		hashKey = -hashKey
+	}
+	inputBuffer, found := p.partition2Buffer.Floor(hashKey+1, hashKey+1)
+	if !found {
+		return fmt.Errorf("no input buffer found for key %d", hashKey)
+	}
+
+	return inputBuffer.Enqueue(&mq_pb.DataMessage{
+		Key:   key,
+		Value: value,
+		TsNs:  time.Now().UnixNano(),
+	})
+}
+
+func (p *TopicPublisher) PublishRecord(key []byte, recordValue *schema_pb.RecordValue) error {
+	// serialize record value
+	value, err := proto.Marshal(recordValue)
+	if err != nil {
+		return fmt.Errorf("failed to marshal record value: %w", err)
+	}
+
+	return p.doPublish(key, value)
+}
+
+func (p *TopicPublisher) FinishPublish() error {
+	if inputBuffers, found := p.partition2Buffer.AllIntersections(0, pub_balancer.MaxPartitionCount); found {
+		for _, inputBuffer := range inputBuffers {
+			inputBuffer.Enqueue(&mq_pb.DataMessage{
+				TsNs: time.Now().UnixNano(),
+				Ctrl: &mq_pb.ControlMessage{
+					IsClose:       true,
+					PublisherName: p.config.PublisherName,
+				},
+			})
+		}
+	}
+
+	return nil
+}

weed/mq/client/pub_client/publisher.go

@@ -0,0 +1,74 @@
+package pub_client
+
+import (
+	"log"
+	"sync"
+
+	"github.com/rdleal/intervalst/interval"
+	"github.com/seaweedfs/seaweedfs/weed/mq/pub_balancer"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+	"github.com/seaweedfs/seaweedfs/weed/util/buffered_queue"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+type PublisherConfiguration struct {
+	Topic          topic.Topic
+	PartitionCount int32
+	Brokers        []string
+	PublisherName  string // for debugging
+	RecordType     *schema_pb.RecordType
+}
+
+type PublishClient struct {
+	mq_pb.SeaweedMessaging_PublishMessageClient
+	Broker string
+	Err    error
+}
+type TopicPublisher struct {
+	partition2Buffer *interval.SearchTree[*buffered_queue.BufferedQueue[*mq_pb.DataMessage], int32]
+	grpcDialOption   grpc.DialOption
+	sync.Mutex       // protects grpc
+	config           *PublisherConfiguration
+	jobs             []*EachPartitionPublishJob
+}
+
+func NewTopicPublisher(config *PublisherConfiguration) (tp *TopicPublisher, err error) {
+	tp = &TopicPublisher{
+		partition2Buffer: interval.NewSearchTree[*buffered_queue.BufferedQueue[*mq_pb.DataMessage]](func(a, b int32) int {
+			return int(a - b)
+		}),
+		grpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
+		config:         config,
+	}
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		if err = tp.startSchedulerThread(&wg); err != nil {
+			log.Println(err)
+			return
+		}
+	}()
+
+	wg.Wait()
+
+	return
+}
+
+func (p *TopicPublisher) Shutdown() error {
+
+	if inputBuffers, found := p.partition2Buffer.AllIntersections(0, pub_balancer.MaxPartitionCount); found {
+		for _, inputBuffer := range inputBuffers {
+			inputBuffer.CloseInput()
+		}
+	}
+
+	for _, job := range p.jobs {
+		job.wg.Wait()
+	}
+
+	return nil
+}

weed/mq/client/pub_client/scheduler.go

@@ -0,0 +1,299 @@
+package pub_client
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"sort"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"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/util/buffered_queue"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
+)
+
+type EachPartitionError struct {
+	*mq_pb.BrokerPartitionAssignment
+	Err        error
+	generation int
+}
+
+type EachPartitionPublishJob struct {
+	*mq_pb.BrokerPartitionAssignment
+	stopChan   chan bool
+	wg         sync.WaitGroup
+	generation int
+	inputQueue *buffered_queue.BufferedQueue[*mq_pb.DataMessage]
+}
+
+func (p *TopicPublisher) startSchedulerThread(wg *sync.WaitGroup) error {
+
+	if err := p.doConfigureTopic(); err != nil {
+		wg.Done()
+		return fmt.Errorf("configure topic %s: %v", p.config.Topic, err)
+	}
+
+	log.Printf("start scheduler thread for topic %s", p.config.Topic)
+
+	generation := 0
+	var errChan chan EachPartitionError
+	for {
+		glog.V(0).Infof("lookup partitions gen %d topic %s", generation+1, p.config.Topic)
+		if assignments, err := p.doLookupTopicPartitions(); err == nil {
+			generation++
+			glog.V(0).Infof("start generation %d with %d assignments", generation, len(assignments))
+			if errChan == nil {
+				errChan = make(chan EachPartitionError, len(assignments))
+			}
+			p.onEachAssignments(generation, assignments, errChan)
+		} else {
+			glog.Errorf("lookup topic %s: %v", p.config.Topic, err)
+			time.Sleep(5 * time.Second)
+			continue
+		}
+
+		if generation == 1 {
+			wg.Done()
+		}
+
+		// wait for any error to happen. If so, consume all remaining errors, and retry
+		for {
+			select {
+			case eachErr := <-errChan:
+				glog.Errorf("gen %d publish to topic %s partition %v: %v", eachErr.generation, p.config.Topic, eachErr.Partition, eachErr.Err)
+				if eachErr.generation < generation {
+					continue
+				}
+				break
+			}
+		}
+	}
+}
+
+func (p *TopicPublisher) onEachAssignments(generation int, assignments []*mq_pb.BrokerPartitionAssignment, errChan chan EachPartitionError) {
+	// TODO assuming this is not re-configured so the partitions are fixed.
+	sort.Slice(assignments, func(i, j int) bool {
+		return assignments[i].Partition.RangeStart < assignments[j].Partition.RangeStart
+	})
+	var jobs []*EachPartitionPublishJob
+	hasExistingJob := len(p.jobs) == len(assignments)
+	for i, assignment := range assignments {
+		if assignment.LeaderBroker == "" {
+			continue
+		}
+		if hasExistingJob {
+			var existingJob *EachPartitionPublishJob
+			existingJob = p.jobs[i]
+			if pb.ServerAddress(existingJob.BrokerPartitionAssignment.LeaderBroker).Equals(pb.ServerAddress(assignment.LeaderBroker)) {
+				existingJob.generation = generation
+				jobs = append(jobs, existingJob)
+				continue
+			} else {
+				if existingJob.LeaderBroker != "" {
+					close(existingJob.stopChan)
+					existingJob.LeaderBroker = ""
+					existingJob.wg.Wait()
+				}
+			}
+		}
+
+		// start a go routine to publish to this partition
+		job := &EachPartitionPublishJob{
+			BrokerPartitionAssignment: assignment,
+			stopChan:                  make(chan bool, 1),
+			generation:                generation,
+			inputQueue:                buffered_queue.NewBufferedQueue[*mq_pb.DataMessage](1024),
+		}
+		job.wg.Add(1)
+		go func(job *EachPartitionPublishJob) {
+			defer job.wg.Done()
+			if err := p.doPublishToPartition(job); err != nil {
+				log.Printf("publish to %s partition %v: %v", p.config.Topic, job.Partition, err)
+				errChan <- EachPartitionError{assignment, err, generation}
+			}
+		}(job)
+		jobs = append(jobs, job)
+		// TODO assuming this is not re-configured so the partitions are fixed.
+		// better just reuse the existing job
+		p.partition2Buffer.Insert(assignment.Partition.RangeStart, assignment.Partition.RangeStop, job.inputQueue)
+	}
+	p.jobs = jobs
+}
+
+func (p *TopicPublisher) doPublishToPartition(job *EachPartitionPublishJob) error {
+
+	log.Printf("connecting to %v for topic partition %+v", job.LeaderBroker, job.Partition)
+
+	grpcConnection, err := grpc.NewClient(job.LeaderBroker, grpc.WithTransportCredentials(insecure.NewCredentials()), p.grpcDialOption)
+	if err != nil {
+		return fmt.Errorf("dial broker %s: %v", job.LeaderBroker, err)
+	}
+	brokerClient := mq_pb.NewSeaweedMessagingClient(grpcConnection)
+	stream, err := brokerClient.PublishMessage(context.Background())
+	if err != nil {
+		return fmt.Errorf("create publish client: %w", err)
+	}
+	publishClient := &PublishClient{
+		SeaweedMessaging_PublishMessageClient: stream,
+		Broker:                                job.LeaderBroker,
+	}
+	if err = publishClient.Send(&mq_pb.PublishMessageRequest{
+		Message: &mq_pb.PublishMessageRequest_Init{
+			Init: &mq_pb.PublishMessageRequest_InitMessage{
+				Topic:          p.config.Topic.ToPbTopic(),
+				Partition:      job.Partition,
+				AckInterval:    128,
+				FollowerBroker: job.FollowerBroker,
+				PublisherName:  p.config.PublisherName,
+			},
+		},
+	}); err != nil {
+		return fmt.Errorf("send init message: %w", err)
+	}
+	// process the hello message
+	resp, err := stream.Recv()
+	if err != nil {
+		return fmt.Errorf("recv init response: %w", err)
+	}
+	if resp.Error != "" {
+		return fmt.Errorf("init response error: %v", resp.Error)
+	}
+
+	var publishedTsNs int64
+	hasMoreData := int32(1)
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for {
+			ackResp, err := publishClient.Recv()
+			if err != nil {
+				e, _ := status.FromError(err)
+				if e.Code() == codes.Unknown && e.Message() == "EOF" {
+					log.Printf("publish to %s EOF", publishClient.Broker)
+					return
+				}
+				publishClient.Err = err
+				log.Printf("publish1 to %s error: %v\n", publishClient.Broker, err)
+				return
+			}
+			if ackResp.Error != "" {
+				publishClient.Err = fmt.Errorf("ack error: %v", ackResp.Error)
+				log.Printf("publish2 to %s error: %v\n", publishClient.Broker, ackResp.Error)
+				return
+			}
+			if ackResp.AckTsNs > 0 {
+				log.Printf("ack %d published %d hasMoreData:%d", ackResp.AckTsNs, atomic.LoadInt64(&publishedTsNs), atomic.LoadInt32(&hasMoreData))
+			}
+			if atomic.LoadInt64(&publishedTsNs) <= ackResp.AckTsNs && atomic.LoadInt32(&hasMoreData) == 0 {
+				return
+			}
+		}
+	}()
+
+	publishCounter := 0
+	for data, hasData := job.inputQueue.Dequeue(); hasData; data, hasData = job.inputQueue.Dequeue() {
+		if data.Ctrl != nil && data.Ctrl.IsClose {
+			// need to set this before sending to brokers, to avoid timing issue
+			atomic.StoreInt32(&hasMoreData, 0)
+		}
+		if err := publishClient.Send(&mq_pb.PublishMessageRequest{
+			Message: &mq_pb.PublishMessageRequest_Data{
+				Data: data,
+			},
+		}); err != nil {
+			return fmt.Errorf("send publish data: %w", err)
+		}
+		publishCounter++
+		atomic.StoreInt64(&publishedTsNs, data.TsNs)
+	}
+	if publishCounter > 0 {
+		wg.Wait()
+	} else {
+		// CloseSend would cancel the context on the server side
+		if err := publishClient.CloseSend(); err != nil {
+			return fmt.Errorf("close send: %w", err)
+		}
+	}
+
+	log.Printf("published %d messages to %v for topic partition %+v", publishCounter, job.LeaderBroker, job.Partition)
+
+	return nil
+}
+
+func (p *TopicPublisher) doConfigureTopic() (err error) {
+	if len(p.config.Brokers) == 0 {
+		return fmt.Errorf("topic configuring found no bootstrap brokers")
+	}
+	var lastErr error
+	for _, brokerAddress := range p.config.Brokers {
+		err = pb.WithBrokerGrpcClient(false,
+			brokerAddress,
+			p.grpcDialOption,
+			func(client mq_pb.SeaweedMessagingClient) error {
+				_, err := client.ConfigureTopic(context.Background(), &mq_pb.ConfigureTopicRequest{
+					Topic:             p.config.Topic.ToPbTopic(),
+					PartitionCount:    p.config.PartitionCount,
+					MessageRecordType: p.config.RecordType, // Flat schema
+				})
+				return err
+			})
+		if err == nil {
+			lastErr = nil
+			return nil
+		} else {
+			lastErr = fmt.Errorf("%s: %v", brokerAddress, err)
+		}
+	}
+
+	if lastErr != nil {
+		return fmt.Errorf("doConfigureTopic %s: %v", p.config.Topic, err)
+	}
+	return nil
+}
+
+func (p *TopicPublisher) doLookupTopicPartitions() (assignments []*mq_pb.BrokerPartitionAssignment, err error) {
+	if len(p.config.Brokers) == 0 {
+		return nil, fmt.Errorf("lookup found no bootstrap brokers")
+	}
+	var lastErr error
+	for _, brokerAddress := range p.config.Brokers {
+		err := pb.WithBrokerGrpcClient(false,
+			brokerAddress,
+			p.grpcDialOption,
+			func(client mq_pb.SeaweedMessagingClient) error {
+				lookupResp, err := client.LookupTopicBrokers(context.Background(),
+					&mq_pb.LookupTopicBrokersRequest{
+						Topic: p.config.Topic.ToPbTopic(),
+					})
+				glog.V(0).Infof("lookup topic %s: %v", p.config.Topic, lookupResp)
+
+				if err != nil {
+					return err
+				}
+
+				if len(lookupResp.BrokerPartitionAssignments) == 0 {
+					return fmt.Errorf("no broker partition assignments")
+				}
+
+				assignments = lookupResp.BrokerPartitionAssignments
+
+				return nil
+			})
+		if err == nil {
+			return assignments, nil
+		} else {
+			lastErr = err
+		}
+	}
+
+	return nil, fmt.Errorf("lookup topic %s: %v", p.config.Topic, lastErr)
+
+}

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
+}