gsh-digital-services/seaweedFS
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

api for sub

Browse Source
This commit is contained in:
chrislu
2023-09-04 21:43:50 -07:00
parent 9e4f985698
commit ba67e6ca29
7 changed files with 189 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
weed/mq/client/sub_client/subscribe.go
View File

@@ -8,33 +8,34 @@ import (
"sync"
)
type EachMessageFunc func(key, value []byte) (shouldContinue bool)
type FinalFunc func()
// Subscribe subscribes to a topic's specified partitions.
// If a partition is moved to another broker, the subscriber will automatically reconnect to the new broker.
func (sub *TopicSubscriber) Subscribe(eachMessageFn EachMessageFunc, finalFn FinalFunc) error {
func (sub *TopicSubscriber) Subscribe() error {
var wg sync.WaitGroup
for _, brokerPartitionAssignment := range sub.brokerPartitionAssignments {
brokerAddress := brokerPartitionAssignment.LeaderBroker
grpcConnection, err := pb.GrpcDial(context.Background(), brokerAddress, true, sub.grpcDialOption)
grpcConnection, err := pb.GrpcDial(context.Background(), brokerAddress, true, sub.SubscriberConfig.GrpcDialOption)
if err != nil {
return fmt.Errorf("dial broker %s: %v", brokerAddress, err)
}
brokerClient := mq_pb.NewSeaweedMessagingClient(grpcConnection)
subscribeClient, err := brokerClient.Subscribe(context.Background(), &mq_pb.SubscribeRequest{
Consumer: &mq_pb.SubscribeRequest_Consumer{
ConsumerGroup: sub.config.ConsumerGroup,
ConsumerId: sub.config.ConsumerId,
ConsumerGroup: sub.SubscriberConfig.GroupId,
ConsumerId: sub.SubscriberConfig.GroupInstanceId,
},
Cursor: &mq_pb.SubscribeRequest_Cursor{
Topic: &mq_pb.Topic{
Namespace: sub.namespace,
Name: sub.topic,
Namespace: sub.ContentConfig.Namespace,
Name: sub.ContentConfig.Topic,
},
Partition: &mq_pb.Partition{
RingSize: brokerPartitionAssignment.Partition.RingSize,
RangeStart: brokerPartitionAssignment.Partition.RangeStart,
RangeStop: brokerPartitionAssignment.Partition.RangeStop,
},
Filter: sub.ContentConfig.Filter,
},
})
if err != nil {
@@ -43,8 +44,8 @@ func (sub *TopicSubscriber) Subscribe(eachMessageFn EachMessageFunc, finalFn Fin
wg.Add(1)
go func() {
defer wg.Done()
if finalFn != nil {
defer finalFn()
if sub.OnCompletionFunc != nil {
defer sub.OnCompletionFunc()
}
for {
resp, err := subscribeClient.Recv()
@@ -57,7 +58,7 @@ func (sub *TopicSubscriber) Subscribe(eachMessageFn EachMessageFunc, finalFn Fin
}
switch m := resp.Message.(type) {
case *mq_pb.SubscribeResponse_Data:
if !eachMessageFn(m.Data.Key, m.Data.Value) {
if !sub.OnEachMessageFunc(m.Data.Key, m.Data.Value) {
return
}
case *mq_pb.SubscribeResponse_Ctrl: