replication/driver: automatic retries on connectivity-related errors

rpc/dataconn/dataconn_client.go

@@ -213,3 +213,23 @@ func (c *Client) ReqRecv(ctx context.Context, req *pdu.ReceiveReq, streamCopier
 
 	return res.res, cause
 }
+
+
+func (c *Client) ReqPing(ctx context.Context, req *pdu.PingReq) (*pdu.PingRes, error) {
+	conn, err := c.getWire(ctx)
+	if err != nil {
+		return nil, err
+	}
+			defer c.putWire(conn)
+
+	if err := c.send(ctx, conn, EndpointPing, req, nil); err != nil {
+		return nil, err
+	}
+
+	var res pdu.PingRes
+	if err := c.recv(ctx, conn, &res); err != nil {
+		return nil, err
+	}
+
+	return &res, nil
+}

rpc/dataconn/dataconn_server.go

@@ -25,6 +25,8 @@ type Handler interface {
 	// It is guaranteed that Server calls Receive with a stream that holds the IdleConnTimeout
 	// configured in ServerConfig.Shared.IdleConnTimeout.
 	Receive(ctx context.Context, r *pdu.ReceiveReq, receive zfs.StreamCopier) (*pdu.ReceiveRes, error)
+	// PingDataconn handles a PingReq
+	PingDataconn(ctx context.Context, r *pdu.PingReq) (*pdu.PingRes, error)
 }
 
 type Logger = logger.Logger
@@ -125,6 +127,13 @@ func (s *Server) serveConn(nc *transport.AuthConn) {
 			return
 		}
 		res, handlerErr = s.h.Receive(ctx, &req, &streamCopier{streamConn: c, closeStreamOnClose: false}) // SHADOWING
+	case EndpointPing:
+		var req pdu.PingReq
+		if err := proto.Unmarshal(reqStructured, &req); err != nil {
+			s.log.WithError(err).Error("cannot unmarshal ping request")
+			return
+		}
+		res, handlerErr = s.h.PingDataconn(ctx, &req) // SHADOWING
 	default:
 		s.log.WithField("endpoint", endpoint).Error("unknown endpoint")
 		handlerErr = fmt.Errorf("requested endpoint does not exist")

rpc/dataconn/dataconn_shared.go

@@ -10,6 +10,7 @@ import (
 )
 
 const (
+	EndpointPing string = "/v1/ping"
 	EndpointSend string = "/v1/send"
 	EndpointRecv string = "/v1/recv"
 )

rpc/dataconn/microbenchmark/microbenchmark.go

@@ -77,6 +77,12 @@ func (devNullHandler) Receive(ctx context.Context, r *pdu.ReceiveReq, stream zfs
 	return &res, err
 }
 
+func (devNullHandler) PingDataconn(ctx context.Context, r *pdu.PingReq) (*pdu.PingRes, error) {
+	return &pdu.PingRes{
+		Echo: r.GetMessage(),
+	}, nil
+}
+
 type tcpConnecter struct {
 	addr string
 }

rpc/rpc_client.go

@@ -2,11 +2,16 @@ package rpc
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"net"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"google.golang.org/grpc"
 
+	"github.com/google/uuid"
 	"github.com/zrepl/zrepl/replication/logic"
 	"github.com/zrepl/zrepl/replication/logic/pdu"
 	"github.com/zrepl/zrepl/rpc/dataconn"
@@ -101,6 +106,67 @@ func (c *Client) ReplicationCursor(ctx context.Context, in *pdu.ReplicationCurso
 	return c.controlClient.ReplicationCursor(ctx, in)
 }
 
+func (c *Client) WaitForConnectivity(ctx context.Context) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	msg := uuid.New().String()
+	req := pdu.PingReq{Message: msg}
+	var ctrlOk, dataOk int32
+	loggers := GetLoggersOrPanic(ctx)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	checkRes := func(res *pdu.PingRes, err error, logger Logger, okVar *int32) {
+		if err == nil && res.GetEcho() != req.GetMessage() {
+			err = errors.New("pilot message not echoed correctly")
+		}
+		if err == context.Canceled {
+			err = nil
+		}
+		if err != nil {
+			logger.WithError(err).Error("ping failed")
+			atomic.StoreInt32(okVar, 0)
+			cancel()
+		} else {
+			atomic.StoreInt32(okVar, 1)
+		}
+	}
+	go func() {
+		defer wg.Done()
+		ctrl, ctrlErr := c.controlClient.Ping(ctx, &req, grpc.FailFast(false))
+		checkRes(ctrl, ctrlErr, loggers.Control, &ctrlOk)
+	}()
+	go func() {
+		defer wg.Done()
+		for ctx.Err() == nil {
+			data, dataErr := c.dataClient.ReqPing(ctx, &req)
+			// dataClient uses transport.Connecter, which doesn't expose FailFast(false)
+			// => we need to mask dial timeouts
+			if err, ok := dataErr.(interface{ Temporary() bool }); ok && err.Temporary() {
+				continue
+			}
+			// it's not a dial timeout, 
+			checkRes(data, dataErr, loggers.Data, &dataOk)
+			return
+		}
+	}()
+	wg.Wait()
+	var what string
+	if ctrlOk == 1 && dataOk == 1 {
+		return nil
+	}
+	if ctrlOk == 0 {
+		what += "control"
+	}
+	if dataOk == 0 {
+		if len(what) > 0 {
+			what += " and data"
+		} else {
+			what += "data"
+		}
+	}
+	return fmt.Errorf("%s rpc failed to respond to ping rpcs", what)
+}
+
 func (c *Client) ResetConnectBackoff() {
 	c.controlConn.ResetConnectBackoff()
 }