Merge pull request #6700 from hashicorp/multi_progressbar

Split progressbars per object
This commit is contained in:
Megan Marsh 2018-09-19 15:22:14 -07:00 committed by GitHub
commit 5470c08345
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 132 additions and 104 deletions

View File

@ -95,7 +95,7 @@ func NewDownloadClient(c *DownloadConfig, ui packer.Ui) *DownloadClient {
type Downloader interface {
Resume()
Cancel()
ProgressBar() packer.ProgressBar
ProgressBar(identifier string) packer.ProgressBar
}
// A LocalDownloader is responsible for converting a uri to a local path
@ -315,7 +315,7 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
total := current + resp.ContentLength
bar := d.ProgressBar()
bar := d.ProgressBar(src.String())
bar.Start(total)
defer bar.Finish()
bar.Add(current)
@ -427,7 +427,7 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
return err
}
bar := d.ProgressBar()
bar := d.ProgressBar(src.String())
bar.Start(fi.Size())
defer bar.Finish()
@ -529,7 +529,7 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
return err
}
bar := d.ProgressBar()
bar := d.ProgressBar(src.String())
bar.Start(fi.Size())
defer bar.Finish()
@ -561,6 +561,14 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
return err
}
func (d *HTTPDownloader) ProgressBar() packer.ProgressBar { return d.Ui.ProgressBar() }
func (d *FileDownloader) ProgressBar() packer.ProgressBar { return d.Ui.ProgressBar() }
func (d *SMBDownloader) ProgressBar() packer.ProgressBar { return d.Ui.ProgressBar() }
func (d *HTTPDownloader) ProgressBar(identifier string) packer.ProgressBar {
return d.Ui.ProgressBar(identifier)
}
func (d *FileDownloader) ProgressBar(identifier string) packer.ProgressBar {
return d.Ui.ProgressBar(identifier)
}
func (d *SMBDownloader) ProgressBar(identifier string) packer.ProgressBar {
return d.Ui.ProgressBar(identifier)
}

View File

