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

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