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/util/buffered_queue"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"sync"
	"time"
)

type PublisherConfiguration struct {
	Topic                     topic.Topic
	CreateTopic               bool
	CreateTopicPartitionCount int32
}

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 {
	return &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,
	}
}

func (p *TopicPublisher) Shutdown() error {

	if inputBuffers, found := p.partition2Buffer.AllIntersections(0, pub_balancer.MaxPartitionCount); found {
		for _, inputBuffer := range inputBuffers {
			inputBuffer.CloseInput()
		}
	}
	time.Sleep(1100 * time.Millisecond)

	return nil
}