pool.go 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. // +build linux darwin freebsd netbsd openbsd solaris dragonfly windows
  2. package pb
  3. import (
  4. "io"
  5. "sync"
  6. "time"
  7. )
  8. // Create and start new pool with given bars
  9. // You need call pool.Stop() after work
  10. func StartPool(pbs ...*ProgressBar) (pool *Pool, err error) {
  11. pool = new(Pool)
  12. if err = pool.start(); err != nil {
  13. return
  14. }
  15. pool.Add(pbs...)
  16. return
  17. }
  18. type Pool struct {
  19. Output io.Writer
  20. RefreshRate time.Duration
  21. bars []*ProgressBar
  22. lastBarsCount int
  23. quit chan int
  24. m sync.Mutex
  25. finishOnce sync.Once
  26. }
  27. // Add progress bars.
  28. func (p *Pool) Add(pbs ...*ProgressBar) {
  29. p.m.Lock()
  30. defer p.m.Unlock()
  31. for _, bar := range pbs {
  32. bar.ManualUpdate = true
  33. bar.NotPrint = true
  34. bar.Start()
  35. p.bars = append(p.bars, bar)
  36. }
  37. }
  38. func (p *Pool) start() (err error) {
  39. p.RefreshRate = DefaultRefreshRate
  40. quit, err := lockEcho()
  41. if err != nil {
  42. return
  43. }
  44. p.quit = make(chan int)
  45. go p.writer(quit)
  46. return
  47. }
  48. func (p *Pool) writer(finish chan int) {
  49. var first = true
  50. for {
  51. select {
  52. case <-time.After(p.RefreshRate):
  53. if p.print(first) {
  54. p.print(false)
  55. finish <- 1
  56. return
  57. }
  58. first = false
  59. case <-p.quit:
  60. finish <- 1
  61. return
  62. }
  63. }
  64. }
  65. // Restore terminal state and close pool
  66. func (p *Pool) Stop() error {
  67. // Wait until one final refresh has passed.
  68. time.Sleep(p.RefreshRate)
  69. p.finishOnce.Do(func() {
  70. close(p.quit)
  71. })
  72. return unlockEcho()
  73. }