conn,task: add config about keepalive (pingcap#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
  }