br/cmd: implement the operator pause-gc-and-schedulers command

br/cmd/br/BUILD.bazel

@@ -7,6 +7,7 @@ go_library(
         "cmd.go",
         "debug.go",
         "main.go",
+        "operator.go",
         "restore.go",
         "stream.go",
     ],
@@ -26,6 +27,7 @@ go_library(
         "//br/pkg/streamhelper/config",
         "//br/pkg/summary",
         "//br/pkg/task",
+        "//br/pkg/task/operator",
         "//br/pkg/trace",
         "//br/pkg/utils",
         "//br/pkg/version/build",

br/cmd/br/main.go

@@ -49,6 +49,7 @@ func main() {
 		NewBackupCommand(),
 		NewRestoreCommand(),
 		NewStreamCommand(),
+		newOpeartorCommand(),
 	)
 	// Outputs cmd.Print to stdout.
 	rootCmd.SetOut(os.Stdout)

br/cmd/br/operator.go

@@ -0,0 +1,49 @@
+// Copyright 2023 PingCAP, Inc. Licensed under Apache-2.0.
+
+package main
+
+import (
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/br/pkg/task"
+	"github.com/pingcap/tidb/br/pkg/task/operator"
+	"github.com/pingcap/tidb/br/pkg/utils"
+	"github.com/pingcap/tidb/br/pkg/version/build"
+	"github.com/spf13/cobra"
+)
+
+func newOpeartorCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "operator <subcommand>",
+		Short: "utilities for operators like tidb-operator.",
+		PersistentPreRunE: func(c *cobra.Command, args []string) error {
+			if err := Init(c); err != nil {
+				return errors.Trace(err)
+			}
+			build.LogInfo(build.BR)
+			utils.LogEnvVariables()
+			task.LogArguments(c)
+			return nil
+		},
+		Hidden: true,
+	}
+	cmd.AddCommand(newPauseGcCommand())
+	return cmd
+}
+
+func newPauseGcCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "pause-gc-and-schedulers",
+		Short: "pause gc and schedulers to the ts until the program exits.",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			cfg := operator.PauseGcConfig{}
+			if err := cfg.ParseFromFlags(cmd.Flags()); err != nil {
+				return err
+			}
+			ctx := GetDefaultContext()
+			return operator.PauseGCAndScheduler(ctx, &cfg)
+		},
+	}
+	operator.DefineFlagsForPauseGcConfig(cmd.Flags())
+	return cmd
+}

br/pkg/pdutil/pd.go

@@ -277,7 +277,7 @@ func NewPdController(
 	}
 	if failure != nil {
 		return nil, errors.Annotatef(berrors.ErrPDUpdateFailed,
-			"pd address (%s) not available, please check network", pdAddrs)
+			"pd address (%s) not available, error is %s, please check network", pdAddrs, failure)
 	}
 
 	version := parseVersion(versionBytes)

br/pkg/task/common.go

@@ -138,6 +138,15 @@ func (tls *TLSConfig) ToTLSConfig() (*tls.Config, error) {
 	return tlsConfig, nil
 }
 
+// Convert the TLS config to the PD security option.
+func (tls *TLSConfig) ToPDSecurityOption() pd.SecurityOption {
+	securityOption := pd.SecurityOption{}
+	securityOption.CAPath = tls.CA
+	securityOption.CertPath = tls.Cert
+	securityOption.KeyPath = tls.Key
+	return securityOption
+}
+
 // ParseFromFlags parses the TLS config from the flag set.
 func (tls *TLSConfig) ParseFromFlags(flags *pflag.FlagSet) (err error) {
 	tls.CA, tls.Cert, tls.Key, err = ParseTLSTripleFromFlags(flags)

br/pkg/task/operator/BUILD.bazel

@@ -0,0 +1,21 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+go_library(
+    name = "operator",
+    srcs = [
+        "cmd.go",
+        "config.go",
+    ],
+    importpath = "github.com/pingcap/tidb/br/pkg/task/operator",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//br/pkg/logutil",
+        "//br/pkg/pdutil",
+        "//br/pkg/task",
+        "//br/pkg/utils",
+        "@com_github_pingcap_log//:log",
+        "@com_github_spf13_pflag//:pflag",
+        "@org_golang_x_sync//errgroup",
+        "@org_uber_go_zap//:zap",
+    ],
+)

br/pkg/task/operator/cmd.go

