From dfeca9f0dfdfc902a60ef340d6e733d167ba66bd Mon Sep 17 00:00:00 2001 From: Adrien Delorme Date: Wed, 12 Sep 2018 16:21:58 -0700 Subject: [PATCH 1/2] split progressbars per object --- common/download.go | 22 +++++-- packer/progressbar.go | 100 +++++++++++++---------------- packer/rpc/ui.go | 69 +++++++++++++------- packer/rpc/ui_test.go | 4 +- packer/ui.go | 33 ++++++---- provisioner/ansible/provisioner.go | 2 +- provisioner/file/provisioner.go | 2 +- 7 files changed, 128 insertions(+), 104 deletions(-) diff --git a/common/download.go b/common/download.go index 78f31f25d..bea2e8960 100644 --- a/common/download.go +++ b/common/download.go @@ -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) +} diff --git a/packer/progressbar.go b/packer/progressbar.go index fc8f9e1ef..072d6d142 100644 --- a/packer/progressbar.go +++ b/packer/progressbar.go @@ -1,10 +1,8 @@ package packer import ( - "fmt" "io" "sync" - "sync/atomic" "github.com/cheggaaa/pb" ) @@ -18,69 +16,47 @@ 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() + + if spb.pool == nil { + spb.pool = pb.NewPool() + go spb.cleanup() } - spb.prefix() + + bar := NewProgressBar(identifier) + bar.Prefix(identifier) + bar.finishCb = spb.wg.Done + spb.bars = append(spb.bars, bar) + + spb.pool.Add(bar.ProgressBar) + spb.pool.Start() + return bar } // BasicProgressBar is packer's basic progress bar. @@ -88,6 +64,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 +80,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) } diff --git a/packer/rpc/ui.go b/packer/rpc/ui.go index c9d3ef176..d8db6fb80 100644 --- a/packer/rpc/ui.go +++ b/packer/rpc/ui.go @@ -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 } diff --git a/packer/rpc/ui_test.go b/packer/rpc/ui_test.go index 078b85e44..8c70500de 100644 --- a/packer/rpc/ui_test.go +++ b/packer/rpc/ui_test.go @@ -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.") } diff --git a/packer/ui.go b/packer/ui.go index 340c56827..73592c490 100644 --- a/packer/ui.go +++ b/packer/ui.go @@ -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) } diff --git a/provisioner/ansible/provisioner.go b/provisioner/ansible/provisioner.go index 84d3d71b5..38e06e79f 100644 --- a/provisioner/ansible/provisioner.go +++ b/provisioner/ansible/provisioner.go @@ -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) } diff --git a/provisioner/file/provisioner.go b/provisioner/file/provisioner.go index 1be9f1885..0dbb8e5aa 100644 --- a/provisioner/file/provisioner.go +++ b/provisioner/file/provisioner.go @@ -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() From be27775a3c502a7a6fbf9b47b16788d5f26028c2 Mon Sep 17 00:00:00 2001 From: Adrien Delorme Date: Wed, 12 Sep 2018 17:04:10 -0700 Subject: [PATCH 2/2] start the pool only once --- packer/progressbar.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packer/progressbar.go b/packer/progressbar.go index 072d6d142..a709f284a 100644 --- a/packer/progressbar.go +++ b/packer/progressbar.go @@ -44,9 +44,11 @@ func (spb *StackableProgressBar) New(identifier string) ProgressBar { spb.wg.Add(1) defer spb.mtx.Unlock() + start := false if spb.pool == nil { spb.pool = pb.NewPool() go spb.cleanup() + start = true } bar := NewProgressBar(identifier) @@ -55,7 +57,9 @@ func (spb *StackableProgressBar) New(identifier string) ProgressBar { spb.bars = append(spb.bars, bar) spb.pool.Add(bar.ProgressBar) - spb.pool.Start() + if start { + spb.pool.Start() + } return bar }