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
|
package main
import (
"io"
"path/filepath"
"sync"
"github.com/cheggaaa/pb"
getter "github.com/hashicorp/go-getter"
)
// defaultProgressBar is the default instance of a cheggaaa
// progress bar.
var defaultProgressBar getter.ProgressTracker = &ProgressBar{}
// ProgressBar wraps a github.com/cheggaaa/pb.Pool
// in order to display download progress for one or multiple
// downloads.
//
// If two different instance of ProgressBar try to
// display a progress only one will be displayed.
// It is therefore recommended to use DefaultProgressBar
type ProgressBar struct {
// lock everything below
lock sync.Mutex
pool *pb.Pool
pbs int
}
func ProgressBarConfig(bar *pb.ProgressBar, prefix string) {
bar.SetUnits(pb.U_BYTES)
bar.Prefix(prefix)
}
// TrackProgress instantiates a new progress bar that will
// display the progress of stream until closed.
// total can be 0.
func (cpb *ProgressBar) TrackProgress(src string, currentSize, totalSize int64, stream io.ReadCloser) io.ReadCloser {
cpb.lock.Lock()
defer cpb.lock.Unlock()
newPb := pb.New64(totalSize)
newPb.Set64(currentSize)
ProgressBarConfig(newPb, filepath.Base(src))
if cpb.pool == nil {
cpb.pool = pb.NewPool()
cpb.pool.Start()
}
cpb.pool.Add(newPb)
reader := newPb.NewProxyReader(stream)
cpb.pbs++
return &readCloser{
Reader: reader,
close: func() error {
cpb.lock.Lock()
defer cpb.lock.Unlock()
newPb.Finish()
cpb.pbs--
if cpb.pbs <= 0 {
cpb.pool.Stop()
cpb.pool = nil
}
return nil
},
}
}
type readCloser struct {
io.Reader
close func() error
}
func (c *readCloser) Close() error { return c.close() }
|
