add subscribe follower

weed/mq/broker/broker_grpc_sub_follow.go

@@ -0,0 +1,76 @@
+package broker
+
+import (
+	"fmt"
+	"github.com/seaweedfs/seaweedfs/weed/filer"
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/util"
+	"io"
+	"time"
+)
+
+
+func (b *MessageQueueBroker) SubscribeFollowMe(stream mq_pb.SeaweedMessaging_SubscribeFollowMeServer) (err error) {
+	var req *mq_pb.SubscribeFollowMeRequest
+	req, err = stream.Recv()
+	if err != nil {
+		return err
+	}
+	initMessage := req.GetInit()
+	if initMessage == nil {
+		return fmt.Errorf("missing init message")
+	}
+
+	// create an in-memory offset
+	subscriberProgress := &SubscriberProgress{
+		Topic: topic.FromPbTopic(initMessage.Topic),
+		Partition: topic.FromPbPartition(initMessage.Partition),
+		ConsumerGroup: "consumer_group",
+	}
+
+	// follow each published messages
+	for {
+		// receive a message
+		req, err = stream.Recv()
+		if err != nil {
+			if err == io.EOF {
+				err = nil
+				break
+			}
+			glog.V(0).Infof("topic %v partition %v subscribe stream error: %v", initMessage.Topic, initMessage.Partition, err)
+			break
+		}
+
+		// Process the received message
+		if ackMessage := req.GetAck(); ackMessage != nil {
+			subscriberProgress.Offset = ackMessage.TsNs
+			println("offset", subscriberProgress.Offset)
+		} else if closeMessage := req.GetClose(); closeMessage != nil {
+			glog.V(0).Infof("topic %v partition %v subscribe stream closed: %v", initMessage.Topic, initMessage.Partition, closeMessage)
+			return nil
+		} else {
+			glog.Errorf("unknown message: %v", req)
+		}
+	}
+
+	t, p := topic.FromPbTopic(initMessage.Topic), topic.FromPbPartition(initMessage.Partition)
+
+	topicDir := fmt.Sprintf("%s/%s/%s", filer.TopicsDir, t.Namespace, t.Name)
+	partitionGeneration := time.Unix(0, p.UnixTimeNs).UTC().Format(topic.TIME_FORMAT)
+	partitionDir := fmt.Sprintf("%s/%s/%04d-%04d", topicDir, partitionGeneration, p.RangeStart, p.RangeStop)
+	offsetFileName := fmt.Sprintf("%s.offset", subscriberProgress.ConsumerGroup)
+
+	offsetBytes := make([]byte, 8)
+	util.Uint64toBytes(offsetBytes, uint64(subscriberProgress.Offset))
+
+	err = b.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
+		return filer.SaveInsideFiler(client, partitionDir, offsetFileName, offsetBytes)
+	})
+
+	glog.V(0).Infof("shut down follower for %v %v", t, p)
+
+	return err
+}

weed/mq/broker/subscriber_progress.go

@@ -0,0 +1,10 @@
+package broker
+
+import "github.com/seaweedfs/seaweedfs/weed/mq/topic"
+
+type SubscriberProgress struct {
+	topic.Topic
+	topic.Partition
+	ConsumerGroup string
+	Offset int64
+}

weed/pb/mq.proto

@@ -45,11 +45,13 @@ service SeaweedMessaging {
     // data plane for each topic partition
     rpc PublishMessage (stream PublishMessageRequest) returns (stream PublishMessageResponse) {
     }
-    rpc SubscribeMessage (SubscribeMessageRequest) returns (stream SubscribeMessageResponse) {
+    rpc SubscribeMessage (stream SubscribeMessageRequest) returns (stream SubscribeMessageResponse) {
     }
     // The lead broker asks a follower broker to follow itself
     rpc PublishFollowMe (stream PublishFollowMeRequest) returns (stream PublishFollowMeResponse) {
     }
+    rpc SubscribeFollowMe (stream SubscribeFollowMeRequest) returns (SubscribeFollowMeResponse) {
+    }
 }
 
 //////////////////////////////////////////////////
@@ -262,6 +264,25 @@ message SubscribeMessageResponse {
         DataMessage data = 2;
     }
 }
