[DRAFT ]PMM-7000 Reconnect gRPC stream.

client/channel/channel.go

@@ -17,16 +17,21 @@
 package channel
 
 import (
+	"context"
+	"io"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/golang/protobuf/proto" //nolint:staticcheck
 	"github.com/percona/pmm/api/agentpb"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
 	protostatus "google.golang.org/genproto/googleapis/rpc/status"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
 	grpcstatus "google.golang.org/grpc/status"
 )
 
@@ -64,8 +69,9 @@ type Response struct {
 //
 // All exported methods are thread-safe.
 type Channel struct { //nolint:maligned
-	s agentpb.Agent_ConnectClient
-	l *logrus.Entry
+	s  agentpb.Agent_ConnectClient
+	l  *logrus.Entry
+	MD *agentpb.ServerConnectMetadata
 
 	mRecv, mSend prometheus.Counter
 
@@ -80,15 +86,46 @@ type Channel struct { //nolint:maligned
 	closeOnce sync.Once
 	closeWait chan struct{}
 	closeErr  error
+
+	reconnect chan bool
+	done      chan bool
 }
 
 // New creates new two-way communication channel with given stream.
 //
 // Stream should not be used by the caller after channel is created.
-func New(stream agentpb.Agent_ConnectClient) *Channel {
+func New(conn *grpc.ClientConn, l *logrus.Entry, ID, version string) *Channel {
+	// gRPC stream is created without lifetime timeout.
+	// However, we need to cancel it if two-way communication channel can't be established
+	// when pmm-managed is down. A separate timer is used for that.
+	streamCtx, streamCancel := context.WithCancel(context.Background())
+	l.Info("Establishing two-way communication channel ...")
+	streamCtx = agentpb.AddAgentConnectMetadata(streamCtx, &agentpb.AgentConnectMetadata{
+		ID:      ID,
+		Version: version,
+	})
+
+	stream, err := agentpb.NewAgentClient(conn).Connect(streamCtx)
+	if err != nil {
+		l.Errorf("Failed to establish two-way communication channel: %s.", err)
+		streamCancel()
+		conn.Close()
+		return nil
+	}
+
+	md, err := agentpb.ReceiveServerConnectMetadata(stream)
+	l.Debugf("Received server metadata: %+v. Error: %+v.", md, err)
+	if err != nil {
+		l.Errorf("Failed to receive server metadata: %s.", err)
+		streamCancel()
+		conn.Close()
+		return nil
+	}
+
 	s := &Channel{
-		s: stream,
-		l: logrus.WithField("component", "channel"), // only for debug logging
+		s:  stream,
+		l:  logrus.WithField("component", "channel"), // only for debug logging
+		MD: md,
 
 		mRecv: prometheus.NewCounter(prometheus.CounterOpts{
 			Namespace: prometheusNamespace,
@@ -106,13 +143,76 @@ func New(stream agentpb.Agent_ConnectClient) *Channel {
 		responses: make(map[uint32]chan Response),
 		requests:  make(chan *ServerRequest, serverRequestsCap),
 
+		reconnect: make(chan bool),
+		done:      make(chan bool),
+
 		closeWait: make(chan struct{}),
 	}
 
 	go s.runReceiver()
+	go func() {
+		for {
+			select {
+			case <-s.reconnect:
+				if !s.waitUntilReady(conn, l) {
+					l.Errorf("Failed to establish two-way communication channel.")
+				}
+				streamCtx, streamCancel := context.WithCancel(context.Background())
+				streamCtx = agentpb.AddAgentConnectMetadata(streamCtx, &agentpb.AgentConnectMetadata{
+					ID:      ID,
+					Version: version,
+				})
+
+				stream, err := agentpb.NewAgentClient(conn).Connect(streamCtx)
+				if err != nil {
+					l.Errorf("Failed to establish two-way communication channel: %s.", err)
+					streamCancel()
+					conn.Close()
+					return
+				}
+				s.s = stream
+				s.runReceiver()
+				// streamCancel()
+			case <-s.done:
+				l.Infoln("Done.")
+				streamCancel()
+				close(s.requests)
+				return
+			}
+		}
+	}()
+
 	return s
 }
 
+// waitUntilReady implements queues the RPCs until the channel is READY
+// https://github.com/grpc/grpc/blob/master/doc/wait-for-ready.md
+func (c *Channel) waitUntilReady(conn *grpc.ClientConn, l *logrus.Entry) bool {
+	//define how long you want to wait for connection to be restored before giving up
+	ctx, cancel := context.WithTimeout(context.Background(), 720*time.Hour) // 30 days
+	defer cancel()
+
+	currentState := conn.GetState()
+	stillConnecting := true
+
+	for currentState != connectivity.Ready && stillConnecting {
+		if currentState == connectivity.Shutdown { // pmm-server was shutted down
+			c.done <- true
+			return false
+		}
+		//will return true when state has changed from thisState, false if timeout
+		stillConnecting = conn.WaitForStateChange(ctx, currentState)
+		currentState = conn.GetState()
+	}
+
+	if stillConnecting == false {
+		l.Error("Connection attempt has timed out.")
+		return false
+	}
+
+	return true
+}
+
 // close marks channel as closed with given error - only once.
 func (c *Channel) close(err error) {
 	c.closeOnce.Do(func() {
@@ -205,16 +305,20 @@ func (c *Channel) send(msg *agentpb.AgentMessage) {
 // runReader receives messages from server
 func (c *Channel) runReceiver() {
 	defer func() {
-		close(c.requests)
 		c.l.Debug("Exiting receiver goroutine.")
 	}()
 
 	for {
 		msg, err := c.s.Recv()
-		if err != nil {
+		if err == io.EOF {
+			c.done <- true
 			c.close(errors.Wrap(err, "failed to receive message"))
 			return
 		}
+		if err != nil {
+			c.reconnect <- true
+			return
+		}
 		c.mRecv.Inc()
 
 		// do not use default compact representation for large/complex messages

client/client.go

@@ -35,6 +35,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/keepalive"
 	grpcstatus "google.golang.org/grpc/status"
 
 	"github.com/percona/pmm-agent/actions" // TODO https://jira.percona.com/browse/PMM-7206
@@ -550,10 +551,9 @@ func (c *Client) getActionTimeout(req *agentpb.StartActionRequest) time.Duration
 }
 
 type dialResult struct {
-	conn         *grpc.ClientConn
-	streamCancel context.CancelFunc
-	channel      *channel.Channel
-	md           *agentpb.ServerConnectMetadata
+	conn    *grpc.ClientConn
+	channel *channel.Channel
+	md      *agentpb.ServerConnectMetadata
 }
 
 // dial tries to connect to the server once.
@@ -562,6 +562,11 @@ func dial(dialCtx context.Context, cfg *config.Config, l *logrus.Entry) (*dialRe
 	opts := []grpc.DialOption{
 		grpc.WithBlock(),
 		grpc.WithUserAgent("pmm-agent/" + version.Version),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:                10 * time.Second,
+			Timeout:             100 * time.Millisecond,
+			PermitWithoutStream: true,
+		}),
 	}
 	if cfg.Server.WithoutTLS {
 		opts = append(opts, grpc.WithInsecure())
@@ -595,43 +600,16 @@ func dial(dialCtx context.Context, cfg *config.Config, l *logrus.Entry) (*dialRe
 	}
 	l.Infof("Connected to %s.", cfg.Server.Address)
 
-	// gRPC stream is created without lifetime timeout.
-	// However, we need to cancel it if two-way communication channel can't be established
-	// when pmm-managed is down. A separate timer is used for that.
-	streamCtx, streamCancel := context.WithCancel(context.Background())
-	teardown := func() {
-		streamCancel()
-		if err := conn.Close(); err != nil {
-			l.Debugf("Connection closed: %s.", err)
-			return
-		}
-		l.Debugf("Connection closed.")
-	}
-	d, ok := dialCtx.Deadline()
+	_, ok := dialCtx.Deadline()
 	if !ok {
 		panic("no deadline in dialCtx")
 	}
-	streamCancelT := time.AfterFunc(time.Until(d), streamCancel)
-	defer streamCancelT.Stop()
-
-	l.Info("Establishing two-way communication channel ...")
-	start := time.Now()
-	streamCtx = agentpb.AddAgentConnectMetadata(streamCtx, &agentpb.AgentConnectMetadata{
-		ID:      cfg.ID,
-		Version: version.Version,
-	})
-	stream, err := agentpb.NewAgentClient(conn).Connect(streamCtx)
-	if err != nil {
-		l.Errorf("Failed to establish two-way communication channel: %s.", err)
-		teardown()
-		return nil, errors.Wrap(err, "failed to connect")
-	}
 
 	// So far, nginx can handle all that itself without pmm-managed.
 	// We need to exchange one pair of messages (ping/pong) for metadata headers to reach pmm-managed
 	// to ensure that pmm-managed is alive and that Agent ID is valid.
-
-	channel := channel.New(stream)
+	start := time.Now()
+	channel := channel.New(conn, l, cfg.ID, version.Version)
 	_, clockDrift, err := getNetworkInformation(channel) // ping/pong
 	if err != nil {
 		msg := err.Error()
@@ -642,36 +620,28 @@ func dial(dialCtx context.Context, cfg *config.Config, l *logrus.Entry) (*dialRe
 		}
 
 		l.Errorf("Failed to establish two-way communication channel: %s.", msg)
-		teardown()
+		conn.Close()
 		return nil, err
 	}
 
-	// read metadata header after receiving pong
-	md, err := agentpb.ReceiveServerConnectMetadata(stream)
-	l.Debugf("Received server metadata: %+v. Error: %+v.", md, err)
-	if err != nil {
-		l.Errorf("Failed to receive server metadata: %s.", err)
-		teardown()
-		return nil, errors.Wrap(err, "failed to receive server metadata")
-	}
-	if md.ServerVersion == "" {
-		l.Errorf("Server metadata does not contain server version.")
-		teardown()
-		return nil, errors.New("empty server version in metadata")
-	}
-
 	level := logrus.InfoLevel
 	if clockDrift > clockDriftWarning || -clockDrift > clockDriftWarning {
 		level = logrus.WarnLevel
 	}
 	l.Logf(level, "Two-way communication channel established in %s. Estimated clock drift: %s.",
 		time.Since(start), clockDrift)
 
+	if channel.MD.ServerVersion == "" {
+		l.Errorf("Server metadata does not contain server version.")
+		conn.Close()
+		return nil, errors.New("empty server version in metadata")
+	}
+
 	return &dialResult{
-		conn:         conn,
-		streamCancel: streamCancel,
-		channel:      channel,
-		md:           md}, nil
+		conn:    conn,
+		channel: channel,
+		md:      channel.MD,
+	}, nil
 }
 
 func getNetworkInformation(channel *channel.Channel) (latency, clockDrift time.Duration, err error) {

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/go-openapi/runtime v0.19.20
 	github.com/go-sql-driver/mysql v1.5.0
 	github.com/golang/protobuf v1.4.3
+	github.com/google/martian v2.1.0+incompatible
 	github.com/google/uuid v1.1.2 // indirect
 	github.com/grpc-ecosystem/grpc-gateway v1.15.1
 	github.com/hashicorp/go-version v1.3.0

go.sum

@@ -182,6 +182,8 @@ github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/go-cmp v0.5.0 h1:/QaMHBdZ26BB3SSst0Iwl10Epc+xhTquomWX0oZEB6w=
 github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/martian v2.1.0+incompatible h1:/CP5g8u/VJHijgedC/Legn3BAbAaWPgecwXBIDzw5no=
+github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXiIaQGbYVAs8BPL6v8lEs=
 github.com/google/uuid v1.0.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y=