publisher is able to fully send data and receive acks

still needs to close the pipes cleanly

weed/mq/broker/broker_grpc_pub.go

@@ -7,11 +7,11 @@ import (
 	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/peer"
 	"io"
 	"math/rand"
 	"net"
-	"sync/atomic"
 )
 
 // PUB
@@ -41,6 +41,7 @@ func (b *MessageQueueBroker) PublishMessage(stream mq_pb.SeaweedMessaging_Publis
 	// 3. write to the filer
 
 	var localTopicPartition *topic.LocalPartition
+	var isGenerated bool
 	req, err := stream.Recv()
 	if err != nil {
 		return err
@@ -49,62 +50,122 @@ func (b *MessageQueueBroker) PublishMessage(stream mq_pb.SeaweedMessaging_Publis
 	// TODO check whether current broker should be the leader for the topic partition
 	ackInterval := 1
 	initMessage := req.GetInit()
-	var t topic.Topic
-	var p topic.Partition
-	if initMessage != nil {
-		t, p = topic.FromPbTopic(initMessage.Topic), topic.FromPbPartition(initMessage.Partition)
-		localTopicPartition, _, err = b.GetOrGenLocalPartition(t, p)
-		if err != nil {
-			response.Error = fmt.Sprintf("topic %v partition %v not setup", initMessage.Topic, initMessage.Partition)
-			glog.Errorf("topic %v partition %v not setup", initMessage.Topic, initMessage.Partition)
-			return stream.Send(response)
-		}
-		ackInterval = int(initMessage.AckInterval)
-		for _, follower := range initMessage.FollowerBrokers {
-			followErr := b.withBrokerClient(false, pb.ServerAddress(follower), func(client mq_pb.SeaweedMessagingClient) error {
-				_, err := client.PublishFollowMe(context.Background(), &mq_pb.PublishFollowMeRequest{
-					Topic:      initMessage.Topic,
-					Partition:  initMessage.Partition,
-					BrokerSelf: string(b.option.BrokerAddress()),
-				})
-				return err
-			})
-			if followErr != nil {
-				response.Error = fmt.Sprintf("follower %v failed: %v", follower, followErr)
-				glog.Errorf("follower %v failed: %v", follower, followErr)
-				return stream.Send(response)
-			}
-		}
-		stream.Send(response)
-	} else {
+	if initMessage == nil {
 		response.Error = fmt.Sprintf("missing init message")
 		glog.Errorf("missing init message")
 		return stream.Send(response)
 	}
 
+	// get or generate a local partition
+	t, p := topic.FromPbTopic(initMessage.Topic), topic.FromPbPartition(initMessage.Partition)
+	localTopicPartition, isGenerated, err = b.GetOrGenLocalPartition(t, p)
+	if err != nil {
+		response.Error = fmt.Sprintf("topic %v partition %v not setup", initMessage.Topic, initMessage.Partition)
+		glog.Errorf("topic %v partition %v not setup", initMessage.Topic, initMessage.Partition)
+		return stream.Send(response)
+	}
+	ackInterval = int(initMessage.AckInterval)
+
+	// connect to follower brokers
+	var followerStream mq_pb.SeaweedMessaging_PublishFollowMeClient
+	var grpcConnection *grpc.ClientConn
+	if isGenerated && len(initMessage.FollowerBrokers) > 0 {
+		follower := initMessage.FollowerBrokers[0]
+		ctx := stream.Context()
+		grpcConnection, err = pb.GrpcDial(ctx, follower, true, b.grpcDialOption)
+		if err != nil {
+			response.Error = fmt.Sprintf("fail to dial %s: %v", follower, err)
+			glog.Errorf("fail to dial %s: %v", follower, err)
+			return stream.Send(response)
+		}
+		followerClient := mq_pb.NewSeaweedMessagingClient(grpcConnection)
+		followerStream, err = followerClient.PublishFollowMe(ctx)
+		if err != nil {
+			response.Error = fmt.Sprintf("fail to create publish client: %v", err)
+			glog.Errorf("fail to create publish client: %v", err)
+			return stream.Send(response)
+		}
+		if err = followerStream.Send(&mq_pb.PublishFollowMeRequest{
+			Message: &mq_pb.PublishFollowMeRequest_Init{
+				Init: &mq_pb.PublishFollowMeRequest_InitMessage{
+					Topic:      initMessage.Topic,
+					Partition:  initMessage.Partition,
+				},
+			},
+		}); err != nil {
+			return err
+		}
+	}
+
+	// process each published messages
 	clientName := fmt.Sprintf("%v-%4d/%s/%v", findClientAddress(stream.Context()), rand.Intn(10000), initMessage.Topic, initMessage.Partition)
 	localTopicPartition.Publishers.AddPublisher(clientName, topic.NewLocalPublisher())
 
 	ackCounter := 0
 	var ackSequence int64
-	var isStopping int32
-	respChan := make(chan *mq_pb.PublishMessageResponse, 128)
 	defer func() {
-		atomic.StoreInt32(&isStopping, 1)
-		respChan <- &mq_pb.PublishMessageResponse{
-			AckSequence: ackSequence,
-		}
-		close(respChan)
-		localTopicPartition.Publishers.RemovePublisher(clientName)
-		glog.V(0).Infof("topic %v partition %v published %d messges Publisher:%d Subscriber:%d", initMessage.Topic, initMessage.Partition, ackSequence, localTopicPartition.Publishers.Size(), localTopicPartition.Subscribers.Size())
-		if localTopicPartition.MaybeShutdownLocalPartition() {
-			b.localTopicManager.RemoveTopicPartition(t, p)
+		if followerStream == nil {
+			// remove the publisher
+			localTopicPartition.Publishers.RemovePublisher(clientName)
+			glog.V(0).Infof("topic %v partition %v published %d messges Publisher:%d Subscriber:%d", initMessage.Topic, initMessage.Partition, ackSequence, localTopicPartition.Publishers.Size(), localTopicPartition.Subscribers.Size())
+			if localTopicPartition.MaybeShutdownLocalPartition() {
+				b.localTopicManager.RemoveTopicPartition(t, p)
+			}
 		}
 	}()
-	go func() {
-		for resp := range respChan {
-			if err := stream.Send(resp); err != nil {
-				glog.Errorf("Error sending response %v: %v", resp, err)
+
+	if followerStream != nil {
+		go func() {
+			defer func() {
+				println("stop receiving ack from follower")
+
+				// remove the publisher
+				localTopicPartition.Publishers.RemovePublisher(clientName)
+				glog.V(0).Infof("topic %v partition %v published %d messges Publisher:%d Subscriber:%d", initMessage.Topic, initMessage.Partition, ackSequence, localTopicPartition.Publishers.Size(), localTopicPartition.Subscribers.Size())
+				if localTopicPartition.MaybeShutdownLocalPartition() {
+					b.localTopicManager.RemoveTopicPartition(t, p)
+				}
+				println("closing grpcConnection to follower")
+				grpcConnection.Close()
+			}()
+
+			for {
+				ack, err := followerStream.Recv()
+				if err != nil {
+					glog.Errorf("Error receiving response: %v", err)
+					return
+				}
+				ackSequence = ack.AckTsNs
+				println("recv ack", ack.AckTsNs)
+				if err := stream.Send(&mq_pb.PublishMessageResponse{
+					AckSequence: ack.AckTsNs,
+				}); err != nil {
+					glog.Errorf("Error sending response %v: %v", ack, err)
+					return
+				}
+			}
+		}()
+	}
+
+	// send a hello message
+	stream.Send(&mq_pb.PublishMessageResponse{})
+
+	var receivedSequence, acknowledgedSequence  int64
+
+	defer func() {
+		if followerStream != nil {
+			//if err := followerStream.CloseSend(); err != nil {
+			//	glog.Errorf("Error closing follower stream: %v", err)
+			//}
+		} else {
+			if acknowledgedSequence < receivedSequence {
+				acknowledgedSequence = receivedSequence
+				response := &mq_pb.PublishMessageResponse{
+					AckSequence: acknowledgedSequence,
+				}
+				if err := stream.Send(response); err != nil {
+					glog.Errorf("Error sending response %v: %v", response, err)
+				}
 			}
 		}
 	}()
@@ -122,22 +183,57 @@ func (b *MessageQueueBroker) PublishMessage(stream mq_pb.SeaweedMessaging_Publis
 		}
 
 		// Process the received message
-		if dataMessage := req.GetData(); dataMessage != nil {
-			localTopicPartition.Publish(dataMessage)
+		dataMessage := req.GetData()
+		if dataMessage == nil {
+			continue
 		}
 
-		ackCounter++
-		ackSequence++
-		if ackCounter >= ackInterval {
-			ackCounter = 0
-			// send back the ack
-			response := &mq_pb.PublishMessageResponse{
-				AckSequence: ackSequence,
+		// send to the local partition
+		localTopicPartition.Publish(dataMessage)
+		receivedSequence = dataMessage.TsNs
+
+		// maybe send to the follower
+		if followerStream != nil {
+			println("recv", string(dataMessage.Key), dataMessage.TsNs)
+			if followErr := followerStream.Send(&mq_pb.PublishFollowMeRequest{
+				Message: &mq_pb.PublishFollowMeRequest_Data{
+					Data: dataMessage,
+				},
+			}); followErr != nil {
+				return followErr
+			}
+		} else {
+			ackCounter++
+			if ackCounter >= ackInterval {
+				ackCounter = 0
+				// send back the ack directly
+				acknowledgedSequence = receivedSequence
+				response := &mq_pb.PublishMessageResponse{
+					AckSequence: acknowledgedSequence,
+				}
+				if err := stream.Send(response); err != nil {
+					glog.Errorf("Error sending response %v: %v", response, err)
+				}
 			}
-			respChan <- response
 		}
 	}
 
+	if followerStream != nil {
+		// send close to the follower
+		if followErr := followerStream.Send(&mq_pb.PublishFollowMeRequest{
+			Message: &mq_pb.PublishFollowMeRequest_Close{
+				Close: &mq_pb.PublishFollowMeRequest_CloseMessage{},
+			},
+		}); followErr != nil {
+			return followErr
+		}
+		println("closing follower stream")
+
+		//if err := followerStream.CloseSend(); err != nil {
+		//	glog.Errorf("Error closing follower stream: %v", err)
+		//}
+	}
+
 	glog.V(0).Infof("topic %v partition %v publish stream closed.", initMessage.Topic, initMessage.Partition)
 
 	return nil

weed/mq/broker/broker_grpc_pub_follow.go

@@ -1,12 +1,49 @@
 package broker
 
 import (
-	"context"
+	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"io"
 )
 
-func (b *MessageQueueBroker) PublishFollowMe(c context.Context, request *mq_pb.PublishFollowMeRequest) (*mq_pb.PublishFollowMeResponse, error) {
-	glog.V(0).Infof("PublishFollowMe %v", request)
-	return &mq_pb.PublishFollowMeResponse{}, nil
+func (b *MessageQueueBroker) PublishFollowMe(stream mq_pb.SeaweedMessaging_PublishFollowMeServer) error {
+	req, err := stream.Recv()
+	if err != nil {
+		return err
+	}
+	initMessage := req.GetInit()
+	if initMessage == nil {
+		return fmt.Errorf("missing init message")
+	}
+
+	// t, p := topic.FromPbTopic(initMessage.Topic), topic.FromPbPartition(initMessage.Partition)
+	// follow each published messages
+	for {
+		// receive a message
+		req, err := stream.Recv()
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			glog.V(0).Infof("topic %v partition %v publish stream error: %v", initMessage.Topic, initMessage.Partition, err)
+			return err
+		}
+
+		// Process the received message
+		if dataMessage := req.GetData(); dataMessage != nil {
+			// send back the ack
+			if err := stream.Send(&mq_pb.PublishFollowMeResponse{
+				AckTsNs: dataMessage.TsNs,
+			}); err != nil {
+				// TODO save un-acked messages to disk
+				glog.Errorf("Error sending response %v: %v", dataMessage, err)
+			}
+			println("ack", string(dataMessage.Key), dataMessage.TsNs)
+		} else if closeMessage := req.GetClose(); closeMessage != nil {
+			glog.V(0).Infof("topic %v partition %v publish stream closed: %v", initMessage.Topic, initMessage.Partition, closeMessage)
+			break
+		}
+	}
+	return nil
 }

weed/mq/client/pub_client/publish.go

@@ -5,6 +5,7 @@ import (
 	"github.com/seaweedfs/seaweedfs/weed/mq/pub_balancer"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"github.com/seaweedfs/seaweedfs/weed/util"
+	"time"
 )
 
 func (p *TopicPublisher) Publish(key, value []byte) error {
@@ -20,5 +21,6 @@ func (p *TopicPublisher) Publish(key, value []byte) error {
 	return inputBuffer.Enqueue(&mq_pb.DataMessage{
 		Key:   key,
 		Value: value,
+		TsNs:  time.Now().UnixNano(),
 	})
 }

weed/mq/client/pub_client/scheduler.go

@@ -150,6 +150,7 @@ func (p *TopicPublisher) doPublishToPartition(job *EachPartitionPublishJob) erro
 	}); err != nil {
 		return fmt.Errorf("send init message: %v", err)
 	}
+	// process the hello message
 	resp, err := stream.Recv()
 	if err != nil {
 		return fmt.Errorf("recv init response: %v", err)
@@ -158,21 +159,25 @@ func (p *TopicPublisher) doPublishToPartition(job *EachPartitionPublishJob) erro
 		return fmt.Errorf("init response error: %v", resp.Error)
 	}
 
+	var wg sync.WaitGroup
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		for {
 			ackResp, err := publishClient.Recv()
 			if err != nil {
 				e, _ := status.FromError(err)
 				if e.Code() == codes.Unknown && e.Message() == "EOF" {
+					log.Printf("publish to %s EOF", publishClient.Broker)
 					return
 				}
 				publishClient.Err = err
-				fmt.Printf("publish1 to %s error: %v\n", publishClient.Broker, err)
+				log.Printf("publish1 to %s error: %v\n", publishClient.Broker, err)
 				return
 			}
 			if ackResp.Error != "" {
 				publishClient.Err = fmt.Errorf("ack error: %v", ackResp.Error)
-				fmt.Printf("publish2 to %s error: %v\n", publishClient.Broker, ackResp.Error)
+				log.Printf("publish2 to %s error: %v\n", publishClient.Broker, ackResp.Error)
 				return
 			}
 			if ackResp.AckSequence > 0 {
@@ -193,11 +198,12 @@ func (p *TopicPublisher) doPublishToPartition(job *EachPartitionPublishJob) erro
 		publishCounter++
 	}
 
-	if err := publishClient.CloseSend(); err != nil {
-		return fmt.Errorf("close send: %v", err)
-	}
+	// CloseSend would cancel the context on the server side
+	//if err := publishClient.CloseSend(); err != nil {
+	//	return fmt.Errorf("close send: %v", err)
+	//}
 
-	time.Sleep(3 * time.Second)
+	wg.Wait()
 
 	log.Printf("published %d messages to %v for topic partition %+v", publishCounter, job.LeaderBroker, job.Partition)