refactoring: parkedBars, NewWithContext
Vladimir Bauer
7 years ago
| 5 | 5 | "fmt" |
| 6 | 6 | "io" |
| 7 | 7 | "io/ioutil" |
| 8 | "runtime" | |
| 8 | "log" | |
| 9 | 9 | "strings" |
| 10 | 10 | "sync" |
| 11 | "sync/atomic" | |
| 12 | 11 | "time" |
| 13 | 12 | "unicode/utf8" |
| 14 | 13 | |
| 35 | 34 | priority int |
| 36 | 35 | index int |
| 37 | 36 | |
| 38 | runningBar *Bar | |
| 39 | cacheState *bState | |
| 40 | operateState chan func(*bState) | |
| 41 | bFrameCh chan *bFrame | |
| 42 | syncTableCh chan [][]chan int | |
| 43 | completed chan bool | |
| 44 | forceRefreshCh chan time.Time | |
| 37 | runningBar *Bar | |
| 38 | cacheState *bState | |
| 39 | operateState chan func(*bState) | |
| 40 | bFrameCh chan *bFrame | |
| 41 | syncTableCh chan [][]chan int | |
| 42 | completed chan bool | |
| 43 | forceRefresh chan<- time.Time | |
| 45 | 44 | |
| 46 | 45 | // done is closed by Bar's goroutine, after cacheState is written |
| 47 | 46 | done chan struct{} |
| 49 | 48 | shutdown chan struct{} |
| 50 | 49 | |
| 51 | 50 | arbitraryCurrent struct { |
| 52 | lock uint32 | |
| 51 | sync.Mutex | |
| 53 | 52 | current int64 |
| 54 | 53 | } |
| 54 | ||
| 55 | dlogger *log.Logger | |
| 55 | 56 | } |
| 56 | 57 | |
| 57 | 58 | type ( |
| 90 | 91 | func newBar( |
| 91 | 92 | ctx context.Context, |
| 92 | 93 | wg *sync.WaitGroup, |
| 93 | forceRefreshCh chan time.Time, | |
| 94 | filler Filler, | |
| 95 | id, width int, | |
| 96 | total int64, | |
| 97 | options ...BarOption, | |
| 94 | forceRefresh chan<- time.Time, | |
| 95 | bs *bState, | |
| 96 | dlogger *log.Logger, | |
| 98 | 97 | ) *Bar { |
| 99 | if total <= 0 { | |
| 100 | total = time.Now().Unix() | |
| 101 | } | |
| 102 | ||
| 103 | s := &bState{ | |
| 104 | filler: filler, | |
| 105 | id: id, | |
| 106 | priority: id, | |
| 107 | width: width, | |
| 108 | total: total, | |
| 109 | } | |
| 110 | ||
| 111 | for _, opt := range options { | |
| 112 | if opt != nil { | |
| 113 | opt(s) | |
| 114 | } | |
| 115 | } | |
| 116 | ||
| 117 | s.bufP = bytes.NewBuffer(make([]byte, 0, width)) | |
| 118 | s.bufB = bytes.NewBuffer(make([]byte, 0, width)) | |
| 119 | s.bufA = bytes.NewBuffer(make([]byte, 0, width)) | |
| 120 | if s.extender != nil { | |
| 121 | s.bufE = bytes.NewBuffer(make([]byte, 0, width)) | |
| 122 | } | |
| 123 | ||
| 124 | b := &Bar{ | |
| 125 | priority: s.priority, | |
| 126 | runningBar: s.runningBar, | |
| 127 | operateState: make(chan func(*bState)), | |
| 128 | bFrameCh: make(chan *bFrame, 1), | |
| 129 | syncTableCh: make(chan [][]chan int), | |
| 130 | completed: make(chan bool), | |
| 131 | done: make(chan struct{}), | |
| 132 | shutdown: make(chan struct{}), | |
| 133 | forceRefreshCh: forceRefreshCh, | |
| 134 | } | |
| 135 | ||
| 136 | if b.runningBar != nil { | |
| 137 | b.priority = b.runningBar.priority | |
| 138 | } | |
| 139 | ||
| 140 | go b.serve(ctx, wg, s) | |
| 141 | return b | |
| 98 | ||
| 99 | bs.bufP = bytes.NewBuffer(make([]byte, 0, bs.width)) | |
| 100 | bs.bufB = bytes.NewBuffer(make([]byte, 0, bs.width)) | |
| 101 | bs.bufA = bytes.NewBuffer(make([]byte, 0, bs.width)) | |
| 102 | if bs.extender != nil { | |
| 103 | bs.bufE = bytes.NewBuffer(make([]byte, 0, bs.width)) | |
| 104 | } | |
| 105 | ||
| 106 | bar := &Bar{ | |
| 107 | priority: bs.priority, | |
| 108 | runningBar: bs.runningBar, | |
| 109 | operateState: make(chan func(*bState)), | |
| 110 | bFrameCh: make(chan *bFrame, 1), | |
| 111 | syncTableCh: make(chan [][]chan int), | |
| 112 | completed: make(chan bool), | |
| 113 | done: make(chan struct{}), | |
| 114 | shutdown: make(chan struct{}), | |
| 115 | forceRefresh: forceRefresh, | |
| 116 | dlogger: dlogger, | |
| 117 | } | |
| 118 | ||
| 119 | go bar.serve(ctx, wg, bs) | |
| 120 | return bar | |
| 142 | 121 | } |
| 143 | 122 | |
| 144 | 123 | // RemoveAllPrependers removes all prepend functions. |
| 211 | 190 | if final && !s.toComplete { |
| 212 | 191 | s.current = s.total |
| 213 | 192 | s.toComplete = true |
| 214 | go b.forceRefresh() | |
| 193 | go b.refreshNowTillShutdown() | |
| 215 | 194 | } |
| 216 | 195 | }: |
| 217 | 196 | return true |
| 225 | 204 | if current <= 0 { |
| 226 | 205 | return |
| 227 | 206 | } |
| 228 | for !atomic.CompareAndSwapUint32(&b.arbitraryCurrent.lock, 0, 1) { | |
| 229 | runtime.Gosched() | |
| 230 | } | |
| 207 | b.arbitraryCurrent.Lock() | |
| 231 | 208 | last := b.arbitraryCurrent.current |
| 232 | 209 | b.IncrBy(int(current-last), wdd...) |
| 233 | 210 | b.arbitraryCurrent.current = current |
| 234 | atomic.StoreUint32(&b.arbitraryCurrent.lock, 0) | |
| 211 | b.arbitraryCurrent.Unlock() | |
| 235 | 212 | } |
| 236 | 213 | |
| 237 | 214 | // Increment is a shorthand for b.IncrBy(1). |
| 249 | 226 | if s.current >= s.total && !s.toComplete { |
| 250 | 227 | s.current = s.total |
| 251 | 228 | s.toComplete = true |
| 252 | go b.forceRefresh() | |
| 229 | go b.refreshNowTillShutdown() | |
| 253 | 230 | } |
| 254 | 231 | for _, ar := range s.amountReceivers { |
| 255 | 232 | ar.NextAmount(n, wdd...) |
| 307 | 284 | // recovering if user defined decorator panics for example |
| 308 | 285 | if p := recover(); p != nil { |
| 309 | 286 | s.panicMsg = fmt.Sprintf("panic: %v", p) |
| 310 | fmt.Fprintf(debugOut, "%s %s bar id %02d %v\n", "[mpb]", time.Now(), s.id, s.panicMsg) | |
| 287 | b.dlogger.Println(s.panicMsg) | |
| 311 | 288 | b.bFrameCh <- &bFrame{ |
| 312 | 289 | rd: strings.NewReader(fmt.Sprintf(fmt.Sprintf("%%.%ds\n", tw), s.panicMsg)), |
| 313 | 290 | toShutdown: true, |
| 410 | 387 | return table |
| 411 | 388 | } |
| 412 | 389 | |
| 413 | func (b *Bar) forceRefresh() { | |
| 390 | func (b *Bar) refreshNowTillShutdown() { | |
| 414 | 391 | for { |
| 415 | 392 | select { |
| 416 | case b.forceRefreshCh <- time.Now(): | |
| 393 | case b.forceRefresh <- time.Now(): | |
| 417 | 394 | case <-b.shutdown: |
| 418 | 395 | return |
| 419 | 396 | } |
| 47 | 47 | } |
| 48 | 48 | } |
| 49 | 49 | |
| 50 | // WithContext provided context will be used for cancellation purposes. | |
| 50 | // WithContext deprecated and has no effect, please use NewWithContext instead. | |
| 51 | 51 | func WithContext(ctx context.Context) ContainerOption { |
| 52 | 52 | return func(s *pState) { |
| 53 | if ctx == nil { | |
| 54 | return | |
| 55 | } | |
| 56 | s.ctx = ctx | |
| 57 | 53 | } |
| 58 | 54 | } |
| 59 | 55 |
| 5 | 5 | "fmt" |
| 6 | 6 | "io" |
| 7 | 7 | "io/ioutil" |
| 8 | "log" | |
| 8 | 9 | "os" |
| 9 | 10 | "sync" |
| 10 | 11 | "time" |
| 21 | 22 | |
| 22 | 23 | // Progress represents the container that renders Progress bars |
| 23 | 24 | type Progress struct { |
| 25 | ctx context.Context | |
| 24 | 26 | uwg *sync.WaitGroup |
| 25 | 27 | cwg *sync.WaitGroup |
| 26 | 28 | bwg *sync.WaitGroup |
| 27 | 29 | operateState chan func(*pState) |
| 28 | 30 | done chan struct{} |
| 31 | forceRefresh chan time.Time | |
| 32 | dlogger *log.Logger | |
| 29 | 33 | } |
| 30 | 34 | |
| 31 | 35 | type pState struct { |
| 37 | 41 | rr time.Duration |
| 38 | 42 | pMatrix map[int][]chan int |
| 39 | 43 | aMatrix map[int][]chan int |
| 40 | forceRefreshCh chan time.Time | |
| 41 | 44 | output io.Writer |
| 42 | 45 | |
| 43 | 46 | // following are provided/overrided by user |
| 44 | ctx context.Context | |
| 45 | 47 | uwg *sync.WaitGroup |
| 46 | 48 | manualRefreshCh <-chan time.Time |
| 47 | 49 | shutdownNotifier chan struct{} |
| 48 | waitBars map[*Bar]*Bar | |
| 50 | parkedBars map[*Bar]*Bar | |
| 49 | 51 | debugOut io.Writer |
| 50 | 52 | } |
| 51 | 53 | |
| 52 | // New creates new Progress instance, which orchestrates bars rendering | |
| 53 | // process. Accepts mpb.ContainerOption funcs for customization. | |
| 54 | // New creates new Progress container instance. It's not possible to | |
| 55 | // reuse instance after *Progress.Wait() method has been called. | |
| 54 | 56 | func New(options ...ContainerOption) *Progress { |
| 57 | return NewWithContext(context.Background(), options...) | |
| 58 | } | |
| 59 | ||
| 60 | // NewWithContext creates new Progress container instance with provided | |
| 61 | // context. It's not possible to reuse instance after *Progress.Wait() | |
| 62 | // method has been called. | |
| 63 | func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress { | |
| 55 | 64 | pq := make(priorityQueue, 0) |
| 56 | 65 | heap.Init(&pq) |
| 57 | 66 | |
| 58 | 67 | s := &pState{ |
| 59 | ctx: context.Background(), | |
| 60 | bHeap: &pq, | |
| 61 | width: pwidth, | |
| 62 | rr: prr, | |
| 63 | waitBars: make(map[*Bar]*Bar), | |
| 64 | forceRefreshCh: make(chan time.Time), | |
| 65 | debugOut: ioutil.Discard, | |
| 66 | output: os.Stdout, | |
| 68 | bHeap: &pq, | |
| 69 | width: pwidth, | |
| 70 | rr: prr, | |
| 71 | parkedBars: make(map[*Bar]*Bar), | |
| 72 | output: os.Stdout, | |
| 73 | debugOut: ioutil.Discard, | |
| 67 | 74 | } |
| 68 | 75 | |
| 69 | 76 | for _, opt := range options { |
| 73 | 80 | } |
| 74 | 81 | |
| 75 | 82 | p := &Progress{ |
| 83 | ctx: ctx, | |
| 76 | 84 | uwg: s.uwg, |
| 77 | 85 | cwg: new(sync.WaitGroup), |
| 78 | 86 | bwg: new(sync.WaitGroup), |
| 79 | 87 | operateState: make(chan func(*pState)), |
| 88 | forceRefresh: make(chan time.Time), | |
| 80 | 89 | done: make(chan struct{}), |
| 90 | dlogger: log.New(s.debugOut, "[mpb] ", log.Lshortfile), | |
| 81 | 91 | } |
| 82 | 92 | p.cwg.Add(1) |
| 83 | 93 | go p.serve(s, cwriter.New(s.output)) |
| 100 | 110 | |
| 101 | 111 | // Add creates a bar which renders itself by provided filler. |
| 102 | 112 | func (p *Progress) Add(total int64, filler Filler, options ...BarOption) *Bar { |
| 113 | if total <= 0 { | |
| 114 | total = time.Now().Unix() | |
| 115 | } | |
| 116 | if filler == nil { | |
| 117 | filler = newDefaultBarFiller() | |
| 118 | } | |
| 103 | 119 | p.bwg.Add(1) |
| 104 | 120 | result := make(chan *Bar) |
| 105 | 121 | select { |
| 106 | case p.operateState <- func(s *pState) { | |
| 107 | b := newBar(s.ctx, p.bwg, s.forceRefreshCh, filler, s.idCounter, s.width, total, options...) | |
| 108 | if b.runningBar != nil { | |
| 109 | s.waitBars[b.runningBar] = b | |
| 122 | case p.operateState <- func(ps *pState) { | |
| 123 | logPrefix := fmt.Sprintf("%sbar#%02d ", p.dlogger.Prefix(), ps.idCounter) | |
| 124 | dlogger := log.New(ps.debugOut, logPrefix, log.Lshortfile) | |
| 125 | bs := &bState{ | |
| 126 | total: total, | |
| 127 | filler: filler, | |
| 128 | priority: ps.idCounter, | |
| 129 | id: ps.idCounter, | |
| 130 | width: ps.width, | |
| 131 | } | |
| 132 | for _, opt := range options { | |
| 133 | if opt != nil { | |
| 134 | opt(bs) | |
| 135 | } | |
| 136 | } | |
| 137 | bar := newBar(p.ctx, p.bwg, p.forceRefresh, bs, dlogger) | |
| 138 | if bar.runningBar != nil { | |
| 139 | if bar.priority == ps.idCounter { | |
| 140 | bar.priority = bar.runningBar.priority | |
| 141 | } | |
| 142 | ps.parkedBars[bar.runningBar] = bar | |
| 110 | 143 | } else { |
| 111 | heap.Push(s.bHeap, b) | |
| 112 | s.heapUpdated = true | |
| 113 | } | |
| 114 | s.idCounter++ | |
| 115 | result <- b | |
| 144 | heap.Push(ps.bHeap, bar) | |
| 145 | ps.heapUpdated = true | |
| 146 | } | |
| 147 | ps.idCounter++ | |
| 148 | result <- bar | |
| 116 | 149 | }: |
| 117 | 150 | return <-result |
| 118 | 151 | case <-p.done: |
| 184 | 217 | manualOrTickCh, cleanUp := s.manualOrTick() |
| 185 | 218 | defer cleanUp() |
| 186 | 219 | |
| 187 | refreshCh := fanInRefreshSrc(p.done, s.forceRefreshCh, manualOrTickCh) | |
| 220 | refreshCh := fanInRefreshSrc(p.done, p.forceRefresh, manualOrTickCh) | |
| 188 | 221 | |
| 189 | 222 | for { |
| 190 | 223 | select { |
| 198 | 231 | return |
| 199 | 232 | } |
| 200 | 233 | if err := s.render(cw); err != nil { |
| 201 | fmt.Fprintf(s.debugOut, "[mpb] %s %v\n", time.Now(), err) | |
| 234 | p.dlogger.Println(err) | |
| 202 | 235 | } |
| 203 | 236 | } |
| 204 | 237 | } |
| 235 | 268 | // only after the bar with completed state has been flushed. this |
| 236 | 269 | // ensures no bar ends up with less than 100% rendered. |
| 237 | 270 | s.shutdownPending = append(s.shutdownPending, bar) |
| 238 | if replacementBar, ok := s.waitBars[bar]; ok { | |
| 271 | if replacementBar, ok := s.parkedBars[bar]; ok { | |
| 239 | 272 | heap.Push(s.bHeap, replacementBar) |
| 240 | 273 | s.heapUpdated = true |
| 241 | delete(s.waitBars, bar) | |
| 274 | delete(s.parkedBars, bar) | |
| 242 | 275 | } |
| 243 | 276 | if frame.removeOnComplete { |
| 244 | 277 | s.heapUpdated = true |
| 75 | 75 | func TestWithContext(t *testing.T) { |
| 76 | 76 | ctx, cancel := context.WithCancel(context.Background()) |
| 77 | 77 | shutdown := make(chan struct{}) |
| 78 | p := mpb.New( | |
| 78 | p := mpb.NewWithContext(ctx, | |
| 79 | 79 | mpb.WithOutput(ioutil.Discard), |
| 80 | mpb.WithContext(ctx), | |
| 81 | 80 | mpb.WithRefreshRate(50*time.Millisecond), |
| 82 | 81 | mpb.WithShutdownNotifier(shutdown), |
| 83 | 82 | ) |