forked from cheggaaa/pb
/
pool.go
133 lines (118 loc) · 2.34 KB
/
pool.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
// +build linux darwin freebsd netbsd openbsd solaris dragonfly windows plan9 aix
package pb
import (
"errors"
"io"
"sync"
"time"
"github.com/cbehopkins/pb/v3/termutil"
)
// Create and start new pool with given bars
// You need call pool.Stop() after work
func StartPool(pbs ...*ProgressBar) (pool *Pool, err error) {
pool = new(Pool)
if err = pool.Start(); err != nil {
return
}
pool.Add(pbs...)
return
}
// NewPool initialises a pool with progress bars, but
// doesn't start it. You need to call Start manually
func NewPool(pbs ...*ProgressBar) (pool *Pool) {
pool = new(Pool)
pool.Add(pbs...)
return
}
type Pool struct {
Output io.Writer
RefreshRate time.Duration
bars []*ProgressBar
lastBarsCount int
shutdownCh chan struct{}
workerCh chan struct{}
m sync.Mutex
finishOnce sync.Once
width int
}
// Add progress bars.
func (p *Pool) Add(pbs ...*ProgressBar) {
p.m.Lock()
defer p.m.Unlock()
for _, bar := range pbs {
bar.Set(Static, true)
bar.Start()
p.bars = append(p.bars, bar)
}
}
// Remove an active bar
func (p *Pool) Remove(pb *ProgressBar) error {
p.m.Lock()
defer p.m.Unlock()
for i, bar := range p.bars {
if bar == pb {
bar.Finish()
p.bars[i] = p.bars[len(p.bars)-1]
p.bars = p.bars[:len(p.bars)-1]
return nil
}
}
return errors.New("bar not found")
}
// Start the pool, this must be called for things to display
func (p *Pool) Start() (err error) {
p.RefreshRate = defaultRefreshRate
p.shutdownCh, err = termutil.RawModeOn()
if err != nil {
return
}
p.workerCh = make(chan struct{})
go p.writer()
return
}
func (p *Pool) SetWidth(w int) *Pool {
p.m.Lock()
defer p.m.Unlock()
p.width = w
return p
}
func (p *Pool) Width() int {
p.m.Lock()
defer p.m.Unlock()
return p.width
}
func (p *Pool) writer() {
var first = true
defer func() {
if !first {
p.print(false)
} else {
p.print(true)
p.print(false)
}
close(p.workerCh)
}()
for {
select {
case <-time.After(p.RefreshRate):
if p.print(first) {
p.print(false)
return
}
first = false
case <-p.shutdownCh:
return
}
}
}
// Restore terminal state and close pool
func (p *Pool) Stop() error {
p.finishOnce.Do(func() {
if p.shutdownCh != nil {
close(p.shutdownCh)
}
})
// Wait for the worker to complete
<-p.workerCh
return termutil.RawModeOff()
}