Reduce allocs by utilizing bytes.Buffer
Vladimir Bauer
8 years ago
| 0 | 0 | package mpb |
| 1 | 1 | |
| 2 | 2 | import ( |
| 3 | "bytes" | |
| 3 | 4 | "fmt" |
| 4 | 5 | "io" |
| 6 | "os" | |
| 5 | 7 | "sync" |
| 6 | 8 | "time" |
| 7 | 9 | "unicode/utf8" |
| 43 | 45 | till int64 |
| 44 | 46 | } |
| 45 | 47 | 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 | |
| 64 | 67 | } |
| 65 | 68 | ) |
| 66 | 69 | |
| 70 | 73 | etaAlpha: etaAlpha, |
| 71 | 74 | } |
| 72 | 75 | |
| 73 | if total <= 0 { | |
| 74 | s.simpleSpinner = getSpinner() | |
| 75 | } | |
| 76 | // if total <= 0 { | |
| 77 | // s.simpleSpinner = getSpinner() | |
| 78 | // } | |
| 76 | 79 | |
| 77 | 80 | for _, opt := range options { |
| 78 | 81 | opt(&s) |
| 79 | 82 | } |
| 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)) | |
| 80 | 87 | |
| 81 | 88 | b := &Bar{ |
| 82 | 89 | quit: make(chan struct{}), |
| 244 | 251 | } |
| 245 | 252 | }: |
| 246 | 253 | case <-time.After(prr): |
| 247 | return | |
| 248 | 254 | } |
| 249 | 255 | } |
| 250 | 256 | |
| 251 | 257 | func (b *Bar) server(s state, wg *sync.WaitGroup, cancel <-chan struct{}) { |
| 252 | ||
| 253 | 258 | defer func() { |
| 254 | 259 | b.cacheState = s |
| 255 | 260 | close(b.done) |
| 272 | 277 | } |
| 273 | 278 | |
| 274 | 279 | 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) | |
| 276 | 281 | |
| 277 | 282 | go func() { |
| 278 | 283 | defer func() { |
| 279 | 284 | // recovering if external decorators panic |
| 280 | 285 | if p := recover(); p != nil { |
| 281 | ch <- []byte(fmt.Sprintln(p)) | |
| 286 | fmt.Fprintf(os.Stderr, "bar panic: %q\n", p) | |
| 282 | 287 | } |
| 283 | close(ch) | |
| 284 | 288 | }() |
| 285 | 289 | var st state |
| 286 | 290 | result := make(chan state, 1) |
| 287 | 291 | select { |
| 288 | 292 | case b.ops <- func(s *state) { |
| 289 | result <- *s | |
| 290 | 293 | if s.completed { |
| 291 | <-flushed | |
| 294 | fmt.Fprintln(os.Stderr, "bar completed") | |
| 295 | // <-flushed | |
| 292 | 296 | b.Complete() |
| 293 | 297 | } |
| 298 | result <- *s | |
| 294 | 299 | }: |
| 295 | 300 | st = <-result |
| 296 | 301 | case <-b.done: |
| 297 | 302 | st = b.cacheState |
| 298 | 303 | } |
| 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()) | |
| 300 | 307 | buf = append(buf, '\n') |
| 301 | 308 | ch <- buf |
| 309 | close(ch) | |
| 302 | 310 | }() |
| 303 | 311 | |
| 304 | 312 | return ch |
| 317 | 325 | s.timePerItem = time.Duration((s.etaAlpha * lastItemEstimate) + (1-s.etaAlpha)*float64(s.timePerItem)) |
| 318 | 326 | } |
| 319 | 327 | |
| 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) { | |
| 324 | 329 | if termWidth <= 0 { |
| 325 | 330 | termWidth = s.width |
| 326 | 331 | } |
| 328 | 333 | stat := newStatistics(s) |
| 329 | 334 | |
| 330 | 335 | // render prepend functions to the left of the bar |
| 331 | var prependBlock []byte | |
| 336 | s.bufP.Reset() | |
| 332 | 337 | 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(' ') | |
| 335 | 343 | } |
| 336 | 344 | |
| 337 | 345 | // render append functions to the right of the bar |
| 338 | var appendBlock []byte | |
| 346 | s.bufA.Reset() | |
| 339 | 347 | 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 | ||
| 354 | 351 | 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]) | |
| 366 | 388 | } |
| 367 | 389 | } 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]) | |
| 378 | 406 | } |
| 379 | 407 | |
| 380 | 408 | func concatenateBlocks(buf []byte, blocks ...[]byte) []byte { |
| 381 | 409 | for _, block := range blocks { |
| 382 | 410 | buf = append(buf, block...) |
| 383 | 411 | } |
| 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 | ||
| 429 | 412 | return buf |
| 430 | 413 | } |
| 431 | 414 | |