*: enhance progress bar (pingcap#54)

Signed-off-by: Neil Shen <[email protected]>

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 == "" {

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
 			}

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)
 			}

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

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=

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()
+}

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))
+}