can pub and sub

weed/mq/client/sub_client/lookup.go

@@ -5,70 +5,30 @@ import (
 	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
-	"google.golang.org/grpc"
 )
 
-func (p *TopicSubscriber) doLookup(
-	brokerAddress string, grpcDialOption grpc.DialOption) error {
+func (sub *TopicSubscriber) doLookup(brokerAddress string) error {
 	err := pb.WithBrokerGrpcClient(true,
 		brokerAddress,
-		grpcDialOption,
+		sub.grpcDialOption,
 		func(client mq_pb.SeaweedMessagingClient) error {
 			lookupResp, err := client.LookupTopicBrokers(context.Background(),
 				&mq_pb.LookupTopicBrokersRequest{
 					Topic: &mq_pb.Topic{
-						Namespace: p.namespace,
-						Name:      p.topic,
+						Namespace: sub.namespace,
+						Name:      sub.topic,
 					},
-					IsForPublish: true,
+					IsForPublish: false,
 				})
 			if err != nil {
 				return err
 			}
-			for _, brokerPartitionAssignment := range lookupResp.BrokerPartitionAssignments {
-				// partition => broker
-				p.partition2Broker.Insert(
-					brokerPartitionAssignment.Partition.RangeStart,
-					brokerPartitionAssignment.Partition.RangeStop,
-					brokerPartitionAssignment.LeaderBroker)
-
-				// broker => publish client
-				// send init message
-				// save the publishing client
-				brokerAddress := brokerPartitionAssignment.LeaderBroker
-				grpcConnection, err := pb.GrpcDial(context.Background(), brokerAddress, true, grpcDialOption)
-				if err != nil {
-					return fmt.Errorf("dial broker %s: %v", brokerAddress, err)
-				}
-				brokerClient := mq_pb.NewSeaweedMessagingClient(grpcConnection)
-				publishClient, err := brokerClient.Publish(context.Background())
-				if err != nil {
-					return fmt.Errorf("create publish client: %v", err)
-				}
-				p.broker2PublishClient.Set(brokerAddress, publishClient)
-				if err = publishClient.Send(&mq_pb.PublishRequest{
-					Message: &mq_pb.PublishRequest_Init{
-						Init: &mq_pb.PublishRequest_InitMessage{
-							Topic: &mq_pb.Topic{
-								Namespace: p.namespace,
-								Name:      p.topic,
-							},
-							Partition: &mq_pb.Partition{
-								RingSize:   brokerPartitionAssignment.Partition.RingSize,
-								RangeStart: brokerPartitionAssignment.Partition.RangeStart,
-								RangeStop:  brokerPartitionAssignment.Partition.RangeStop,
-							},
-						},
-					},
-				}); err != nil {
-					return fmt.Errorf("send init message: %v", err)
-				}
-			}
+			sub.brokerPartitionAssignments = lookupResp.BrokerPartitionAssignments
 			return nil
 		})
 
 	if err != nil {
-		return fmt.Errorf("lookup topic %s/%s: %v", p.namespace, p.topic, err)
+		return fmt.Errorf("lookup topic %s/%s: %v", sub.namespace, sub.topic, err)
 	}
 	return nil
 }

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
 }

weed/mq/client/sub_client/subscriber.go

@@ -0,0 +1,36 @@
+package sub_client
+
+import (
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+type SubscriberConfiguration struct {
+	ConsumerGroup string
+	ConsumerId    string
+}
+
+type TopicSubscriber struct {
+	config                     *SubscriberConfiguration
+	namespace                  string
+	topic                      string
+	brokerPartitionAssignments []*mq_pb.BrokerPartitionAssignment
+	grpcDialOption             grpc.DialOption
+}
+
+func NewTopicSubscriber(config *SubscriberConfiguration, namespace, topic string) *TopicSubscriber {
+	return &TopicSubscriber{
+		config:         config,
+		namespace:      namespace,
+		topic:          topic,
+		grpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
+	}
+}
+
+func (sub *TopicSubscriber) Connect(bootstrapBroker string) error {
+	if err := sub.doLookup(bootstrapBroker); err != nil {
+		return err
+	}
+	return nil
+}