You can not select more than 25 topics
			Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
		
		
		
		
		
			
		
			
				
					
					
						
							73 lines
						
					
					
						
							1.8 KiB
						
					
					
				
			
		
		
		
			
			
			
		
		
	
	
							73 lines
						
					
					
						
							1.8 KiB
						
					
					
				
								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) (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
							 | 
						|
								}
							 |