pool.go 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  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. // NewPool initialises a pool with progress bars, but
  19. // doesn't start it. You need to call Start manually
  20. func NewPool(pbs ...*ProgressBar) (pool *Pool) {
  21. pool = new(Pool)
  22. pool.Add(pbs...)
  23. return
  24. }
  25. type Pool struct {
  26. Output io.Writer
  27. RefreshRate time.Duration
  28. bars []*ProgressBar
  29. lastBarsCount int
  30. shutdownCh chan struct{}
  31. workerCh chan struct{}
  32. m sync.Mutex
  33. finishOnce sync.Once
  34. }
  35. // Add progress bars.
  36. func (p *Pool) Add(pbs ...*ProgressBar) {
  37. p.m.Lock()
  38. defer p.m.Unlock()
  39. for _, bar := range pbs {
  40. bar.ManualUpdate = true
  41. bar.NotPrint = true
  42. bar.Start()
  43. p.bars = append(p.bars, bar)
  44. }
  45. }
  46. func (p *Pool) Start() (err error) {
  47. p.RefreshRate = DefaultRefreshRate
  48. p.shutdownCh, err = lockEcho()
  49. if err != nil {
  50. return
  51. }
  52. p.workerCh = make(chan struct{})
  53. go p.writer()
  54. return
  55. }
  56. func (p *Pool) writer() {
  57. var first = true
  58. defer func() {
  59. if first == false {
  60. p.print(false)
  61. } else {
  62. p.print(true)
  63. p.print(false)
  64. }
  65. close(p.workerCh)
  66. }()
  67. for {
  68. select {
  69. case <-time.After(p.RefreshRate):
  70. if p.print(first) {
  71. p.print(false)
  72. return
  73. }
  74. first = false
  75. case <-p.shutdownCh:
  76. return
  77. }
  78. }
  79. }
  80. // Restore terminal state and close pool
  81. func (p *Pool) Stop() error {
  82. p.finishOnce.Do(func() {
  83. close(p.shutdownCh)
  84. })
  85. // Wait for the worker to complete
  86. select {
  87. case <-p.workerCh:
  88. }
  89. return unlockEcho()
  90. }