refactoring suppressBar example
Vladimir Bauer
2 years ago
| 2 | 2 | import ( |
| 3 | 3 | "errors" |
| 4 | 4 | "fmt" |
| 5 | "io" | |
| 5 | "math" | |
| 6 | 6 | "math/rand" |
| 7 | 7 | "sync" |
| 8 | 8 | "time" |
| 9 | 9 | |
| 10 | "github.com/mattn/go-runewidth" | |
| 11 | 10 | "github.com/vbauerster/mpb/v8" |
| 12 | 11 | "github.com/vbauerster/mpb/v8/decor" |
| 13 | 12 | ) |
| 15 | 14 | func main() { |
| 16 | 15 | p := mpb.New() |
| 17 | 16 | |
| 18 | total := 100 | |
| 19 | msgCh := make(chan string) | |
| 20 | resumeCh := make(chan struct{}) | |
| 21 | nextCh := make(chan struct{}, 1) | |
| 22 | ew := &errorWrapper{} | |
| 17 | total, numBars := 100, 3 | |
| 18 | err := new(errorWrapper) | |
| 23 | 19 | timer := time.AfterFunc(2*time.Second, func() { |
| 24 | ew.reset(errors.New("timeout")) | |
| 20 | err.set(errors.New("timeout"), rand.Intn(numBars)) | |
| 25 | 21 | }) |
| 26 | 22 | defer timer.Stop() |
| 27 | bar := p.AddBar(int64(total), | |
| 28 | mpb.BarFillerMiddleware(func(base mpb.BarFiller) mpb.BarFiller { | |
| 29 | var msg *string | |
| 30 | var times int | |
| 31 | return mpb.BarFillerFunc(func(w io.Writer, st decor.Statistics) error { | |
| 32 | if msg == nil { | |
| 33 | select { | |
| 34 | case m := <-msgCh: | |
| 35 | msg = &m | |
| 36 | times = 10 | |
| 37 | nextCh <- struct{}{} | |
| 38 | default: | |
| 39 | } | |
| 40 | return base.Fill(w, st) | |
| 23 | ||
| 24 | for i := 0; i < numBars; i++ { | |
| 25 | msgCh := make(chan string, 1) | |
| 26 | bar := p.AddBar(int64(total), | |
| 27 | mpb.PrependDecorators(newTitleDecorator(fmt.Sprintf("Bar#%d:", i), msgCh, 16)), | |
| 28 | mpb.AppendDecorators(decor.Percentage(decor.WCSyncWidth)), | |
| 29 | ) | |
| 30 | // simulating some work | |
| 31 | barID := i | |
| 32 | go func() { | |
| 33 | max := 100 * time.Millisecond | |
| 34 | for i := 0; i < total; i++ { | |
| 35 | if err.check(barID) { | |
| 36 | msgCh <- fmt.Sprintf("%s at %d, retrying...", err.Error(), i) | |
| 37 | bar.SetRefill(int64(i)) | |
| 38 | err.reset() | |
| 39 | i-- | |
| 40 | continue | |
| 41 | 41 | } |
| 42 | switch { | |
| 43 | case times == 0, st.Completed, st.Aborted: | |
| 44 | defer func() { | |
| 45 | msg = nil | |
| 46 | }() | |
| 47 | resumeCh <- struct{}{} | |
| 48 | default: | |
| 49 | times-- | |
| 50 | } | |
| 51 | _, err := io.WriteString(w, runewidth.Truncate(*msg, st.AvailableWidth, "…")) | |
| 52 | nextCh <- struct{}{} | |
| 53 | return err | |
| 54 | }) | |
| 55 | }), | |
| 56 | mpb.PrependDecorators(decor.Name("my bar:")), | |
| 57 | mpb.AppendDecorators(newCustomPercentage(nextCh)), | |
| 58 | ) | |
| 59 | // simulating some work | |
| 60 | go func() { | |
| 61 | rng := rand.New(rand.NewSource(time.Now().UnixNano())) | |
| 62 | max := 100 * time.Millisecond | |
| 63 | for i := 0; i < total; i++ { | |
| 64 | time.Sleep(time.Duration(rng.Intn(10)+1) * max / 10) | |
| 65 | if ew.isErr() { | |
| 66 | msgCh <- fmt.Sprintf("%s at %d, retrying...", ew.Error(), i) | |
| 67 | i-- | |
| 68 | bar.SetRefill(int64(i)) | |
| 69 | ew.reset(nil) | |
| 70 | <-resumeCh | |
| 71 | continue | |
| 42 | time.Sleep(time.Duration(rand.Intn(10)+1) * max / 10) | |
| 43 | bar.Increment() | |
| 72 | 44 | } |
| 73 | bar.Increment() | |
| 74 | } | |
| 75 | }() | |
| 45 | }() | |
| 46 | } | |
| 76 | 47 | |
| 77 | 48 | p.Wait() |
| 78 | 49 | } |
| 79 | 50 | |
| 80 | 51 | type errorWrapper struct { |
| 81 | 52 | sync.RWMutex |
| 82 | err error | |
| 53 | err error | |
| 54 | barID int | |
| 83 | 55 | } |
| 84 | 56 | |
| 85 | 57 | func (ew *errorWrapper) Error() string { |
| 88 | 60 | return ew.err.Error() |
| 89 | 61 | } |
| 90 | 62 | |
| 91 | func (ew *errorWrapper) isErr() bool { | |
| 63 | func (ew *errorWrapper) check(barID int) bool { | |
| 92 | 64 | ew.RLock() |
| 93 | 65 | defer ew.RUnlock() |
| 94 | return ew.err != nil | |
| 66 | return ew.err != nil && ew.barID == barID | |
| 95 | 67 | } |
| 96 | 68 | |
| 97 | func (ew *errorWrapper) reset(err error) { | |
| 69 | func (ew *errorWrapper) set(err error, barID int) { | |
| 98 | 70 | ew.Lock() |
| 99 | 71 | ew.err = err |
| 72 | ew.barID = barID | |
| 100 | 73 | ew.Unlock() |
| 101 | 74 | } |
| 102 | 75 | |
| 103 | type percentage struct { | |
| 104 | decor.Decorator | |
| 105 | suspend <-chan struct{} | |
| 76 | func (ew *errorWrapper) reset() { | |
| 77 | ew.Lock() | |
| 78 | ew.err = nil | |
| 79 | ew.Unlock() | |
| 106 | 80 | } |
| 107 | 81 | |
| 108 | func (d percentage) Decor(s decor.Statistics) (string, int) { | |
| 109 | select { | |
| 110 | case <-d.suspend: | |
| 111 | return d.Format("") | |
| 112 | default: | |
| 113 | return d.Decorator.Decor(s) | |
| 82 | type title struct { | |
| 83 | decor.Decorator | |
| 84 | name string | |
| 85 | msgCh <-chan string | |
| 86 | msg string | |
| 87 | count int | |
| 88 | limit int | |
| 89 | } | |
| 90 | ||
| 91 | func (d *title) Decor(stat decor.Statistics) (string, int) { | |
| 92 | if d.count == 0 { | |
| 93 | select { | |
| 94 | case msg := <-d.msgCh: | |
| 95 | d.count = d.limit | |
| 96 | d.msg = msg | |
| 97 | default: | |
| 98 | return d.Decorator.Decor(stat) | |
| 99 | } | |
| 100 | } | |
| 101 | d.count-- | |
| 102 | _, _ = d.Format("") | |
| 103 | return fmt.Sprintf("%s %s", d.name, d.msg), math.MaxInt | |
| 104 | } | |
| 105 | ||
| 106 | func newTitleDecorator(name string, msgCh <-chan string, limit int) decor.Decorator { | |
| 107 | return &title{ | |
| 108 | Decorator: decor.Name(name), | |
| 109 | name: name, | |
| 110 | msgCh: msgCh, | |
| 111 | limit: limit, | |
| 114 | 112 | } |
| 115 | 113 | } |
| 116 | ||
| 117 | func newCustomPercentage(nextCh <-chan struct{}) decor.Decorator { | |
| 118 | return percentage{ | |
| 119 | Decorator: decor.Percentage(), | |
| 120 | suspend: nextCh, | |
| 121 | } | |
| 122 | } | |