Codebase list golang-github-vbauerster-mpb / 7316c86
Reduce allocs by utilizing bytes.Buffer Vladimir Bauer 8 years ago
1 changed file(s) with 100 addition(s) and 117 deletion(s). Raw diff Collapse all Expand all
+100
-117
bar.go less more
00 package mpb
11
22 import (
3 "bytes"
34 "fmt"
45 "io"
6 "os"
57 "sync"
68 "time"
79 "unicode/utf8"
4345 till int64
4446 }
4547 state struct {
46 id int
47 width int
48 format fmtRunes
49 etaAlpha float64
50 total int64
51 current int64
52 trimLeftSpace bool
53 trimRightSpace bool
54 completed bool
55 aborted bool
56 startTime time.Time
57 timeElapsed time.Duration
58 blockStartTime time.Time
59 timePerItem time.Duration
60 appendFuncs []decor.DecoratorFunc
61 prependFuncs []decor.DecoratorFunc
62 simpleSpinner func() byte
63 refill *refill
48 id int
49 width int
50 format fmtRunes
51 etaAlpha float64
52 total int64
53 current int64
54 trimLeftSpace bool
55 trimRightSpace bool
56 completed bool
57 aborted bool
58 startTime time.Time
59 timeElapsed time.Duration
60 blockStartTime time.Time
61 timePerItem time.Duration
62 appendFuncs []decor.DecoratorFunc
63 prependFuncs []decor.DecoratorFunc
64 simpleSpinner func() byte
65 refill *refill
66 bufP, bufB, bufA *bytes.Buffer
6467 }
6568 )
6669
7073 etaAlpha: etaAlpha,
7174 }
7275
73 if total <= 0 {
74 s.simpleSpinner = getSpinner()
75 }
76 // if total <= 0 {
77 // s.simpleSpinner = getSpinner()
78 // }
7679
7780 for _, opt := range options {
7881 opt(&s)
7982 }
83
84 s.bufP = bytes.NewBuffer(make([]byte, 0, s.width/2))
85 s.bufB = bytes.NewBuffer(make([]byte, 0, s.width))
86 s.bufA = bytes.NewBuffer(make([]byte, 0, s.width/2))
8087
8188 b := &Bar{
8289 quit: make(chan struct{}),
244251 }
245252 }:
246253 case <-time.After(prr):
247 return
248254 }
249255 }
250256
251257 func (b *Bar) server(s state, wg *sync.WaitGroup, cancel <-chan struct{}) {
252
253258 defer func() {
254259 b.cacheState = s
255260 close(b.done)
272277 }
273278
274279 func (b *Bar) render(tw int, flushed chan struct{}, prependWs, appendWs *widthSync) <-chan []byte {
275 ch := make(chan []byte, 1)
280 ch := make(chan []byte)
276281
277282 go func() {
278283 defer func() {
279284 // recovering if external decorators panic
280285 if p := recover(); p != nil {
281 ch <- []byte(fmt.Sprintln(p))
286 fmt.Fprintf(os.Stderr, "bar panic: %q\n", p)
282287 }
283 close(ch)
284288 }()
285289 var st state
286290 result := make(chan state, 1)
287291 select {
288292 case b.ops <- func(s *state) {
289 result <- *s
290293 if s.completed {
291 <-flushed
294 fmt.Fprintln(os.Stderr, "bar completed")
295 // <-flushed
292296 b.Complete()
293297 }
298 result <- *s
294299 }:
295300 st = <-result
296301 case <-b.done:
297302 st = b.cacheState
298303 }
299 buf := draw(&st, tw, prependWs, appendWs)
304 st.draw(tw, prependWs, appendWs)
305 buf := make([]byte, 0, st.bufP.Len()+st.bufB.Len()+st.bufA.Len())
306 buf = concatenateBlocks(buf, st.bufP.Bytes(), st.bufB.Bytes(), st.bufA.Bytes())
300307 buf = append(buf, '\n')
301308 ch <- buf
309 close(ch)
302310 }()
303311
304312 return ch
317325 s.timePerItem = time.Duration((s.etaAlpha * lastItemEstimate) + (1-s.etaAlpha)*float64(s.timePerItem))
318326 }
319327
320 func draw(s *state, termWidth int, prependWs, appendWs *widthSync) []byte {
321 if len(s.prependFuncs) != len(prependWs.Listen) || len(s.appendFuncs) != len(appendWs.Listen) {
322 return []byte{}
323 }
328 func (s *state) draw(termWidth int, prependWs, appendWs *widthSync) {
324329 if termWidth <= 0 {
325330 termWidth = s.width
326331 }
328333 stat := newStatistics(s)
329334
330335 // render prepend functions to the left of the bar
331 var prependBlock []byte
336 s.bufP.Reset()
332337 for i, f := range s.prependFuncs {
333 prependBlock = append(prependBlock,
334 []byte(f(stat, prependWs.Listen[i], prependWs.Result[i]))...)
338 s.bufP.WriteString(f(stat, prependWs.Listen[i], prependWs.Result[i]))
339 }
340
341 if !s.trimLeftSpace {
342 s.bufP.WriteByte(' ')
335343 }
336344
337345 // render append functions to the right of the bar
338 var appendBlock []byte
346 s.bufA.Reset()
339347 for i, f := range s.appendFuncs {
340 appendBlock = append(appendBlock,
341 []byte(f(stat, appendWs.Listen[i], appendWs.Result[i]))...)
342 }
343
344 prependCount := utf8.RuneCount(prependBlock)
345 appendCount := utf8.RuneCount(appendBlock)
346
347 var leftSpace, rightSpace []byte
348 space := []byte{' '}
349
350 if !s.trimLeftSpace {
351 prependCount++
352 leftSpace = space
353 }
348 s.bufA.WriteString(f(stat, appendWs.Listen[i], appendWs.Result[i]))
349 }
350
354351 if !s.trimRightSpace {
355 appendCount++
356 rightSpace = space
357 }
358
359 var barBlock []byte
360 buf := make([]byte, 0, termWidth)
361 segments := fmtRunesToByteSegments(s.format)
362
363 if s.simpleSpinner != nil {
364 for _, block := range [...][]byte{segments[rLeft], {s.simpleSpinner()}, segments[rRight]} {
365 barBlock = append(barBlock, block...)
352 s.bufA.WriteByte(' ')
353 }
354
355 prependCount := utf8.RuneCount(s.bufP.Bytes())
356 appendCount := utf8.RuneCount(s.bufA.Bytes())
357
358 s.fillBar(s.width)
359 barCount := utf8.RuneCount(s.bufB.Bytes())
360 totalCount := prependCount + barCount + appendCount
361 if totalCount > termWidth {
362 shrinkWidth := termWidth - prependCount - appendCount
363 s.fillBar(shrinkWidth)
364 }
365 }
366
367 func (s *state) fillBar(width int) {
368 s.bufB.Reset()
369 if width <= 2 {
370 return
371 }
372
373 // bar s.width without leftEnd and rightEnd runes
374 barWidth := width - 2
375
376 completedWidth := decor.CalcPercentage(s.total, s.current, barWidth)
377
378 s.bufB.WriteRune(s.format[rLeft])
379
380 if s.refill != nil {
381 till := decor.CalcPercentage(s.total, s.refill.till, barWidth)
382 // append refill rune
383 for i := 0; i < till; i++ {
384 s.bufB.WriteRune(s.refill.char)
385 }
386 for i := till; i < completedWidth; i++ {
387 s.bufB.WriteRune(s.format[rFill])
366388 }
367389 } else {
368 barBlock = fillBar(s.total, s.current, s.width, segments, s.refill)
369 barCount := utf8.RuneCount(barBlock)
370 totalCount := prependCount + barCount + appendCount
371 if totalCount > termWidth {
372 shrinkWidth := termWidth - prependCount - appendCount
373 barBlock = fillBar(s.total, s.current, shrinkWidth, segments, s.refill)
374 }
375 }
376
377 return concatenateBlocks(buf, prependBlock, leftSpace, barBlock, rightSpace, appendBlock)
390 for i := 0; i < completedWidth; i++ {
391 s.bufB.WriteRune(s.format[rFill])
392 }
393 }
394
395 if completedWidth < barWidth && completedWidth > 0 {
396 _, size := utf8.DecodeLastRune(s.bufB.Bytes())
397 s.bufB.Truncate(s.bufB.Len() - size)
398 s.bufB.WriteRune(s.format[rTip])
399 }
400
401 for i := completedWidth; i < barWidth; i++ {
402 s.bufB.WriteRune(s.format[rEmpty])
403 }
404
405 s.bufB.WriteRune(s.format[rRight])
378406 }
379407
380408 func concatenateBlocks(buf []byte, blocks ...[]byte) []byte {
381409 for _, block := range blocks {
382410 buf = append(buf, block...)
383411 }
384 return buf
385 }
386
387 func fillBar(total, current int64, width int, fmtBytes fmtByteSegments, rf *refill) []byte {
388 if width < 2 || total <= 0 {
389 return []byte{}
390 }
391
392 // bar width without leftEnd and rightEnd runes
393 barWidth := width - 2
394
395 completedWidth := decor.CalcPercentage(total, current, barWidth)
396
397 buf := make([]byte, 0, width)
398 buf = append(buf, fmtBytes[rLeft]...)
399
400 if rf != nil {
401 till := decor.CalcPercentage(total, rf.till, barWidth)
402 rbytes := make([]byte, utf8.RuneLen(rf.char))
403 utf8.EncodeRune(rbytes, rf.char)
404 // append refill rune
405 for i := 0; i < till; i++ {
406 buf = append(buf, rbytes...)
407 }
408 for i := till; i < completedWidth; i++ {
409 buf = append(buf, fmtBytes[rFill]...)
410 }
411 } else {
412 for i := 0; i < completedWidth; i++ {
413 buf = append(buf, fmtBytes[rFill]...)
414 }
415 }
416
417 if completedWidth < barWidth && completedWidth > 0 {
418 _, size := utf8.DecodeLastRune(buf)
419 buf = buf[:len(buf)-size]
420 buf = append(buf, fmtBytes[rTip]...)
421 }
422
423 for i := completedWidth; i < barWidth; i++ {
424 buf = append(buf, fmtBytes[rEmpty]...)
425 }
426
427 buf = append(buf, fmtBytes[rRight]...)
428
429412 return buf
430413 }
431414