From b55c7e169a9f6190229fc992ed50d1235d3bb105 Mon Sep 17 00:00:00 2001 From: Neil Shen Date: Fri, 15 Nov 2019 14:13:20 +0800 Subject: [PATCH] *: enhance progress bar (#54) Signed-off-by: Neil Shen --- cmd/cmd.go | 10 ++++++ cmd/raw.go | 17 ++++----- cmd/restore.go | 30 ++++++++-------- go.mod | 2 +- go.sum | 4 +-- pkg/utils/progress.go | 71 ++++++++++++++++++++++++++++++++------ pkg/utils/progress_test.go | 43 +++++++++++++++++++++++ 7 files changed, 139 insertions(+), 38 deletions(-) create mode 100644 pkg/utils/progress_test.go diff --git a/cmd/cmd.go b/cmd/cmd.go index 70c0763801323..bb1cd5c52b988 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -5,6 +5,7 @@ import ( "net/http" "net/http/pprof" "sync" + "sync/atomic" "github.com/pingcap/errors" "github.com/pingcap/log" @@ -20,6 +21,7 @@ var ( initOnce = sync.Once{} defaultContext context.Context pdAddress string + hasLogFile uint64 backerOnce = sync.Once{} defaultBacker *meta.Backer @@ -79,6 +81,9 @@ func Init(cmd *cobra.Command) (err error) { if err != nil { return } + if len(conf.File.Filename) != 0 { + atomic.StoreUint64(&hasLogFile, 1) + } lg, p, e := log.InitLogger(conf) if e != nil { err = e @@ -129,6 +134,11 @@ func Init(cmd *cobra.Command) (err error) { return err } +// HasLogFile returns whether we set a log file +func HasLogFile() bool { + return atomic.LoadUint64(&hasLogFile) != uint64(0) +} + // GetDefaultBacker returns the default backer for command line usage. func GetDefaultBacker() (*meta.Backer, error) { if pdAddress == "" { diff --git a/cmd/raw.go b/cmd/raw.go index 944b7773cc8e3..a4761bb0e042d 100644 --- a/cmd/raw.go +++ b/cmd/raw.go @@ -106,14 +106,14 @@ func newFullBackupCommand() *cobra.Command { return err } - progress := utils.NewProgressPrinter( - "Full Backup", int64(approximateRegions)) ctx, cancel := context.WithCancel(defaultBacker.Context()) defer cancel() - progress.GoPrintProgress(ctx) + // Redirect to log if there is no log file to avoid unreadable output. + updateCh := utils.StartProgress( + ctx, "Full Backup", int64(approximateRegions), !HasLogFile()) err = client.BackupRanges( - ranges, u, backupTS, rate, concurrency, progress.UpdateCh()) + ranges, u, backupTS, rate, concurrency, updateCh) if err != nil { return err } @@ -208,6 +208,7 @@ func newTableBackupCommand() *cobra.Command { return errors.New("at least one thread required") } + // TODO: include admin check in progress bar. ranges, err := client.PreBackupTableRanges(db, table, u, backupTS) if err != nil { return err @@ -223,14 +224,14 @@ func newTableBackupCommand() *cobra.Command { approximateRegions += regionCount } - progress := utils.NewProgressPrinter( - "Table Backup", int64(approximateRegions)) ctx, cancel := context.WithCancel(defaultBacker.Context()) defer cancel() - progress.GoPrintProgress(ctx) + // Redirect to log if there is no log file to avoid unreadable output. + updateCh := utils.StartProgress( + ctx, "Table Backup", int64(approximateRegions), !HasLogFile()) err = client.BackupRanges( - ranges, u, backupTS, rate, concurrency, progress.UpdateCh()) + ranges, u, backupTS, rate, concurrency, updateCh) if err != nil { return err } diff --git a/cmd/restore.go b/cmd/restore.go index 0372eb2396fd8..c7ad7354f3f54 100644 --- a/cmd/restore.go +++ b/cmd/restore.go @@ -83,13 +83,13 @@ func newFullRestoreCommand() *cobra.Command { } ranges := restore.GetRanges(files) - progress := utils.NewProgressPrinter( + // Redirect to log if there is no log file to avoid unreadable output. + updateCh := utils.StartProgress( + ctx, "Full Restore", // Split/Scatter + Download/Ingest int64(len(ranges)+len(files)), - ) - progress.GoPrintProgress(ctx) - updateCh := progress.UpdateCh() + !HasLogFile()) rewriteRules := &restore_util.RewriteRules{ Table: tableRules, @@ -109,8 +109,7 @@ func newFullRestoreCommand() *cobra.Command { return errors.Trace(err) } - err = client.RestoreAll( - rewriteRules, progress.UpdateCh()) + err = client.RestoreAll(rewriteRules, updateCh) if err != nil { return errors.Trace(err) } @@ -178,13 +177,13 @@ func newDbRestoreCommand() *cobra.Command { } ranges := restore.GetRanges(files) - progress := utils.NewProgressPrinter( + // Redirect to log if there is no log file to avoid unreadable output. + updateCh := utils.StartProgress( + ctx, "Database Restore", // Split/Scatter + Download/Ingest int64(len(ranges)+len(files)), - ) - progress.GoPrintProgress(ctx) - updateCh := progress.UpdateCh() + !HasLogFile()) err = restore.SplitRanges(ctx, client, ranges, rewriteRules, updateCh) if err != nil { @@ -279,13 +278,13 @@ func newTableRestoreCommand() *cobra.Command { } ranges := restore.GetRanges(table.Files) - progress := utils.NewProgressPrinter( + // Redirect to log if there is no log file to avoid unreadable output. + updateCh := utils.StartProgress( + ctx, "Table Restore", // Split/Scatter + Download/Ingest int64(len(ranges)+len(table.Files)), - ) - progress.GoPrintProgress(ctx) - updateCh := progress.UpdateCh() + !HasLogFile()) err = restore.SplitRanges(ctx, client, ranges, rewriteRules, updateCh) if err != nil { @@ -299,8 +298,7 @@ func newTableRestoreCommand() *cobra.Command { if err != nil { return errors.Trace(err) } - err = client.RestoreTable( - table, rewriteRules, progress.UpdateCh()) + err = client.RestoreTable(table, rewriteRules, updateCh) if err != nil { return errors.Trace(err) } diff --git a/go.mod b/go.mod index 1a270d0e3edd0..f768761cdf597 100644 --- a/go.mod +++ b/go.mod @@ -20,7 +20,7 @@ require ( github.com/pingcap/parser v0.0.0-20191011160321-0c4055ef2c1d github.com/pingcap/pd v1.1.0-beta.0.20191108030828-0e3a054daa85 github.com/pingcap/tidb v1.1.0-beta.0.20191108051109-35f3653fc601 - github.com/pingcap/tidb-tools v3.0.6-0.20191108032853-4bb0acca5cc5+incompatible + github.com/pingcap/tidb-tools v3.0.6-0.20191113022349-48d5e90d3271+incompatible github.com/pingcap/tipb v0.0.0-20191101114505-cbd0e985c780 github.com/prometheus/client_golang v1.0.0 github.com/shurcooL/httpfs v0.0.0-20190707220628-8d4bc4ba7749 // indirect diff --git a/go.sum b/go.sum index be8a8a8baf0c5..472bcc949b21b 100644 --- a/go.sum +++ b/go.sum @@ -233,8 +233,8 @@ github.com/pingcap/pd v1.1.0-beta.0.20191108030828-0e3a054daa85/go.mod h1:ribyi6 github.com/pingcap/tidb v1.1.0-beta.0.20191108051109-35f3653fc601 h1:lzRmKMBJbh4USjVbgwlYgacz5kphaq6SUB5hlVNsn3U= github.com/pingcap/tidb v1.1.0-beta.0.20191108051109-35f3653fc601/go.mod h1:1fZBjQ5znUHrJK2BQEipjyJQOqZGkESaH4zZtJDYCUw= github.com/pingcap/tidb-tools v3.0.6-0.20191106033616-90632dda3863+incompatible/go.mod h1:XGdcy9+yqlDSEMTpOXnwf3hiTeqrV6MN/u1se9N8yIM= -github.com/pingcap/tidb-tools v3.0.6-0.20191108032853-4bb0acca5cc5+incompatible h1:lQgNM4wEOcTVIBVOYduJd4n+Y3W2El381AkN4nVLm3M= -github.com/pingcap/tidb-tools v3.0.6-0.20191108032853-4bb0acca5cc5+incompatible/go.mod h1:XGdcy9+yqlDSEMTpOXnwf3hiTeqrV6MN/u1se9N8yIM= +github.com/pingcap/tidb-tools v3.0.6-0.20191113022349-48d5e90d3271+incompatible h1:/7AeCi01XT4RydwHxtAFrcbZM56aL/sW0o4A1i89Krg= +github.com/pingcap/tidb-tools v3.0.6-0.20191113022349-48d5e90d3271+incompatible/go.mod h1:XGdcy9+yqlDSEMTpOXnwf3hiTeqrV6MN/u1se9N8yIM= github.com/pingcap/tipb v0.0.0-20191101114505-cbd0e985c780 h1:SvFkjLhS/ou97Ey60r8Fq3ZF4wq6wuveWoiLtWLGpek= github.com/pingcap/tipb v0.0.0-20191101114505-cbd0e985c780/go.mod h1:RtkHW8WbcNxj8lsbzjaILci01CtYnYbIkQhjyZWrWVI= github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= diff --git a/pkg/utils/progress.go b/pkg/utils/progress.go index bc6be87e7b38f..095678eb969b7 100644 --- a/pkg/utils/progress.go +++ b/pkg/utils/progress.go @@ -2,25 +2,34 @@ package utils import ( "context" + "io" "time" "github.com/cheggaaa/pb/v3" + "github.com/pingcap/log" + "go.uber.org/zap" ) // ProgressPrinter prints a progress bar type ProgressPrinter struct { - name string - total int64 + name string + total int64 + redirectLog bool updateCh chan struct{} } // NewProgressPrinter returns a new progress printer -func NewProgressPrinter(name string, total int64) *ProgressPrinter { +func NewProgressPrinter( + name string, + total int64, + redirectLog bool, +) *ProgressPrinter { return &ProgressPrinter{ - name: name, - total: total, - updateCh: make(chan struct{}, total/2), + name: name, + total: total, + redirectLog: redirectLog, + updateCh: make(chan struct{}, total/2), } } @@ -29,11 +38,30 @@ func (pp *ProgressPrinter) UpdateCh() chan<- struct{} { return pp.updateCh } -// GoPrintProgress starts a gorouinte and prints progress -func (pp *ProgressPrinter) GoPrintProgress(ctx context.Context) { - tmpl := `{{string . "barName" | red}} {{ bar . "<" "-" (cycle . "↖" "↗" "↘" "↙" ) "." ">"}} {{percent .}}` - bar := pb.ProgressBarTemplate(tmpl).Start64(pp.total) - bar.Set("barName", pp.name) +// goPrintProgress starts a gorouinte and prints progress +func (pp *ProgressPrinter) goPrintProgress( + ctx context.Context, + testWriter io.Writer, // Only for tests +) { + var bar *pb.ProgressBar + if pp.redirectLog || testWriter != nil { + tmpl := `{{percent .}}` + bar = pb.ProgressBarTemplate(tmpl).Start64(pp.total) + bar.SetRefreshRate(time.Second * 10) + bar.Set(pb.Static, false) // Do not update automatically + bar.Set(pb.ReturnSymbol, false) // Do not append '\r' + bar.Set(pb.Terminal, false) // Do not use terminal width + // Hack! set Color to avoid separate progress string + bar.Set(pb.Color, true) + bar.SetWriter(&wrappedWriter{name: pp.name}) + } else { + tmpl := `{{string . "barName" | red}} {{ bar . "<" "-" (cycle . "-" "\\" "|" "/" ) "." ">"}} {{percent .}}` + bar = pb.ProgressBarTemplate(tmpl).Start64(pp.total) + bar.Set("barName", pp.name) + } + if testWriter != nil { + bar.SetWriter(testWriter) + } bar.Start() go func() { @@ -60,3 +88,24 @@ func (pp *ProgressPrinter) GoPrintProgress(ctx context.Context) { } }() } + +type wrappedWriter struct { + name string +} + +func (ww *wrappedWriter) Write(p []byte) (int, error) { + log.Info(ww.name, zap.String("progress", string(p))) + return len(p), nil +} + +// StartProgress starts progress bar. +func StartProgress( + ctx context.Context, + name string, + total int64, + redirectLog bool, +) chan<- struct{} { + progress := NewProgressPrinter(name, total, redirectLog) + progress.goPrintProgress(ctx, nil) + return progress.UpdateCh() +} diff --git a/pkg/utils/progress_test.go b/pkg/utils/progress_test.go new file mode 100644 index 0000000000000..9bfc9e31dda32 --- /dev/null +++ b/pkg/utils/progress_test.go @@ -0,0 +1,43 @@ +package utils + +import ( + "context" + "strings" + + . "github.com/pingcap/check" +) + +type testProgressSuite struct{} + +var _ = Suite(&testProgressSuite{}) + +type testWriter struct { + fn func(string) +} + +func (t *testWriter) Write(p []byte) (int, error) { + t.fn(string(p)) + return len(p), nil +} + +func (r *testProgressSuite) TestProgress(c *C) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var p string + pCh := make(chan string, 2) + progress := NewProgressPrinter("test", 2, false) + progress.goPrintProgress(ctx, &testWriter{ + fn: func(p string) { pCh <- p }, + }) + updateCh := progress.UpdateCh() + updateCh <- struct{}{} + p = <-pCh + c.Assert(strings.Contains(p, "50"), IsTrue, Commentf("%s", p)) + updateCh <- struct{}{} + p = <-pCh + c.Assert(strings.Contains(p, "100"), IsTrue, Commentf("%s", p)) + updateCh <- struct{}{} + p = <-pCh + c.Assert(strings.Contains(p, "100"), IsTrue, Commentf("%s", p)) +}