@@ -0,0 +1,103 @@
+// Copyright 2023 PingCAP, Inc. Licensed under Apache-2.0.
+
+package operator
+
+import (
+	"context"
+	"crypto/tls"
+	"strings"
+	"time"
+
+	"github.com/pingcap/log"
+	"github.com/pingcap/tidb/br/pkg/logutil"
+	"github.com/pingcap/tidb/br/pkg/pdutil"
+	"github.com/pingcap/tidb/br/pkg/task"
+	"github.com/pingcap/tidb/br/pkg/utils"
+	"go.uber.org/zap"
+	"golang.org/x/sync/errgroup"
+)
+
+func dialPD(ctx context.Context, cfg *task.Config) (*pdutil.PdController, error) {
+	pdAddrs := strings.Join(cfg.PD, ",")
+	var tc *tls.Config
+	if cfg.TLS.IsEnabled() {
+		var err error
+		tc, err = cfg.TLS.ToTLSConfig()
+		if err != nil {
+			return nil, err
+		}
+	}
+	mgr, err := pdutil.NewPdController(ctx, pdAddrs, tc, cfg.TLS.ToPDSecurityOption())
+	if err != nil {
+		return nil, err
+	}
+	return mgr, nil
+}
+
+func cleanUpWith(f func(ctx context.Context)) {
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	f(ctx)
+}
+
+// PauseGCAndScheduler blocks the current goroutine and pause the GC safepoint and remove the scheduler by the config.
+// This function will block until the context being canceled.
+func PauseGCAndScheduler(ctx context.Context, cfg *PauseGcConfig) error {
+	mgr, err := dialPD(ctx, &cfg.Config)
+	if err != nil {
+		return err
+	}
+
+	eg, ectx := errgroup.WithContext(ctx)
+
+	eg.Go(func() error { return pauseGCKeeper(ectx, cfg, mgr) })
+	eg.Go(func() error { return pauseSchedulerKeeper(ectx, mgr) })
+
+	return eg.Wait()
+}
+
+func pauseGCKeeper(ctx context.Context, cfg *PauseGcConfig, ctl *pdutil.PdController) error {
+	// Note: should we remove the service safepoint as soon as this exits?
+	sp := utils.BRServiceSafePoint{
+		ID:       utils.MakeSafePointID(),
+		TTL:      int64(cfg.TTL.Seconds()),
+		BackupTS: cfg.SafePoint,
+	}
+	if sp.BackupTS == 0 {
+		rts, err := ctl.GetMinResolvedTS(ctx)
+		if err != nil {
+			return err
+		}
+		log.Info("No service safepoint provided, using the minimal resolved TS.", zap.Uint64("min-resolved-ts", rts))
+		sp.BackupTS = rts
+	}
+	err := utils.StartServiceSafePointKeeper(ctx, ctl.GetPDClient(), sp)
+	if err != nil {
+		return err
+	}
+	log.Info("GC is paused.", zap.Object("safepoint", sp))
+	// Note: in fact we can directly return here.
+	// But the name `keeper` implies once the function exits,
+	// the GC should be resume, so let's block here.
+	<-ctx.Done()
+	return nil
+}
+
+func pauseSchedulerKeeper(ctx context.Context, ctl *pdutil.PdController) error {
+	undo, err := ctl.RemoveAllPDSchedulers(ctx)
+	if undo != nil {
+		defer cleanUpWith(func(ctx context.Context) {
+			if err := undo(ctx); err != nil {
+				log.Warn("failed to restore pd scheduler.", logutil.ShortError(err))
+			}
+		})
+	}
+	if err != nil {
+		return err
+	}
+	log.Info("Schedulers are paused.")
+	// Wait until the context canceled.
+	// So we can properly do the clean up work.
+	<-ctx.Done()
+	return nil
+}

br/pkg/task/operator/config.go

@@ -0,0 +1,41 @@
+// Copyright 2023 PingCAP, Inc. Licensed under Apache-2.0.
+
+package operator
+
+import (
+	"time"
+
+	"github.com/pingcap/tidb/br/pkg/task"
+	"github.com/spf13/pflag"
+)
+
+type PauseGcConfig struct {
+	task.Config
+
+	SafePoint uint64        `json:"safepoint" yaml:"safepoint"`
+	TTL       time.Duration `json:"ttl" yaml:"ttl"`
+}
+
+func DefineFlagsForPauseGcConfig(f *pflag.FlagSet) {
+	_ = f.DurationP("ttl", "i", 5*time.Minute, "The time-to-live of the safepoint.")
+	_ = f.Uint64P("safepoint", "t", 0, "The GC safepoint to be kept.")
+}
+
+// ParseFromFlags fills the config via the flags.
+func (cfg *PauseGcConfig) ParseFromFlags(flags *pflag.FlagSet) error {
+	if err := cfg.Config.ParseFromFlags(flags); err != nil {
+		return err
+	}
+
+	var err error
+	cfg.SafePoint, err = flags.GetUint64("safepoint")
+	if err != nil {
+		return err
+	}
+	cfg.TTL, err = flags.GetDuration("ttl")
+	if err != nil {
+		return err
+	}
+
+	return nil
+}