refactoring p.done = ctx.Done()
Vladimir Bauer
3 years ago
| 22 | 22 | |
| 23 | 23 | // Progress represents a container that renders one or more progress bars. |
| 24 | 24 | type Progress struct { |
| 25 | uwg *sync.WaitGroup | |
| 26 | bwg *sync.WaitGroup | |
| 27 | operateState chan func(*pState) | |
| 28 | interceptIO chan func(io.Writer) | |
| 29 | done, shutdown chan struct{} | |
| 30 | cancel func() | |
| 25 | uwg *sync.WaitGroup | |
| 26 | pwg, bwg sync.WaitGroup | |
| 27 | operateState chan func(*pState) | |
| 28 | interceptIO chan func(io.Writer) | |
| 29 | done <-chan struct{} | |
| 30 | cancel func() | |
| 31 | 31 | } |
| 32 | 32 | |
| 33 | 33 | // pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine. |
| 85 | 85 | } |
| 86 | 86 | } |
| 87 | 87 | |
| 88 | go s.hm.run() | |
| 89 | ||
| 90 | cw := cwriter.New(s.output) | |
| 91 | if (cw.IsTerminal() || s.forceAutoRefresh) && !s.manualRefresh { | |
| 92 | go s.autoRefresh() | |
| 93 | } | |
| 94 | ||
| 88 | 95 | p := &Progress{ |
| 89 | 96 | uwg: s.uwg, |
| 90 | bwg: new(sync.WaitGroup), | |
| 91 | 97 | operateState: make(chan func(*pState)), |
| 92 | 98 | interceptIO: make(chan func(io.Writer)), |
| 93 | done: make(chan struct{}), | |
| 94 | shutdown: make(chan struct{}), | |
| 99 | done: ctx.Done(), | |
| 95 | 100 | cancel: cancel, |
| 96 | 101 | } |
| 97 | cw := cwriter.New(s.output) | |
| 98 | go p.serve(s, cw, s.newTicker(cw.IsTerminal(), p.done)) | |
| 99 | go s.hm.run() | |
| 102 | ||
| 103 | p.pwg.Add(1) | |
| 104 | go p.serve(s, cw) | |
| 100 | 105 | return p |
| 101 | 106 | } |
| 102 | 107 | |
| 226 | 231 | // are doing. Proper way to shutdown is to call (*Progress).Wait() instead. |
| 227 | 232 | func (p *Progress) Shutdown() { |
| 228 | 233 | p.cancel() |
| 229 | <-p.shutdown | |
| 230 | } | |
| 231 | ||
| 232 | func (p *Progress) serve(s *pState, cw *cwriter.Writer, tickerC <-chan time.Time) { | |
| 234 | p.pwg.Wait() | |
| 235 | } | |
| 236 | ||
| 237 | func (p *Progress) serve(s *pState, cw *cwriter.Writer) { | |
| 238 | defer p.pwg.Done() | |
| 239 | render := func() error { return s.render(cw) } | |
| 233 | 240 | var err error |
| 234 | render := func() error { return s.render(cw) } | |
| 235 | 241 | |
| 236 | 242 | for { |
| 237 | 243 | select { |
| 239 | 245 | op(s) |
| 240 | 246 | case fn := <-p.interceptIO: |
| 241 | 247 | fn(cw) |
| 242 | case <-tickerC: | |
| 248 | case <-s.refreshCh: | |
| 243 | 249 | e := render() |
| 244 | 250 | if e != nil { |
| 245 | 251 | p.cancel() // cancel all bars |
| 260 | 266 | _, _ = fmt.Fprintln(s.debugOut, err.Error()) |
| 261 | 267 | } |
| 262 | 268 | s.hm.end(s.shutdownNotifier) |
| 263 | close(p.shutdown) | |
| 264 | 269 | return |
| 265 | 270 | } |
| 266 | 271 | } |
| 267 | 272 | } |
| 268 | 273 | |
| 269 | func (s *pState) newTicker(isTerminal bool, done chan struct{}) chan time.Time { | |
| 270 | ch := make(chan time.Time, 1) | |
| 271 | go func() { | |
| 272 | var autoRefresh <-chan time.Time | |
| 273 | if (isTerminal || s.forceAutoRefresh) && !s.manualRefresh { | |
| 274 | if s.renderDelay != nil { | |
| 275 | <-s.renderDelay | |
| 276 | } | |
| 277 | ticker := time.NewTicker(s.refreshRate) | |
| 278 | defer ticker.Stop() | |
| 279 | autoRefresh = ticker.C | |
| 280 | } | |
| 281 | for { | |
| 282 | select { | |
| 283 | case t := <-autoRefresh: | |
| 284 | ch <- t | |
| 285 | case t := <-s.refreshCh: | |
| 286 | ch <- t | |
| 287 | case <-s.ctx.Done(): | |
| 288 | close(done) | |
| 289 | return | |
| 290 | } | |
| 291 | } | |
| 292 | }() | |
| 293 | return ch | |
| 274 | func (s *pState) autoRefresh() { | |
| 275 | if s.renderDelay != nil { | |
| 276 | <-s.renderDelay | |
| 277 | } | |
| 278 | ticker := time.NewTicker(s.refreshRate) | |
| 279 | defer ticker.Stop() | |
| 280 | for { | |
| 281 | select { | |
| 282 | case t := <-ticker.C: | |
| 283 | s.refreshCh <- t | |
| 284 | case <-s.ctx.Done(): | |
| 285 | return | |
| 286 | } | |
| 287 | } | |
| 294 | 288 | } |
| 295 | 289 | |
| 296 | 290 | func (s *pState) render(cw *cwriter.Writer) (err error) { |