seaweedfs/weed/mq/client/pub_client/publisher.go

package pub_client

import (
	"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"
	"log"
	"sync"
)

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) *TopicPublisher {
	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 tp
}

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
}