conn,task: add config about keepalive (#545)

* conn,task: add config about keepalive

* *: fix CI

* task: hide grpc lifetime flags

* *: fix a typo

* task: add adjust to grpc keepalive variables

cmd/debug.go

@@ -357,7 +357,7 @@ func setPDConfigCommand() *cobra.Command {
 				return err
 			}
 
-			mgr, err := task.NewMgr(ctx, tidbGlue, cfg.PD, cfg.TLS, cfg.CheckRequirements)
+			mgr, err := task.NewMgr(ctx, tidbGlue, cfg.PD, cfg.TLS, task.GetKeepalive(&cfg), cfg.CheckRequirements)
 			if err != nil {
 				return err
 			}

pkg/conn/conn.go

@@ -44,6 +44,7 @@ type Mgr struct {
 		mu   sync.Mutex
 		clis map[uint64]*grpc.ClientConn
 	}
+	keepalive   keepalive.ClientParameters
 	ownsStorage bool
 }
 
@@ -109,6 +110,7 @@ func NewMgr(
 	storage tikv.Storage,
 	tlsConf *tls.Config,
 	securityOption pd.SecurityOption,
+	keepalive keepalive.ClientParameters,
 	storeBehavior StoreBehavior,
 	checkRequirements bool,
 ) (*Mgr, error) {
@@ -159,6 +161,7 @@ func NewMgr(
 		ownsStorage:  g.OwnsStorage(),
 	}
 	mgr.grpcClis.clis = make(map[uint64]*grpc.ClientConn)
+	mgr.keepalive = keepalive
 	return mgr, nil
 }
 
@@ -172,8 +175,6 @@ func (mgr *Mgr) getGrpcConnLocked(ctx context.Context, storeID uint64) (*grpc.Cl
 		opt = grpc.WithTransportCredentials(credentials.NewTLS(mgr.tlsConf))
 	}
 	ctx, cancel := context.WithTimeout(ctx, dialTimeout)
-	keepAlive := 10
-	keepAliveTimeout := 3
 	bfConf := backoff.DefaultConfig
 	bfConf.MaxDelay = time.Second * 3
 	addr := store.GetPeerAddress()
