context, lastState
Vladimir Bauer
9 years ago
| 0 | 0 | package mpb |
| 1 | 1 | |
| 2 | 2 | import ( |
| 3 | "context" | |
| 3 | 4 | "io" |
| 4 | 5 | "sync" |
| 5 | 6 | "time" |
| 8 | 9 | |
| 9 | 10 | // Bar represents a progress Bar |
| 10 | 11 | type Bar struct { |
| 11 | width int | |
| 12 | alpha float64 | |
| 12 | width int | |
| 13 | termWidth int | |
| 14 | alpha float64 | |
| 13 | 15 | |
| 14 | 16 | fill byte |
| 15 | 17 | empty byte |
| 17 | 19 | leftEnd byte |
| 18 | 20 | rightEnd byte |
| 19 | 21 | |
| 20 | incrCh chan int | |
| 21 | trimLeftCh chan bool | |
| 22 | trimRightCh chan bool | |
| 23 | redrawReqCh chan *redrawRequest | |
| 24 | currentReqCh chan chan int | |
| 25 | statusReqCh chan chan int | |
| 26 | decoratorCh chan *decorator | |
| 27 | flushedCh chan struct{} | |
| 28 | stopCh chan struct{} | |
| 29 | done chan struct{} | |
| 30 | } | |
| 31 | ||
| 32 | type redrawRequest struct { | |
| 33 | width int | |
| 34 | respCh chan []byte | |
| 22 | incrCh chan int64 | |
| 23 | trimLeftCh chan bool | |
| 24 | trimRightCh chan bool | |
| 25 | stateReqCh chan chan state | |
| 26 | decoratorCh chan *decorator | |
| 27 | flushedCh chan struct{} | |
| 28 | done chan struct{} | |
| 29 | ||
| 30 | lastState state | |
| 35 | 31 | } |
| 36 | 32 | |
| 37 | 33 | // Statistics represents statistics of the progress bar |
| 38 | 34 | // instance of this, sent to DecoratorFunc, as param |
| 39 | 35 | type Statistics struct { |
| 40 | Total, Current, TermWidth int | |
| 36 | Total, Current int64 | |
| 37 | TermWidth int | |
| 41 | 38 | TimeElapsed, TimePerItemEstimate time.Duration |
| 42 | 39 | } |
| 43 | 40 | |
| 45 | 42 | return time.Duration(s.Total-s.Current) * s.TimePerItemEstimate |
| 46 | 43 | } |
| 47 | 44 | |
| 48 | func newBar(total, width int, wg *sync.WaitGroup) *Bar { | |
| 45 | type state struct { | |
| 46 | total, current int64 | |
| 47 | timeElapsed, timePerItem time.Duration | |
| 48 | appendFuncs, prependFuncs []DecoratorFunc | |
| 49 | trimLeftSpace, trimRightSpace bool | |
| 50 | } | |
| 51 | ||
| 52 | func newBar(ctx context.Context, wg *sync.WaitGroup, total int64, width int) *Bar { | |
| 49 | 53 | b := &Bar{ |
| 50 | 54 | fill: '=', |
| 51 | 55 | empty: '-', |
| 55 | 59 | alpha: 0.25, |
| 56 | 60 | width: width, |
| 57 | 61 | |
| 58 | incrCh: make(chan int), | |
| 59 | trimLeftCh: make(chan bool), | |
| 60 | trimRightCh: make(chan bool), | |
| 61 | redrawReqCh: make(chan *redrawRequest), | |
| 62 | currentReqCh: make(chan chan int), | |
| 63 | statusReqCh: make(chan chan int), | |
| 64 | decoratorCh: make(chan *decorator), | |
| 65 | flushedCh: make(chan struct{}), | |
| 66 | stopCh: make(chan struct{}), | |
| 67 | done: make(chan struct{}), | |
| 68 | } | |
| 69 | go b.server(wg, total) | |
| 62 | incrCh: make(chan int64), | |
| 63 | trimLeftCh: make(chan bool), | |
| 64 | trimRightCh: make(chan bool), | |
| 65 | stateReqCh: make(chan chan state), | |
| 66 | decoratorCh: make(chan *decorator), | |
| 67 | flushedCh: make(chan struct{}), | |
| 68 | done: make(chan struct{}), | |
| 69 | } | |
| 70 | go b.server(ctx, wg, total) | |
| 70 | 71 | return b |
| 71 | 72 | } |
| 72 | 73 | |
| 145 | 146 | // Incr increments progress bar |
| 146 | 147 | func (b *Bar) Incr(n int) { |
| 147 | 148 | if !b.isDone() { |
| 148 | b.incrCh <- n | |
| 149 | b.incrCh <- int64(n) | |
| 149 | 150 | } |
| 150 | 151 | } |
| 151 | 152 | |
| 152 | 153 | // Current returns the actual current. |
| 153 | // returns 0 after bar was stopped | |
| 154 | func (b *Bar) Current() int { | |
| 154 | func (b *Bar) Current() int64 { | |
| 155 | 155 | if b.isDone() { |
| 156 | return 0 | |
| 157 | } | |
| 158 | respCh := make(chan int) | |
| 159 | b.currentReqCh <- respCh | |
| 160 | return <-respCh | |
| 161 | } | |
| 162 | ||
| 163 | // Stop stops rendering the bar | |
| 164 | func (b *Bar) Stop() { | |
| 165 | if !b.isDone() { | |
| 166 | b.stopCh <- struct{}{} | |
| 156 | return b.lastState.current | |
| 157 | } | |
| 158 | ch := make(chan state) | |
| 159 | b.stateReqCh <- ch | |
| 160 | state := <-ch | |
| 161 | return state.current | |
| 162 | } | |
| 163 | ||
| 164 | func (b *Bar) stop() { | |
| 165 | if !b.isDone() { | |
| 166 | close(b.done) | |
| 167 | 167 | } |
| 168 | 168 | } |
| 169 | 169 | |
| 193 | 193 | if width <= 0 { |
| 194 | 194 | width = b.width |
| 195 | 195 | } |
| 196 | respCh := make(chan []byte) | |
| 197 | b.redrawReqCh <- &redrawRequest{width, respCh} | |
| 198 | return <-respCh | |
| 199 | } | |
| 200 | ||
| 201 | func (b *Bar) server(wg *sync.WaitGroup, total int) { | |
| 196 | if b.isDone() { | |
| 197 | return b.draw(b.lastState, width) | |
| 198 | } | |
| 199 | ch := make(chan state) | |
| 200 | b.stateReqCh <- ch | |
| 201 | return b.draw(<-ch, width) | |
| 202 | } | |
| 203 | ||
| 204 | func (b *Bar) server(ctx context.Context, wg *sync.WaitGroup, total int64) { | |
| 205 | defer wg.Done() | |
| 206 | var completed bool | |
| 207 | state := state{total: total} | |
| 202 | 208 | timeStarted := time.Now() |
| 203 | 209 | blockStartTime := timeStarted |
| 204 | var timePerItem, timeElapsed time.Duration | |
| 205 | var appendFuncs, prependFuncs []DecoratorFunc | |
| 206 | var completed, wgDoneReported, trimLeftSpace, trimRightSpace bool | |
| 207 | var current int | |
| 208 | 210 | for { |
| 209 | 211 | select { |
| 210 | 212 | case i := <-b.incrCh: |
| 211 | n := current + i | |
| 213 | n := state.current + i | |
| 212 | 214 | if n > total { |
| 213 | current = total | |
| 215 | state.current = total | |
| 214 | 216 | completed = true |
| 217 | blockStartTime = time.Now() | |
| 215 | 218 | break // break out of select |
| 216 | 219 | } |
| 217 | timeElapsed = time.Since(timeStarted) | |
| 218 | timePerItem = calcTimePerItemEstimate(timePerItem, blockStartTime, b.alpha, i) | |
| 219 | blockStartTime = time.Now() | |
| 220 | current = n | |
| 221 | if current == total { | |
| 220 | state.timeElapsed = time.Since(timeStarted) | |
| 221 | state.timePerItem = calcTimePerItemEstimate(state.timePerItem, blockStartTime, b.alpha, i) | |
| 222 | if n == total { | |
| 222 | 223 | completed = true |
| 223 | 224 | } |
| 225 | state.current = n | |
| 226 | blockStartTime = time.Now() | |
| 224 | 227 | case d := <-b.decoratorCh: |
| 225 | 228 | switch d.kind { |
| 226 | 229 | case decoratorAppend: |
| 227 | appendFuncs = append(appendFuncs, d.f) | |
| 230 | state.appendFuncs = append(state.appendFuncs, d.f) | |
| 228 | 231 | case decoratorPrepend: |
| 229 | prependFuncs = append(prependFuncs, d.f) | |
| 232 | state.prependFuncs = append(state.prependFuncs, d.f) | |
| 230 | 233 | } |
| 231 | case respCh := <-b.currentReqCh: | |
| 232 | respCh <- current | |
| 233 | case r := <-b.redrawReqCh: | |
| 234 | stat := &Statistics{total, current, r.width, timeElapsed, timePerItem} | |
| 235 | r.respCh <- b.draw(stat, appendFuncs, prependFuncs, trimLeftSpace, trimRightSpace) | |
| 236 | case respCh := <-b.statusReqCh: | |
| 237 | respCh <- percentage(total, current, 100) | |
| 238 | case result := <-b.trimLeftCh: | |
| 239 | trimLeftSpace = result | |
| 240 | case result := <-b.trimRightCh: | |
| 241 | trimRightSpace = result | |
| 234 | case ch := <-b.stateReqCh: | |
| 235 | ch <- state | |
| 236 | case state.trimLeftSpace = <-b.trimLeftCh: | |
| 237 | case state.trimRightSpace = <-b.trimRightCh: | |
| 242 | 238 | case <-b.flushedCh: |
| 243 | if completed && !wgDoneReported { | |
| 244 | wgDoneReported = true | |
| 245 | wg.Done() | |
| 239 | if completed { | |
| 240 | b.lastState = state | |
| 241 | b.stop() | |
| 242 | return | |
| 246 | 243 | } |
| 247 | case <-b.stopCh: | |
| 248 | if !wgDoneReported { | |
| 249 | wg.Done() | |
| 250 | } | |
| 251 | close(b.done) | |
| 244 | case <-ctx.Done(): | |
| 245 | b.lastState = state | |
| 246 | b.stop() | |
| 252 | 247 | return |
| 253 | 248 | } |
| 254 | 249 | } |
| 255 | 250 | } |
| 256 | 251 | |
| 257 | func (b *Bar) draw(stat *Statistics, appendFuncs, prependFuncs []DecoratorFunc, trimLeftSpace, trimRightSpace bool) []byte { | |
| 258 | ||
| 259 | buf := make([]byte, 0, stat.TermWidth) | |
| 260 | ||
| 261 | barBlock := b.fillBar(stat.Total, stat.Current, b.width) | |
| 252 | func (b *Bar) draw(s state, termWidth int) []byte { | |
| 253 | ||
| 254 | buf := make([]byte, 0, termWidth) | |
| 255 | stat := &Statistics{ | |
| 256 | Total: s.total, | |
| 257 | Current: s.current, | |
| 258 | TermWidth: termWidth, | |
| 259 | TimeElapsed: s.timeElapsed, | |
| 260 | TimePerItemEstimate: s.timePerItem, | |
| 261 | } | |
| 262 | ||
| 263 | barBlock := b.fillBar(s.total, s.current, b.width) | |
| 262 | 264 | |
| 263 | 265 | // render append functions to the right of the bar |
| 264 | 266 | var appendBlock []byte |
| 265 | for _, f := range appendFuncs { | |
| 267 | for _, f := range s.appendFuncs { | |
| 266 | 268 | appendBlock = append(appendBlock, []byte(f(stat))...) |
| 267 | 269 | } |
| 268 | 270 | |
| 269 | 271 | // render prepend functions to the left of the bar |
| 270 | 272 | var prependBlock []byte |
| 271 | for _, f := range prependFuncs { | |
| 273 | for _, f := range s.prependFuncs { | |
| 272 | 274 | prependBlock = append(prependBlock, []byte(f(stat))...) |
| 273 | 275 | } |
| 274 | 276 | |
| 279 | 281 | var leftSpace, rightSpace []byte |
| 280 | 282 | space := []byte{' '} |
| 281 | 283 | |
| 282 | if !trimLeftSpace { | |
| 284 | if !s.trimLeftSpace { | |
| 283 | 285 | prependCount++ |
| 284 | 286 | leftSpace = space |
| 285 | 287 | } |
| 286 | if !trimRightSpace { | |
| 288 | if !s.trimRightSpace { | |
| 287 | 289 | appendCount++ |
| 288 | 290 | rightSpace = space |
| 289 | 291 | } |
| 290 | 292 | |
| 291 | 293 | totalCount := prependCount + barCount + appendCount |
| 292 | if totalCount >= stat.TermWidth { | |
| 293 | newWidth := stat.TermWidth - prependCount - appendCount | |
| 294 | barBlock = b.fillBar(stat.Total, stat.Current, newWidth-1) | |
| 294 | if totalCount >= termWidth { | |
| 295 | newWidth := termWidth - prependCount - appendCount | |
| 296 | barBlock = b.fillBar(s.total, s.current, newWidth-1) | |
| 295 | 297 | } |
| 296 | 298 | |
| 297 | 299 | for _, block := range [...][]byte{prependBlock, leftSpace, barBlock, rightSpace, appendBlock} { |
| 301 | 303 | return buf |
| 302 | 304 | } |
| 303 | 305 | |
| 304 | func (b *Bar) fillBar(total, current, width int) []byte { | |
| 306 | func (b *Bar) fillBar(total, current int64, width int) []byte { | |
| 305 | 307 | if width < 2 { |
| 306 | 308 | return []byte{b.leftEnd, b.rightEnd} |
| 307 | 309 | } |
| 335 | 337 | } |
| 336 | 338 | |
| 337 | 339 | func (b *Bar) status() int { |
| 338 | respCh := make(chan int) | |
| 339 | b.statusReqCh <- respCh | |
| 340 | return <-respCh | |
| 340 | var total, current int64 | |
| 341 | if b.isDone() { | |
| 342 | total = b.lastState.total | |
| 343 | current = b.lastState.current | |
| 344 | } else { | |
| 345 | ch := make(chan state) | |
| 346 | b.stateReqCh <- ch | |
| 347 | state := <-ch | |
| 348 | total = state.total | |
| 349 | current = state.current | |
| 350 | } | |
| 351 | return percentage(total, current, 100) | |
| 341 | 352 | } |
| 342 | 353 | |
| 343 | 354 | // SortableBarSlice satisfies sort interface |
| 349 | 360 | |
| 350 | 361 | func (p SortableBarSlice) Swap(i, j int) { p[i], p[j] = p[j], p[i] } |
| 351 | 362 | |
| 352 | func calcTimePerItemEstimate(tpie time.Duration, blockStartTime time.Time, alpha float64, items int) time.Duration { | |
| 363 | func calcTimePerItemEstimate(tpie time.Duration, blockStartTime time.Time, alpha float64, items int64) time.Duration { | |
| 353 | 364 | lastBlockTime := time.Since(blockStartTime) |
| 354 | 365 | lastItemEstimate := float64(lastBlockTime) / float64(items) |
| 355 | 366 | return time.Duration((alpha * lastItemEstimate) + (1-alpha)*float64(tpie)) |
| 356 | 367 | } |
| 357 | 368 | |
| 358 | func percentage(total, current, ratio int) int { | |
| 369 | func percentage(total, current int64, ratio int) int { | |
| 359 | 370 | if total == 0 { |
| 360 | 371 | return 0 |
| 361 | 372 | } |