github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/task/restore_raw.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package task
     4  
     5  import (
     6  	"context"
     7  
     8  	"github.com/pingcap/br/pkg/metautil"
     9  
    10  	"github.com/pingcap/errors"
    11  	"github.com/pingcap/log"
    12  	"github.com/spf13/cobra"
    13  	"github.com/spf13/pflag"
    14  
    15  	berrors "github.com/pingcap/br/pkg/errors"
    16  	"github.com/pingcap/br/pkg/glue"
    17  	"github.com/pingcap/br/pkg/restore"
    18  	"github.com/pingcap/br/pkg/summary"
    19  )
    20  
    21  // RestoreRawConfig is the configuration specific for raw kv restore tasks.
    22  type RestoreRawConfig struct {
    23  	RawKvConfig
    24  	RestoreCommonConfig
    25  }
    26  
    27  // DefineRawRestoreFlags defines common flags for the backup command.
    28  func DefineRawRestoreFlags(command *cobra.Command) {
    29  	command.Flags().StringP(flagKeyFormat, "", "hex", "start/end key format, support raw|escaped|hex")
    30  	command.Flags().StringP(flagTiKVColumnFamily, "", "default", "restore specify cf, correspond to tikv cf")
    31  	command.Flags().StringP(flagStartKey, "", "", "restore raw kv start key, key is inclusive")
    32  	command.Flags().StringP(flagEndKey, "", "", "restore raw kv end key, key is exclusive")
    33  
    34  	DefineRestoreCommonFlags(command.PersistentFlags())
    35  }
    36  
    37  // ParseFromFlags parses the backup-related flags from the flag set.
    38  func (cfg *RestoreRawConfig) ParseFromFlags(flags *pflag.FlagSet) error {
    39  	var err error
    40  	cfg.Online, err = flags.GetBool(flagOnline)
    41  	if err != nil {
    42  		return errors.Trace(err)
    43  	}
    44  	err = cfg.RestoreCommonConfig.ParseFromFlags(flags)
    45  	if err != nil {
    46  		return errors.Trace(err)
    47  	}
    48  	return cfg.RawKvConfig.ParseFromFlags(flags)
    49  }
    50  
    51  func (cfg *RestoreRawConfig) adjust() {
    52  	cfg.Config.adjust()
    53  	cfg.RestoreCommonConfig.adjust()
    54  
    55  	if cfg.Concurrency == 0 {
    56  		cfg.Concurrency = defaultRestoreConcurrency
    57  	}
    58  }
    59  
    60  // RunRestoreRaw starts a raw kv restore task inside the current goroutine.
    61  func RunRestoreRaw(c context.Context, g glue.Glue, cmdName string, cfg *RestoreRawConfig) (err error) {
    62  	cfg.adjust()
    63  
    64  	defer summary.Summary(cmdName)
    65  	ctx, cancel := context.WithCancel(c)
    66  	defer cancel()
    67  
    68  	// Restore raw does not need domain.
    69  	needDomain := false
    70  	mgr, err := NewMgr(ctx, g, cfg.PD, cfg.TLS, GetKeepalive(&cfg.Config), cfg.CheckRequirements, needDomain)
    71  	if err != nil {
    72  		return errors.Trace(err)
    73  	}
    74  	defer mgr.Close()
    75  
    76  	keepaliveCfg := GetKeepalive(&cfg.Config)
    77  	// sometimes we have pooled the connections.
    78  	// sending heartbeats in idle times is useful.
    79  	keepaliveCfg.PermitWithoutStream = true
    80  	client, err := restore.NewRestoreClient(g, mgr.GetPDClient(), mgr.GetStorage(), mgr.GetTLSConfig(), keepaliveCfg)
    81  	if err != nil {
    82  		return errors.Trace(err)
    83  	}
    84  	defer client.Close()
    85  	client.SetRateLimit(cfg.RateLimit)
    86  	client.SetConcurrency(uint(cfg.Concurrency))
    87  	if cfg.Online {
    88  		client.EnableOnline()
    89  	}
    90  	client.SetSwitchModeInterval(cfg.SwitchModeInterval)
    91  
    92  	u, s, backupMeta, err := ReadBackupMeta(ctx, metautil.MetaFile, &cfg.Config)
    93  	if err != nil {
    94  		return errors.Trace(err)
    95  	}
    96  	reader := metautil.NewMetaReader(backupMeta, s)
    97  	if err = client.InitBackupMeta(c, backupMeta, u, s, reader); err != nil {
    98  		return errors.Trace(err)
    99  	}
   100  
   101  	if !client.IsRawKvMode() {
   102  		return errors.Annotate(berrors.ErrRestoreModeMismatch, "cannot do raw restore from transactional data")
   103  	}
   104  
   105  	files, err := client.GetFilesInRawRange(cfg.StartKey, cfg.EndKey, cfg.CF)
   106  	if err != nil {
   107  		return errors.Trace(err)
   108  	}
   109  	archiveSize := reader.ArchiveSize(ctx, files)
   110  	g.Record(summary.RestoreDataSize, archiveSize)
   111  
   112  	if len(files) == 0 {
   113  		log.Info("all files are filtered out from the backup archive, nothing to restore")
   114  		return nil
   115  	}
   116  	summary.CollectInt("restore files", len(files))
   117  
   118  	ranges, _, err := restore.MergeFileRanges(
   119  		files, cfg.MergeSmallRegionKeyCount, cfg.MergeSmallRegionKeyCount)
   120  	if err != nil {
   121  		return errors.Trace(err)
   122  	}
   123  
   124  	// Redirect to log if there is no log file to avoid unreadable output.
   125  	// TODO: How to show progress?
   126  	updateCh := g.StartProgress(
   127  		ctx,
   128  		"Raw Restore",
   129  		// Split/Scatter + Download/Ingest
   130  		int64(len(ranges)+len(files)),
   131  		!cfg.LogProgress)
   132  
   133  	// RawKV restore does not need to rewrite keys.
   134  	rewrite := &restore.RewriteRules{}
   135  	err = restore.SplitRanges(ctx, client, ranges, rewrite, updateCh)
   136  	if err != nil {
   137  		return errors.Trace(err)
   138  	}
   139  
   140  	restoreSchedulers, err := restorePreWork(ctx, client, mgr)
   141  	if err != nil {
   142  		return errors.Trace(err)
   143  	}
   144  	defer restorePostWork(ctx, client, restoreSchedulers)
   145  
   146  	err = client.RestoreRaw(ctx, cfg.StartKey, cfg.EndKey, files, updateCh)
   147  	if err != nil {
   148  		return errors.Trace(err)
   149  	}
   150  
   151  	// Restore has finished.
   152  	updateCh.Close()
   153  
   154  	// Set task summary to success status.
   155  	summary.SetSuccessStatus(true)
   156  	return nil
   157  }