add publisher shutdown

weed/mq/client/pub_client/lookup.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"github.com/seaweedfs/seaweedfs/weed/pb"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func (p *TopicPublisher) doLookup(brokerAddress string) error {
@@ -100,6 +102,10 @@ func (p *TopicPublisher) doConnect(partition *mq_pb.Partition, brokerAddress str
 		for {
 			_, err := publishClient.Recv()
 			if err != nil {
+				e, ok := status.FromError(err)
+				if ok && e.Code() == codes.Unknown && e.Message() == "EOF" {
+					return
+				}
 				publishClient.Err = err
 				fmt.Printf("publish to %s error: %v\n", publishClient.Broker, err)
 				return

weed/mq/client/pub_client/publisher.go

@@ -2,10 +2,12 @@ package pub_client
 
 import (
 	"github.com/rdleal/intervalst/interval"
+	"github.com/seaweedfs/seaweedfs/weed/mq/broker"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"sync"
+	"time"
 )
 
 type PublisherConfiguration struct {
@@ -41,3 +43,15 @@ func (p *TopicPublisher) Connect(bootstrapBroker string) error {
 	}
 	return nil
 }
+
+func (p *TopicPublisher) Shutdown() error {
+
+	if clients, found := p.partition2Broker.AllIntersections(0, broker.MaxPartitionCount); found {
+		for _, client := range clients {
+			client.CloseSend()
+		}
+	}
+	time.Sleep(1100 * time.Millisecond)
+
+	return nil
+}