+message SubscribeFollowMeRequest {
+    message InitMessage {
+        Topic topic = 1;
+        Partition partition = 2;
+    }
+    message AckMessage {
+        int64 ts_ns = 1;
+    }
+    message CloseMessage {
+    }
+    oneof message {
+        InitMessage init = 1;
+        AckMessage ack = 2;
+        CloseMessage close = 3;
+    }
+}
+message SubscribeFollowMeResponse {
+    int64 ack_ts_ns = 1;
+}
 message ClosePublishersRequest {
     Topic topic = 1;
     int64 unix_time_ns = 2;

weed/pb/mq_pb/mq_grpc.pb.go

@@ -32,6 +32,7 @@ const (
 	SeaweedMessaging_PublishMessage_FullMethodName             = "/messaging_pb.SeaweedMessaging/PublishMessage"
 	SeaweedMessaging_SubscribeMessage_FullMethodName           = "/messaging_pb.SeaweedMessaging/SubscribeMessage"
 	SeaweedMessaging_PublishFollowMe_FullMethodName            = "/messaging_pb.SeaweedMessaging/PublishFollowMe"
+	SeaweedMessaging_SubscribeFollowMe_FullMethodName          = "/messaging_pb.SeaweedMessaging/SubscribeFollowMe"
 )
 
 // SeaweedMessagingClient is the client API for SeaweedMessaging service.
@@ -55,9 +56,10 @@ type SeaweedMessagingClient interface {
 	SubscriberToSubCoordinator(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_SubscriberToSubCoordinatorClient, error)
 	// data plane for each topic partition
 	PublishMessage(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_PublishMessageClient, error)
-	SubscribeMessage(ctx context.Context, in *SubscribeMessageRequest, opts ...grpc.CallOption) (SeaweedMessaging_SubscribeMessageClient, error)
+	SubscribeMessage(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_SubscribeMessageClient, error)
 	// The lead broker asks a follower broker to follow itself
 	PublishFollowMe(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_PublishFollowMeClient, error)
+	SubscribeFollowMe(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_SubscribeFollowMeClient, error)
 }
 
 type seaweedMessagingClient struct {
@@ -233,22 +235,17 @@ func (x *seaweedMessagingPublishMessageClient) Recv() (*PublishMessageResponse,
 	return m, nil
 }
 
-func (c *seaweedMessagingClient) SubscribeMessage(ctx context.Context, in *SubscribeMessageRequest, opts ...grpc.CallOption) (SeaweedMessaging_SubscribeMessageClient, error) {
+func (c *seaweedMessagingClient) SubscribeMessage(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_SubscribeMessageClient, error) {
 	stream, err := c.cc.NewStream(ctx, &SeaweedMessaging_ServiceDesc.Streams[3], SeaweedMessaging_SubscribeMessage_FullMethodName, opts...)
 	if err != nil {
 		return nil, err
 	}
 	x := &seaweedMessagingSubscribeMessageClient{stream}
-	if err := x.ClientStream.SendMsg(in); err != nil {
-		return nil, err
-	}
-	if err := x.ClientStream.CloseSend(); err != nil {
-		return nil, err
-	}
 	return x, nil
 }
 
 type SeaweedMessaging_SubscribeMessageClient interface {
+	Send(*SubscribeMessageRequest) error
 	Recv() (*SubscribeMessageResponse, error)
 	grpc.ClientStream
 }
@@ -257,6 +254,10 @@ type seaweedMessagingSubscribeMessageClient struct {
 	grpc.ClientStream
 }
 
+func (x *seaweedMessagingSubscribeMessageClient) Send(m *SubscribeMessageRequest) error {
+	return x.ClientStream.SendMsg(m)
+}
+
 func (x *seaweedMessagingSubscribeMessageClient) Recv() (*SubscribeMessageResponse, error) {
 	m := new(SubscribeMessageResponse)
 	if err := x.ClientStream.RecvMsg(m); err != nil {
@@ -296,6 +297,40 @@ func (x *seaweedMessagingPublishFollowMeClient) Recv() (*PublishFollowMeResponse
 	return m, nil
 }
 
+func (c *seaweedMessagingClient) SubscribeFollowMe(ctx context.Context, opts ...grpc.CallOption) (SeaweedMessaging_SubscribeFollowMeClient, error) {
+	stream, err := c.cc.NewStream(ctx, &SeaweedMessaging_ServiceDesc.Streams[5], SeaweedMessaging_SubscribeFollowMe_FullMethodName, opts...)
+	if err != nil {
+		return nil, err
+	}
+	x := &seaweedMessagingSubscribeFollowMeClient{stream}
+	return x, nil
+}
+
+type SeaweedMessaging_SubscribeFollowMeClient interface {
+	Send(*SubscribeFollowMeRequest) error
+	CloseAndRecv() (*SubscribeFollowMeResponse, error)
+	grpc.ClientStream
+}
+
+type seaweedMessagingSubscribeFollowMeClient struct {
+	grpc.ClientStream
+}
+
+func (x *seaweedMessagingSubscribeFollowMeClient) Send(m *SubscribeFollowMeRequest) error {
+	return x.ClientStream.SendMsg(m)
+}
+
+func (x *seaweedMessagingSubscribeFollowMeClient) CloseAndRecv() (*SubscribeFollowMeResponse, error) {
+	if err := x.ClientStream.CloseSend(); err != nil {
+		return nil, err
+	}
+	m := new(SubscribeFollowMeResponse)
+	if err := x.ClientStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
 // SeaweedMessagingServer is the server API for SeaweedMessaging service.
 // All implementations must embed UnimplementedSeaweedMessagingServer
 // for forward compatibility
@@ -317,9 +352,10 @@ type SeaweedMessagingServer interface {
 	SubscriberToSubCoordinator(SeaweedMessaging_SubscriberToSubCoordinatorServer) error
 	// data plane for each topic partition
 	PublishMessage(SeaweedMessaging_PublishMessageServer) error
-	SubscribeMessage(*SubscribeMessageRequest, SeaweedMessaging_SubscribeMessageServer) error
+	SubscribeMessage(SeaweedMessaging_SubscribeMessageServer) error
 	// The lead broker asks a follower broker to follow itself
 	PublishFollowMe(SeaweedMessaging_PublishFollowMeServer) error
+	SubscribeFollowMe(SeaweedMessaging_SubscribeFollowMeServer) error
 	mustEmbedUnimplementedSeaweedMessagingServer()
 }
 
@@ -360,12 +396,15 @@ func (UnimplementedSeaweedMessagingServer) SubscriberToSubCoordinator(SeaweedMes
 func (UnimplementedSeaweedMessagingServer) PublishMessage(SeaweedMessaging_PublishMessageServer) error {
 	return status.Errorf(codes.Unimplemented, "method PublishMessage not implemented")
 }
-func (UnimplementedSeaweedMessagingServer) SubscribeMessage(*SubscribeMessageRequest, SeaweedMessaging_SubscribeMessageServer) error {
+func (UnimplementedSeaweedMessagingServer) SubscribeMessage(SeaweedMessaging_SubscribeMessageServer) error {
 	return status.Errorf(codes.Unimplemented, "method SubscribeMessage not implemented")
 }
 func (UnimplementedSeaweedMessagingServer) PublishFollowMe(SeaweedMessaging_PublishFollowMeServer) error {
 	return status.Errorf(codes.Unimplemented, "method PublishFollowMe not implemented")
 }
+func (UnimplementedSeaweedMessagingServer) SubscribeFollowMe(SeaweedMessaging_SubscribeFollowMeServer) error {
+	return status.Errorf(codes.Unimplemented, "method SubscribeFollowMe not implemented")
+}
 func (UnimplementedSeaweedMessagingServer) mustEmbedUnimplementedSeaweedMessagingServer() {}
 
 // UnsafeSeaweedMessagingServer may be embedded to opt out of forward compatibility for this service.
@@ -602,15 +641,12 @@ func (x *seaweedMessagingPublishMessageServer) Recv() (*PublishMessageRequest, e
 }
 
 func _SeaweedMessaging_SubscribeMessage_Handler(srv interface{}, stream grpc.ServerStream) error {
-	m := new(SubscribeMessageRequest)
-	if err := stream.RecvMsg(m); err != nil {
-		return err
-	}
-	return srv.(SeaweedMessagingServer).SubscribeMessage(m, &seaweedMessagingSubscribeMessageServer{stream})
+	return srv.(SeaweedMessagingServer).SubscribeMessage(&seaweedMessagingSubscribeMessageServer{stream})
 }
 
 type SeaweedMessaging_SubscribeMessageServer interface {
 	Send(*SubscribeMessageResponse) error
+	Recv() (*SubscribeMessageRequest, error)
 	grpc.ServerStream
 }
 
@@ -622,6 +658,14 @@ func (x *seaweedMessagingSubscribeMessageServer) Send(m *SubscribeMessageRespons
 	return x.ServerStream.SendMsg(m)
 }
 
+func (x *seaweedMessagingSubscribeMessageServer) Recv() (*SubscribeMessageRequest, error) {
+	m := new(SubscribeMessageRequest)
+	if err := x.ServerStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
 func _SeaweedMessaging_PublishFollowMe_Handler(srv interface{}, stream grpc.ServerStream) error {
 	return srv.(SeaweedMessagingServer).PublishFollowMe(&seaweedMessagingPublishFollowMeServer{stream})
 }
@@ -648,6 +692,32 @@ func (x *seaweedMessagingPublishFollowMeServer) Recv() (*PublishFollowMeRequest,
 	return m, nil
 }
 
+func _SeaweedMessaging_SubscribeFollowMe_Handler(srv interface{}, stream grpc.ServerStream) error {
+	return srv.(SeaweedMessagingServer).SubscribeFollowMe(&seaweedMessagingSubscribeFollowMeServer{stream})
+}
+
+type SeaweedMessaging_SubscribeFollowMeServer interface {
+	SendAndClose(*SubscribeFollowMeResponse) error
+	Recv() (*SubscribeFollowMeRequest, error)
+	grpc.ServerStream
+}
+
+type seaweedMessagingSubscribeFollowMeServer struct {
+	grpc.ServerStream
+}
+
+func (x *seaweedMessagingSubscribeFollowMeServer) SendAndClose(m *SubscribeFollowMeResponse) error {
+	return x.ServerStream.SendMsg(m)
+}
+
+func (x *seaweedMessagingSubscribeFollowMeServer) Recv() (*SubscribeFollowMeRequest, error) {
+	m := new(SubscribeFollowMeRequest)
+	if err := x.ServerStream.RecvMsg(m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
 // SeaweedMessaging_ServiceDesc is the grpc.ServiceDesc for SeaweedMessaging service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -711,6 +781,7 @@ var SeaweedMessaging_ServiceDesc = grpc.ServiceDesc{
 			StreamName:    "SubscribeMessage",
 			Handler:       _SeaweedMessaging_SubscribeMessage_Handler,
 			ServerStreams: true,
+			ClientStreams: true,
 		},
 		{
 			StreamName:    "PublishFollowMe",
@@ -718,6 +789,11 @@ var SeaweedMessaging_ServiceDesc = grpc.ServiceDesc{
 			ServerStreams: true,
 			ClientStreams: true,
 		},
+		{
+			StreamName:    "SubscribeFollowMe",
+			Handler:       _SeaweedMessaging_SubscribeFollowMe_Handler,
+			ClientStreams: true,
+		},
 	},
 	Metadata: "mq.proto",
 }