-
Notifications
You must be signed in to change notification settings - Fork 5.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
This is an automated cherry-pick of #41670
Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>
- Loading branch information
1 parent
113e2cc
commit 13adb2c
Showing
13 changed files
with
1,071 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,169 @@ | ||
| // Copyright 2022 PingCAP, Inc. Licensed under Apache-2.0. | ||
|
|
||
| package glue | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "io" | ||
| "os" | ||
| "time" | ||
|
|
||
| "github.com/fatih/color" | ||
| "github.com/pingcap/tidb/br/pkg/utils" | ||
| "github.com/vbauerster/mpb/v7" | ||
| "github.com/vbauerster/mpb/v7/decor" | ||
| "golang.org/x/term" | ||
| ) | ||
|
|
||
| const OnlyOneTask int = -1 | ||
|
|
||
| var spinnerText []string = []string{".", "..", "..."} | ||
|
|
||
| type pbProgress struct { | ||
| bar *mpb.Bar | ||
| progress *mpb.Progress | ||
| ops ConsoleOperations | ||
| } | ||
|
|
||
| // Inc increases the progress. This method must be goroutine-safe, and can | ||
| // be called from any goroutine. | ||
| func (p pbProgress) Inc() { | ||
| p.bar.Increment() | ||
| } | ||
|
|
||
| // IncBy increases the progress by n. | ||
| func (p pbProgress) IncBy(n int64) { | ||
| p.bar.IncrBy(int(n)) | ||
| } | ||
|
|
||
| func (p pbProgress) GetCurrent() int64 { | ||
| return p.bar.Current() | ||
| } | ||
|
|
||
| // Close marks the progress as 100% complete and that Inc() can no longer be | ||
| // called. | ||
| func (p pbProgress) Close() { | ||
| if p.bar.Completed() || p.bar.Aborted() { | ||
| return | ||
| } | ||
| p.bar.Abort(false) | ||
| } | ||
|
|
||
| // Wait implements the ProgressWaiter interface. | ||
| func (p pbProgress) Wait(ctx context.Context) error { | ||
| ch := make(chan struct{}) | ||
| go func() { | ||
| p.progress.Wait() | ||
| close(ch) | ||
| }() | ||
| select { | ||
| case <-ctx.Done(): | ||
| return ctx.Err() | ||
| case <-ch: | ||
| return nil | ||
| } | ||
| } | ||
|
|
||
| // ProgressWaiter is the extended `Progress“: which provides a `wait` method to | ||
| // allow caller wait until all unit in the progress finished. | ||
| type ProgressWaiter interface { | ||
| Progress | ||
| Wait(context.Context) error | ||
| } | ||
|
|
||
| type noOPWaiter struct { | ||
| Progress | ||
| } | ||
|
|
||
| func (nw noOPWaiter) Wait(context.Context) error { | ||
| return nil | ||
| } | ||
|
|
||
| // cbOnComplete like `decor.OnComplete`, however allow the message provided by a function. | ||
| func cbOnComplete(decl decor.Decorator, cb func() string) decor.DecorFunc { | ||
| return func(s decor.Statistics) string { | ||
| if s.Completed { | ||
| return cb() | ||
| } | ||
| return decl.Decor(s) | ||
| } | ||
| } | ||
|
|
||
| func (ops ConsoleOperations) OutputIsTTY() bool { | ||
| f, ok := ops.Out().(*os.File) | ||
| if !ok { | ||
| return false | ||
| } | ||
| return term.IsTerminal(int(f.Fd())) | ||
| } | ||
|
|
||
| // StartProgressBar starts a progress bar with the console operations. | ||
| // Note: This function has overlapped function with `glue.StartProgress`, however this supports display extra fields | ||
| // | ||
| // after success, and implement by `mpb` (instead of `pb`). | ||
| // | ||
| // Note': Maybe replace the old `StartProgress` with `mpb` too. | ||
| func (ops ConsoleOperations) StartProgressBar(title string, total int, extraFields ...ExtraField) ProgressWaiter { | ||
| if !ops.OutputIsTTY() { | ||
| return ops.startProgressBarOverDummy(title, total, extraFields...) | ||
| } | ||
| return ops.startProgressBarOverTTY(title, total, extraFields...) | ||
| } | ||
|
|
||
| func (ops ConsoleOperations) startProgressBarOverDummy(title string, total int, extraFields ...ExtraField) ProgressWaiter { | ||
| return noOPWaiter{utils.StartProgress(context.TODO(), title, int64(total), true, nil)} | ||
| } | ||
|
|
||
| func (ops ConsoleOperations) startProgressBarOverTTY(title string, total int, extraFields ...ExtraField) ProgressWaiter { | ||
| pb := mpb.New(mpb.WithOutput(ops.Out()), mpb.WithRefreshRate(400*time.Millisecond)) | ||
| bar := adjustTotal(pb, title, total, extraFields...) | ||
|
|
||
| // If total is zero, finish right now. | ||
| if total == 0 { | ||
| bar.SetTotal(0, true) | ||
| } | ||
|
|
||
| return pbProgress{ | ||
| bar: bar, | ||
| ops: ops, | ||
| progress: pb, | ||
| } | ||
| } | ||
|
|
||
| func adjustTotal(pb *mpb.Progress, title string, total int, extraFields ...ExtraField) *mpb.Bar { | ||
| if total == OnlyOneTask { | ||
| return buildOneTaskBar(pb, title, 1) | ||
| } | ||
| return buildProgressBar(pb, title, total, extraFields...) | ||
| } | ||
|
|
||
| func buildProgressBar(pb *mpb.Progress, title string, total int, extraFields ...ExtraField) *mpb.Bar { | ||
| greenTitle := color.GreenString(title) | ||
| return pb.New(int64(total), | ||
| // Play as if the old BR style. | ||
| mpb.BarStyle().Lbound("<").Filler("-").Padding(".").Rbound(">").Tip("-", "\\", "|", "/", "-").TipOnComplete("-"), | ||
| mpb.BarFillerMiddleware(func(bf mpb.BarFiller) mpb.BarFiller { | ||
| return mpb.BarFillerFunc(func(w io.Writer, reqWidth int, stat decor.Statistics) { | ||
| if stat.Aborted || stat.Completed { | ||
| return | ||
| } | ||
| bf.Fill(w, reqWidth, stat) | ||
| }) | ||
| }), | ||
| mpb.PrependDecorators(decor.OnAbort(decor.OnComplete(decor.Name(greenTitle), fmt.Sprintf("%s ::", title)), fmt.Sprintf("%s ::", title))), | ||
| mpb.AppendDecorators(decor.OnAbort(decor.Any(cbOnComplete(decor.NewPercentage("%02.2f"), printFinalMessage(extraFields))), color.RedString("ABORTED"))), | ||
| ) | ||
| } | ||
|
|
||
| var ( | ||
| spinnerDoneText string = fmt.Sprintf("... %s", color.GreenString("DONE")) | ||
| ) | ||
|
|
||
| func buildOneTaskBar(pb *mpb.Progress, title string, total int) *mpb.Bar { | ||
| return pb.New(int64(total), | ||
| mpb.NopStyle(), | ||
| mpb.PrependDecorators(decor.Name(title)), | ||
| mpb.AppendDecorators(decor.OnAbort(decor.OnComplete(decor.Spinner(spinnerText), spinnerDoneText), color.RedString("ABORTED"))), | ||
| ) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.