can pub and sub

weed/mq/client/sub_client/subscribe.go

@@ -1,28 +1,71 @@
 package sub_client
 
 import (
-	cmap "github.com/orcaman/concurrent-map"
-	"github.com/rdleal/intervalst/interval"
+	"context"
+	"fmt"
+	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"sync"
 )
 
-type SubscriberConfiguration struct {
-}
+type EachMessageFunc func(key, value []byte) (shouldContinue bool)
+type FinalFunc func()
 
-type TopicSubscriber struct {
-	namespace            string
-	topic                string
-	partition2Broker     *interval.SearchTree[string, int32]
-	broker2PublishClient cmap.ConcurrentMap[string, mq_pb.SeaweedMessaging_PublishClient]
-}
-
-func NewTopicSubscriber(config *SubscriberConfiguration, namespace, topic string) *TopicSubscriber {
-	return &TopicSubscriber{
-		namespace: namespace,
-		topic:     topic,
-		partition2Broker: interval.NewSearchTree[string](func(a, b int32) int {
-			return int(a - b)
-		}),
-		broker2PublishClient: cmap.New[mq_pb.SeaweedMessaging_PublishClient](),
+func (sub *TopicSubscriber) Subscribe(eachMessageFn EachMessageFunc, finalFn FinalFunc) error {
+	var wg sync.WaitGroup
+	for _, brokerPartitionAssignment := range sub.brokerPartitionAssignments {
+		brokerAddress := brokerPartitionAssignment.LeaderBroker
+		grpcConnection, err := pb.GrpcDial(context.Background(), brokerAddress, true, sub.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,
+			},
+			Cursor: &mq_pb.SubscribeRequest_Cursor{
+				Topic: &mq_pb.Topic{
+					Namespace: sub.namespace,
+					Name:      sub.topic,
+				},
+				Partition: &mq_pb.Partition{
+					RingSize:   brokerPartitionAssignment.Partition.RingSize,
+					RangeStart: brokerPartitionAssignment.Partition.RangeStart,
+					RangeStop:  brokerPartitionAssignment.Partition.RangeStop,
+				},
+			},
+		})
+		if err != nil {
+			return fmt.Errorf("create subscribe client: %v", err)
+		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if finalFn != nil {
+				defer finalFn()
+			}
+			for {
+				resp, err := subscribeClient.Recv()
+				if err != nil {
+					fmt.Printf("subscribe error: %v\n", err)
+					return
+				}
+				if resp.Message == nil {
+					continue
+				}
+				switch m := resp.Message.(type) {
+				case *mq_pb.SubscribeResponse_Data:
+					if !eachMessageFn(m.Data.Key, m.Data.Value) {
+						return
+					}
+				case *mq_pb.SubscribeResponse_Ctrl:
+					// ignore
+				}
+			}
+		}()
 	}
+	wg.Wait()
+	return nil
 }