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

58 lines
1.3 KiB
Go
Raw Normal View History

2023-08-28 16:02:12 +00:00
package pub_client
2023-08-27 20:13:14 +00:00
2023-08-28 00:50:59 +00:00
import (
2023-08-28 16:02:12 +00:00
"github.com/rdleal/intervalst/interval"
2023-09-19 23:22:41 +00:00
"github.com/seaweedfs/seaweedfs/weed/mq/balancer"
2023-08-28 00:50:59 +00:00
"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
2023-09-05 04:43:30 +00:00
"sync"
2023-09-08 06:55:19 +00:00
"time"
2023-08-28 00:50:59 +00:00
)
2023-08-28 16:02:12 +00:00
type PublisherConfiguration struct {
}
2023-09-05 04:43:30 +00:00
type PublishClient struct {
mq_pb.SeaweedMessaging_PublishClient
Broker string
Err error
}
2023-08-28 16:02:12 +00:00
type TopicPublisher struct {
2023-09-05 04:43:30 +00:00
namespace string
topic string
partition2Broker *interval.SearchTree[*PublishClient, int32]
grpcDialOption grpc.DialOption
sync.Mutex // protects grpc
2023-08-28 16:02:12 +00:00
}
2023-08-28 00:50:59 +00:00
2023-08-28 16:02:12 +00:00
func NewTopicPublisher(namespace, topic string) *TopicPublisher {
return &TopicPublisher{
namespace: namespace,
topic: topic,
2023-09-05 04:43:30 +00:00
partition2Broker: interval.NewSearchTree[*PublishClient](func(a, b int32) int {
2023-08-28 16:02:12 +00:00
return int(a - b)
}),
2023-09-05 04:43:30 +00:00
grpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
2023-08-28 00:50:59 +00:00
}
2023-08-28 16:02:12 +00:00
}
2023-08-28 00:50:59 +00:00
2023-08-28 16:02:12 +00:00
func (p *TopicPublisher) Connect(bootstrapBroker string) error {
2023-09-01 07:36:51 +00:00
if err := p.doLookup(bootstrapBroker); err != nil {
2023-08-28 16:02:12 +00:00
return err
}
return nil
2023-08-27 20:13:14 +00:00
}
2023-09-08 06:55:19 +00:00
func (p *TopicPublisher) Shutdown() error {
2023-09-19 23:22:41 +00:00
if clients, found := p.partition2Broker.AllIntersections(0, balancer.MaxPartitionCount); found {
2023-09-08 06:55:19 +00:00
for _, client := range clients {
client.CloseSend()
}
}
time.Sleep(1100 * time.Millisecond)
return nil
}