packer-cn/packer/progressbar.go

75 lines
1.3 KiB
Go

// +build !solaris
package packer
import (
"io"
"path/filepath"
"sync"
pb "github.com/cheggaaa/pb"
)
func ProgressBarConfig(bar *pb.ProgressBar, prefix string) {
bar.SetUnits(pb.U_BYTES)
bar.Prefix(prefix)
}
// UiProgressBar is a progress bar compatible with go-getter used in our
// UI structs.
type UiProgressBar struct {
lock sync.Mutex
pool *pb.Pool
pbs int
}
func (p *UiProgressBar) TrackProgress(src string, currentSize, totalSize int64, stream io.ReadCloser) io.ReadCloser {
if p == nil {
return stream
}
p.lock.Lock()
defer p.lock.Unlock()
newPb := pb.New64(totalSize)
newPb.Set64(currentSize)
ProgressBarConfig(newPb, filepath.Base(src))
if p.pool == nil {
pool := pb.NewPool()
err := pool.Start()
if err != nil {
// here, we probably cannot lock
// stdout, so let's just return
// stream to avoid any error.
return stream
}
p.pool = pool
}
p.pool.Add(newPb)
reader := newPb.NewProxyReader(stream)
p.pbs++
return &readCloser{
Reader: reader,
close: func() error {
p.lock.Lock()
defer p.lock.Unlock()
newPb.Finish()
p.pbs--
if p.pbs <= 0 {
p.pool.Stop()
p.pool = nil
}
return nil
},
}
}
type readCloser struct {
io.Reader
close func() error
}
func (c *readCloser) Close() error { return c.close() }