@ -1,10 +1,8 @@
package packer
import (
"fmt"
"io"
"sync"
"sync/atomic"
"github.com/cheggaaa/pb"
)
@ -18,69 +16,51 @@ type ProgressBar interface {
Finish()
}
// StackableProgressBar is a progress bar that
// allows to track multiple downloads at once.
// Every call to Start increments a counter that
// will display the number of current loadings.
// Every call to Start will add total to an internal
// total that is the total displayed.
// First call to Start will start a goroutine
// that is waiting for every download to be finished.
// Last call to Finish triggers a cleanup.
// When all active downloads are finished
// StackableProgressBar will clean itself to a default
// state.
// StackableProgressBar is a progress bar pool that
// allows to track multiple advencments at once,
// by displaying multiple bars.
type StackableProgressBar struct {
mtx sync.Mutex // locks in Start & Finish
BasicProgressBar
items int32
total int64
started bool
pool *pb.Pool
bars []*BasicProgressBar
wg sync.WaitGroup
}
var _ ProgressBar = new(StackableProgressBar)
func (spb *StackableProgressBar) cleanup() {
spb.wg.Wait()
func (spb *StackableProgressBar) start() {
spb.BasicProgressBar.ProgressBar = pb.New(0)
spb.BasicProgressBar.ProgressBar.SetUnits(pb.U_BYTES)
spb.BasicProgressBar.ProgressBar.Start()
spb.started = true
}
func (spb *StackableProgressBar) Start(total int64) {
spb.mtx.Lock()
spb.total += total
spb.items++
if !spb.started {
spb.start()
}
spb.SetTotal64(spb.total)
spb.prefix()
spb.mtx.Unlock()
}
func (spb *StackableProgressBar) prefix() {
spb.BasicProgressBar.ProgressBar.Prefix(fmt.Sprintf("%d items: ", atomic.LoadInt32(&spb.items)))
}
func (spb *StackableProgressBar) Finish() {
spb.mtx.Lock()
defer spb.mtx.Unlock()
spb.items--
if spb.items == 0 {
// slef cleanup
spb.BasicProgressBar.ProgressBar.Finish()
spb.BasicProgressBar.ProgressBar = nil
spb.started = false
spb.total = 0
return
spb.pool.Stop()
spb.pool = nil
spb.bars = nil
}
func (spb *StackableProgressBar) New(identifier string) ProgressBar {
spb.mtx.Lock()
spb.wg.Add(1)
defer spb.mtx.Unlock()
start := false
if spb.pool == nil {
spb.pool = pb.NewPool()
go spb.cleanup()
start = true
}
spb.prefix()
bar := NewProgressBar(identifier)
bar.Prefix(identifier)
bar.finishCb = spb.wg.Done
spb.bars = append(spb.bars, bar)
spb.pool.Add(bar.ProgressBar)
if start {
spb.pool.Start()
}
return bar
}
// BasicProgressBar is packer's basic progress bar.
@ -88,6 +68,13 @@ func (spb *StackableProgressBar) Finish() {
// itself at the bottom of a terminal.
type BasicProgressBar struct {
*pb.ProgressBar
finishCb func()
}
func NewProgressBar(identifier string) *BasicProgressBar {
bar := new(BasicProgressBar)
bar.ProgressBar = pb.New(0)
return bar
}
var _ ProgressBar = new(BasicProgressBar)
@ -97,6 +84,13 @@ func (bpb *BasicProgressBar) Start(total int64) {
bpb.ProgressBar.Start()
}
func (bpb *BasicProgressBar) Finish() {
if bpb.finishCb != nil {
bpb.finishCb()
}
bpb.ProgressBar.Finish()
}
func (bpb *BasicProgressBar) Add(current int64) {
bpb.ProgressBar.Add64(current)
}

View File

@ -64,28 +64,36 @@ func (u *Ui) Say(message string) {
}
}
func (u *Ui) ProgressBar() packer.ProgressBar {
if err := u.client.Call("Ui.ProgressBar", new(interface{}), new(interface{})); err != nil {
log.Printf("Error in Ui RPC call: %s", err)
func (u *Ui) ProgressBar(identifier string) packer.ProgressBar {
if err := u.client.Call("Ui.ProgressBar", identifier, new(interface{})); err != nil {
log.Printf("Err or in Ui RPC call: %s", err)
}
return &RemoteProgressBarClient{
id: identifier,
client: u.client,
}
return u // Ui is also a progress bar !!
}
var _ packer.ProgressBar = new(Ui)
func (pb *Ui) Start(total int64) {
pb.client.Call("Ui.Start", total, new(interface{}))
type RemoteProgressBarClient struct {
id string
client *rpc.Client
}
func (pb *Ui) Add(current int64) {
pb.client.Call("Ui.Add", current, new(interface{}))
var _ packer.ProgressBar = new(RemoteProgressBarClient)
func (pb *RemoteProgressBarClient) Start(total int64) {
pb.client.Call(pb.id+".Start", total, new(interface{}))
}
func (pb *Ui) Finish() {
pb.client.Call("Ui.Finish", nil, new(interface{}))
func (pb *RemoteProgressBarClient) Add(current int64) {
pb.client.Call(pb.id+".Add", current, new(interface{}))
}
func (pb *Ui) NewProxyReader(r io.Reader) io.Reader {
func (pb *RemoteProgressBarClient) Finish() {
pb.client.Call(pb.id+".Finish", nil, new(interface{}))
}
func (pb *RemoteProgressBarClient) NewProxyReader(r io.Reader) io.Reader {
return &packer.ProxyReader{Reader: r, ProgressBar: pb}
}
@ -121,25 +129,38 @@ func (u *UiServer) Say(message *string, reply *interface{}) error {
return nil
}
func (u *UiServer) ProgressBar(_ *string, reply *interface{}) error {
// No-op for now, this function might be
// used in the future if we want to use
// different progress bars with identifiers.
u.ui.ProgressBar()
// ProgressBar registers a rpc progress bar server identified by identifier.
// ProgressBar expects identifiers to be unique across runs
// since for examples an iso download should be cached.
func (u *UiServer) ProgressBar(identifier string, reply *interface{}) error {
bar := u.ui.ProgressBar(identifier)
log.Printf("registering progressbar for '%s'", identifier)
err := u.register(identifier, &UiProgressBarServer{bar})
if err != nil {
log.Printf("failed to register a new progress bar rpc server, %s", err)
return err
}
*reply = identifier
return nil
}
func (pb *UiServer) Finish(_ string, _ *interface{}) error {
pb.ui.ProgressBar().Finish()
type UiProgressBarServer struct {
bar packer.ProgressBar
}
func (pb *UiProgressBarServer) Finish(_ string, _ *interface{}) error {
pb.bar.Finish()
return nil
}
func (pb *UiServer) Start(total int64, _ *interface{}) error {
pb.ui.ProgressBar().Start(total)
func (pb *UiProgressBarServer) Start(total int64, _ *interface{}) error {
pb.bar.Start(total)
return nil
}
func (pb *UiServer) Add(current int64, _ *interface{}) error {
pb.ui.ProgressBar().Add(current)
func (pb *UiProgressBarServer) Add(current int64, _ *interface{}) error {
pb.bar.Add(current)
return nil
}

View File

@ -55,7 +55,7 @@ func (u *testUi) Say(message string) {
u.sayMessage = message
}
func (u *testUi) ProgressBar() packer.ProgressBar {
func (u *testUi) ProgressBar(string) packer.ProgressBar {
u.progressBarCalled = true
return u
}
@ -119,7 +119,7 @@ func TestUiRPC(t *testing.T) {
t.Fatalf("bad: %#v", ui.errorMessage)
}
bar := uiClient.ProgressBar()
bar := uiClient.ProgressBar("test")
if ui.progressBarCalled != true {
t.Errorf("ProgressBar not called.")
}

View File

@ -37,19 +37,24 @@ type Ui interface {
Message(string)
Error(string)
Machine(string, ...string)
ProgressBar() ProgressBar
// ProgressBar instanciates a ProgressBar for
// an object named 'identifer'.
// ProgressBar should not be called twice with
// the same identifier.
ProgressBar(identifier string) ProgressBar
}
type NoopUi struct{}
var _ Ui = new(NoopUi)
func (*NoopUi) Ask(string) (string, error) { return "", errors.New("this is a noop ui") }
func (*NoopUi) Say(string) { return }
func (*NoopUi) Message(string) { return }
func (*NoopUi) Error(string) { return }
func (*NoopUi) Machine(string, ...string) { return }
func (*NoopUi) ProgressBar() ProgressBar { return new(NoopProgressBar) }
func (*NoopUi) Ask(string) (string, error) { return "", errors.New("this is a noop ui") }
func (*NoopUi) Say(string) { return }
func (*NoopUi) Message(string) { return }
func (*NoopUi) Error(string) { return }
func (*NoopUi) Machine(string, ...string) { return }
func (*NoopUi) ProgressBar(string) ProgressBar { return new(NoopProgressBar) }
// ColoredUi is a UI that is colored using terminal colors.
type ColoredUi struct {
@ -87,8 +92,8 @@ type BasicUi struct {
var _ Ui = new(BasicUi)
func (bu *BasicUi) ProgressBar() ProgressBar {
return &bu.StackableProgressBar
func (bu *BasicUi) ProgressBar(identifier string) ProgressBar {
return bu.StackableProgressBar.New(identifier)
}
// MachineReadableUi is a UI that only outputs machine-readable output
@ -125,8 +130,8 @@ func (u *ColoredUi) Machine(t string, args ...string) {
u.Ui.Machine(t, args...)
}
func (u *ColoredUi) ProgressBar() ProgressBar {
return u.Ui.ProgressBar() //TODO(adrien): color me
func (u *ColoredUi) ProgressBar(identifier string) ProgressBar {
return u.Ui.ProgressBar(identifier) //TODO(adrien): color me
}
func (u *ColoredUi) colorize(message string, color UiColor, bold bool) string {
@ -182,8 +187,8 @@ func (u *TargetedUI) Machine(t string, args ...string) {
u.Ui.Machine(fmt.Sprintf("%s,%s", u.Target, t), args...)
}
func (u *TargetedUI) ProgressBar() ProgressBar {
return u.Ui.ProgressBar()
func (u *TargetedUI) ProgressBar(identifier string) ProgressBar {
return u.Ui.ProgressBar(identifier)
}
func (u *TargetedUI) prefixLines(arrow bool, message string) string {
@ -339,6 +344,6 @@ func (u *MachineReadableUi) Machine(category string, args ...string) {
}
}
func (u *MachineReadableUi) ProgressBar() ProgressBar {
func (u *MachineReadableUi) ProgressBar(string) ProgressBar {
return new(NoopProgressBar)
}

View File

@ -601,6 +601,6 @@ func (ui *Ui) Machine(t string, args ...string) {
<-ui.sem
}
func (ui *Ui) ProgressBar() packer.ProgressBar {
func (ui *Ui) ProgressBar(string) packer.ProgressBar {
return new(packer.NoopProgressBar)
}

View File

@ -176,7 +176,7 @@ func (p *Provisioner) ProvisionUpload(ui packer.Ui, comm packer.Communicator) er
}
// Get a default progress bar
bar := ui.ProgressBar()
bar := ui.ProgressBar(fmt.Sprintf("Uploading %s => %s", src, dst))
bar.Start(info.Size())
defer bar.Finish()