@@ -186,10 +187,7 @@ func (mgr *Mgr) getGrpcConnLocked(ctx context.Context, storeID uint64) (*grpc.Cl
 		opt,
 		grpc.WithBlock(),
 		grpc.WithConnectParams(grpc.ConnectParams{Backoff: bfConf}),
-		grpc.WithKeepaliveParams(keepalive.ClientParameters{
-			Time:    time.Duration(keepAlive) * time.Second,
-			Timeout: time.Duration(keepAliveTimeout) * time.Second,
-		}),
+		grpc.WithKeepaliveParams(mgr.keepalive),
 	)
 	cancel()
 	if err != nil {

pkg/task/backup.go

@@ -148,6 +148,7 @@ func parseCompressionFlags(flags *pflag.FlagSet) (*CompressionConfig, error) {
 // we should set proper value in this function.
 // so that both binary and TiDB will use same default value.
 func (cfg *BackupConfig) adjustBackupConfig() {
+	cfg.adjust()
 	if cfg.Config.Concurrency == 0 {
 		cfg.Config.Concurrency = defaultBackupConcurrency
 	}
@@ -176,7 +177,7 @@ func RunBackup(c context.Context, g glue.Glue, cmdName string, cfg *BackupConfig
 	if err != nil {
 		return err
 	}
-	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, cfg.CheckRequirements)
+	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(&cfg.Config), cfg.CheckRequirements)
 	if err != nil {
 		return err
 	}

pkg/task/backup_raw.go

@@ -118,6 +118,8 @@ func (cfg *RawKvConfig) ParseBackupConfigFromFlags(flags *pflag.FlagSet) error {
 
 // RunBackupRaw starts a backup task inside the current goroutine.
 func RunBackupRaw(c context.Context, g glue.Glue, cmdName string, cfg *RawKvConfig) error {
+	cfg.adjust()
+
 	defer summary.Summary(cmdName)
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
@@ -126,7 +128,7 @@ func RunBackupRaw(c context.Context, g glue.Glue, cmdName string, cfg *RawKvConf
 	if err != nil {
 		return err
 	}
-	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, cfg.CheckRequirements)
+	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(&cfg.Config), cfg.CheckRequirements)
 	if err != nil {
 		return err
 	}

pkg/task/common.go

@@ -21,6 +21,7 @@ import (
 	pd "github.com/tikv/pd/client"
 	"go.etcd.io/etcd/pkg/transport"
 	"go.uber.org/zap"
+	"google.golang.org/grpc/keepalive"
 
 	"github.com/pingcap/br/pkg/conn"
 	berrors "github.com/pingcap/br/pkg/errors"
@@ -55,8 +56,14 @@ const (
 	flagRemoveTiFlash      = "remove-tiflash"
 	flagCheckRequirement   = "check-requirements"
 	flagSwitchModeInterval = "switch-mode-interval"
-
-	defaultSwitchInterval = 5 * time.Minute
+	// flagGrpcKeepaliveTime is the interval of pinging the server.
+	flagGrpcKeepaliveTime = "grpc-keepalive-time"
+	// flagGrpcKeepaliveTimeout is the max time a grpc conn can keep idel before killed.
+	flagGrpcKeepaliveTimeout = "grpc-keepalive-timeout"
+
+	defaultSwitchInterval       = 5 * time.Minute
+	defaultGRPCKeepaliveTime    = 10 * time.Second
+	defaultGRPCKeepaliveTimeout = 3 * time.Second
 )
 
 // TLSConfig is the common configuration for TLS connection.
@@ -113,6 +120,11 @@ type Config struct {
 	TableFilter        filter.Filter `json:"-" toml:"-"`
 	CheckRequirements  bool          `json:"check-requirements" toml:"check-requirements"`
 	SwitchModeInterval time.Duration `json:"switch-mode-interval" toml:"switch-mode-interval"`
+
+	// GrpcKeepaliveTime is the interval of pinging the server.
+	GRPCKeepaliveTime time.Duration `json:"grpc-keepalive-time" toml:"grpc-keepalive-time"`
+	// GrpcKeepaliveTimeout is the max time a grpc conn can keep idel before killed.
+	GRPCKeepaliveTimeout time.Duration `json:"grpc-keepalive-timeout" toml:"grpc-keepalive-timeout"`
 }
 
 // DefineCommonFlags defines the flags common to all BRIE commands.
@@ -143,6 +155,12 @@ func DefineCommonFlags(flags *pflag.FlagSet) {
 	flags.Bool(flagCheckRequirement, true,
 		"Whether start version check before execute command")
 	flags.Duration(flagSwitchModeInterval, defaultSwitchInterval, "maintain import mode on TiKV during restore")
+	flags.Duration(flagGrpcKeepaliveTime, defaultGRPCKeepaliveTime,
+		"the interval of pinging gRPC peer, must keep the same value with TiKV and PD")
+	flags.Duration(flagGrpcKeepaliveTimeout, defaultGRPCKeepaliveTimeout,
+		"the max time a gRPC connection can keep idle before killed, must keep the same value with TiKV and PD")
+	_ = flags.MarkHidden(flagGrpcKeepaliveTime)
+	_ = flags.MarkHidden(flagGrpcKeepaliveTimeout)
 
 	storage.DefineFlags(flags)
 }
@@ -267,6 +285,14 @@ func (cfg *Config) ParseFromFlags(flags *pflag.FlagSet) error {
 	if err != nil {
 		return errors.Trace(err)
 	}
+	cfg.GRPCKeepaliveTime, err = flags.GetDuration(flagGrpcKeepaliveTime)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	cfg.GRPCKeepaliveTimeout, err = flags.GetDuration(flagGrpcKeepaliveTimeout)
+	if err != nil {
+		return errors.Trace(err)
+	}
 
 	if cfg.SwitchModeInterval <= 0 {
 		return errors.Annotatef(berrors.ErrInvalidArgument, "--switch-mode-interval must be positive, %s is not allowed", cfg.SwitchModeInterval)
@@ -282,6 +308,7 @@ func (cfg *Config) ParseFromFlags(flags *pflag.FlagSet) error {
 func NewMgr(ctx context.Context,
 	g glue.Glue, pds []string,
 	tlsConfig TLSConfig,
+	keepalive keepalive.ClientParameters,
 	checkRequirements bool) (*conn.Mgr, error) {
 	var (
 		tlsConf *tls.Config
@@ -312,7 +339,7 @@ func NewMgr(ctx context.Context,
 	// Is it necessary to remove `StoreBehavior`?
 	return conn.NewMgr(ctx, g,
 		pdAddress, store.(tikv.Storage),
-		tlsConf, securityOption,
+		tlsConf, securityOption, keepalive,
 		conn.SkipTiFlash, checkRequirements)
 }
 
@@ -378,3 +405,22 @@ func LogArguments(cmd *cobra.Command) {
 	})
 	log.Info("arguments", fields...)
 }
+
+// GetKeepalive get the keepalive info from the config.
+func GetKeepalive(cfg *Config) keepalive.ClientParameters {
+	return keepalive.ClientParameters{
+		Time:    cfg.GRPCKeepaliveTime,
+		Timeout: cfg.GRPCKeepaliveTimeout,
+	}
+}
+
+// adjust adjusts the abnormal config value in the current config.
+// useful when not starting BR from CLI (e.g. from BRIE in SQL).
+func (cfg *Config) adjust() {
+	if cfg.GRPCKeepaliveTime == 0 {
+		cfg.GRPCKeepaliveTime = defaultGRPCKeepaliveTime
+	}
+	if cfg.GRPCKeepaliveTimeout == 0 {
+		cfg.GRPCKeepaliveTimeout = defaultGRPCKeepaliveTimeout
+	}
+}

pkg/task/restore.go

@@ -78,6 +78,8 @@ func (cfg *RestoreConfig) ParseFromFlags(flags *pflag.FlagSet) error {
 // we should set proper value in this function.
 // so that both binary and TiDB will use same default value.
 func (cfg *RestoreConfig) adjustRestoreConfig() {
+	cfg.adjust()
+
 	if cfg.Config.Concurrency == 0 {
 		cfg.Config.Concurrency = defaultRestoreConcurrency
 	}
@@ -94,7 +96,7 @@ func RunRestore(c context.Context, g glue.Glue, cmdName string, cfg *RestoreConf
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
 
-	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, cfg.CheckRequirements)
+	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(&cfg.Config), cfg.CheckRequirements)
 	if err != nil {
 		return err
 	}
@@ -370,7 +372,7 @@ func RunRestoreTiflashReplica(c context.Context, g glue.Glue, cmdName string, cf
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
 
-	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, cfg.CheckRequirements)
+	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(&cfg.Config), cfg.CheckRequirements)
 	if err != nil {
 		return err
 	}

pkg/task/restore_log.go

@@ -72,6 +72,8 @@ func (cfg *LogRestoreConfig) ParseFromFlags(flags *pflag.FlagSet) error {
 // we should set proper value in this function.
 // so that both binary and TiDB will use same default value.
 func (cfg *LogRestoreConfig) adjustRestoreConfig() {
+	cfg.adjust()
+
 	if cfg.Config.Concurrency == 0 {
 		cfg.Config.Concurrency = defaultRestoreConcurrency
 	}
@@ -97,7 +99,7 @@ func RunLogRestore(c context.Context, g glue.Glue, cfg *LogRestoreConfig) error
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
 
-	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, cfg.CheckRequirements)
+	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(&cfg.Config), cfg.CheckRequirements)
 	if err != nil {
 		return err
 	}

pkg/task/restore_raw.go

@@ -47,6 +47,8 @@ func (cfg *RestoreRawConfig) ParseFromFlags(flags *pflag.FlagSet) error {
 }
 
 func (cfg *RestoreRawConfig) adjust() {
+	cfg.Config.adjust()
+
 	if cfg.Concurrency == 0 {
 		cfg.Concurrency = defaultRestoreConcurrency
 	}
@@ -60,7 +62,7 @@ func RunRestoreRaw(c context.Context, g glue.Glue, cmdName string, cfg *RestoreR
 	ctx, cancel := context.WithCancel(c)
 	defer cancel()
 
-	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, cfg.CheckRequirements)
+	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(&cfg.Config), cfg.CheckRequirements)
 	if err != nil {
 		return err
 	}