connect and publish

weed/mq/client/sub_client/lookup.go

@@ -0,0 +1,74 @@
+package sub_client
+
+import (
+	"context"
+	"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 {
+	err := pb.WithBrokerGrpcClient(true,
+		brokerAddress,
+		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,
+					},
+					IsForPublish: true,
+				})
+			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)
+				}
+			}
+			return nil
+		})
+
+	if err != nil {
+		return fmt.Errorf("lookup topic %s/%s: %v", p.namespace, p.topic, err)
+	}
+	return nil
+}

weed/mq/client/sub_client/subscribe.go

@@ -0,0 +1,28 @@
+package sub_client
+
+import (
+	cmap "github.com/orcaman/concurrent-map"
+	"github.com/rdleal/intervalst/interval"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+)
+
+type SubscriberConfiguration struct {
+}
+
+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](),
+	}
+}

weed/mq/client/sub_client/subscriber.go

@@ -1,50 +0,0 @@
-package main
-
-import (
-	"context"
-	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/pb"
-	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
-)
-
-func main() {
-	err := pb.WithBrokerGrpcClient(true,
-		"localhost:17777",
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-		func(client mq_pb.SeaweedMessagingClient) error {
-			subClient, err := client.Subscribe(context.Background(), &mq_pb.SubscribeRequest{
-				Init: &mq_pb.SubscribeRequest_InitMessage{
-					Topic: &mq_pb.Topic{
-						Namespace: "test",
-						Name:      "test",
-					},
-				},
-			})
-			if err != nil {
-				return err
-			}
-
-			for {
-				resp, err := subClient.Recv()
-				if err != nil {
-					return err
-				}
-				if resp.GetCtrl() != nil {
-					if resp.GetCtrl().Error != "" {
-						return fmt.Errorf("ctrl error: %v", resp.GetCtrl().Error)
-					}
-				}
-				if resp.GetData() != nil {
-					println(string(resp.GetData().Key), "=>", string(resp.GetData().Value))
-				}
-
-			}
-			return nil
-		})
-
-	if err != nil {
-		fmt.Println(err)
-	}
-}