Consolidated the progress-bar's format into common/step_download.go. Removed DownloadClient's PercentProgress callback since cheggaaa's progress-bar already does that.

This commit is contained in:
Ali Rizvi-Santiago 2017-12-06 17:57:03 -06:00
parent d85883582f
commit 69e5eec1ce
2 changed files with 91 additions and 66 deletions

View File

@ -83,19 +83,19 @@ func HashForType(t string) hash.Hash {
// NewDownloadClient returns a new DownloadClient for the given // NewDownloadClient returns a new DownloadClient for the given
// configuration. // configuration.
func NewDownloadClient(c *DownloadConfig) *DownloadClient { func NewDownloadClient(c *DownloadConfig, bar *pb.ProgressBar) *DownloadClient {
const mtu = 1500 /* ethernet */ - 20 /* ipv4 */ - 20 /* tcp */ const mtu = 1500 /* ethernet */ - 20 /* ipv4 */ - 20 /* tcp */
if c.DownloaderMap == nil { if c.DownloaderMap == nil {
// Create downloader map
c.DownloaderMap = map[string]Downloader{ c.DownloaderMap = map[string]Downloader{
"file": &FileDownloader{bufferSize: nil}, "file": &FileDownloader{progress: bar, bufferSize: nil},
"ftp": &FTPDownloader{userInfo: url.UserPassword("anonymous", "anonymous@"), mtu: mtu}, "ftp": &FTPDownloader{progress: bar, userInfo: url.UserPassword("anonymous", "anonymous@"), mtu: mtu},
"http": &HTTPDownloader{userAgent: c.UserAgent}, "http": &HTTPDownloader{progress: bar, userAgent: c.UserAgent},
"https": &HTTPDownloader{userAgent: c.UserAgent}, "https": &HTTPDownloader{progress: bar, userAgent: c.UserAgent},
"smb": &SMBDownloader{bufferSize: nil}, "smb": &SMBDownloader{progress: bar, bufferSize: nil},
} }
} }
return &DownloadClient{config: c} return &DownloadClient{config: c}
} }
@ -209,14 +209,6 @@ func (d *DownloadClient) Get() (string, error) {
return finalPath, err return finalPath, err
} }
// PercentProgress returns the download progress as a percentage.
func (d *DownloadClient) PercentProgress() int {
if d.downloader == nil {
return -1
}
return int((float64(d.downloader.Progress()) / float64(d.downloader.Total())) * 100)
}
// VerifyChecksum tests that the path matches the checksum for the // VerifyChecksum tests that the path matches the checksum for the
// download. // download.
func (d *DownloadClient) VerifyChecksum(path string) (bool, error) { func (d *DownloadClient) VerifyChecksum(path string) (bool, error) {
@ -239,9 +231,11 @@ func (d *DownloadClient) VerifyChecksum(path string) (bool, error) {
// HTTPDownloader is an implementation of Downloader that downloads // HTTPDownloader is an implementation of Downloader that downloads
// files over HTTP. // files over HTTP.
type HTTPDownloader struct { type HTTPDownloader struct {
progress uint64 current uint64
total uint64 total uint64
userAgent string userAgent string
progress *pb.ProgressBar
} }
func (d *HTTPDownloader) Cancel() { func (d *HTTPDownloader) Cancel() {
@ -261,7 +255,7 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
} }
// Reset our progress // Reset our progress
d.progress = 0 d.current = 0
// Make the request. We first make a HEAD request so we can check // Make the request. We first make a HEAD request so we can check
// if the server supports range queries. If the server/URL doesn't // if the server supports range queries. If the server/URL doesn't
@ -290,7 +284,7 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
if _, err = dst.Seek(0, os.SEEK_END); err == nil { if _, err = dst.Seek(0, os.SEEK_END); err == nil {
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", fi.Size())) req.Header.Set("Range", fmt.Sprintf("bytes=%d-", fi.Size()))
d.progress = uint64(fi.Size()) d.current = uint64(fi.Size())
} }
} }
} }
@ -304,9 +298,11 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
return err return err
} }
d.total = d.progress + uint64(resp.ContentLength) d.total = d.current + uint64(resp.ContentLength)
progressBar := pb.New64(int64(d.Total())).Start()
progressBar.Set64(int64(d.Progress())) d.progress.Total = int64(d.total)
progressBar := d.progress.Start()
progressBar.Set64(int64(d.current))
var buffer [4096]byte var buffer [4096]byte
for { for {
@ -315,8 +311,8 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
return err return err
} }
d.progress += uint64(n) d.current += uint64(n)
progressBar.Set64(int64(d.Progress())) progressBar.Set64(int64(d.current))
if _, werr := dst.Write(buffer[:n]); werr != nil { if _, werr := dst.Write(buffer[:n]); werr != nil {
return werr return werr
@ -326,12 +322,12 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
break break
} }
} }
progressBar.Finish()
return nil return nil
} }
func (d *HTTPDownloader) Progress() uint64 { func (d *HTTPDownloader) Progress() uint64 {
return d.progress return d.current
} }
func (d *HTTPDownloader) Total() uint64 { func (d *HTTPDownloader) Total() uint64 {
@ -345,12 +341,14 @@ type FTPDownloader struct {
mtu uint mtu uint
active bool active bool
progress uint64 current uint64
total uint64 total uint64
progress *pb.ProgressBar
} }
func (d *FTPDownloader) Progress() uint64 { func (d *FTPDownloader) Progress() uint64 {
return d.progress return d.current
} }
func (d *FTPDownloader) Total() uint64 { func (d *FTPDownloader) Total() uint64 {
@ -438,13 +436,15 @@ func (d *FTPDownloader) Download(dst *os.File, src *url.URL) error {
} }
log.Printf("Found file : %s : %v bytes\n", entry.Name, entry.Size) log.Printf("Found file : %s : %v bytes\n", entry.Name, entry.Size)
d.progress = 0 d.current = 0
d.total = entry.Size d.total = entry.Size
progressBar := pb.New64(int64(d.Total())).Start()
d.progress.Total = int64(d.total)
progressBar := d.progress.Start()
// download specified file // download specified file
d.active = true d.active = true
reader, err := cli.RetrFrom(uri.Path, d.progress) reader, err := cli.RetrFrom(uri.Path, d.current)
if err != nil { if err != nil {
return nil return nil
} }
@ -458,8 +458,8 @@ func (d *FTPDownloader) Download(dst *os.File, src *url.URL) error {
break break
} }
d.progress += uint64(n) d.current += uint64(n)
progressBar.Set64(int64(d.Progress())) progressBar.Set64(int64(d.current))
} }
d.active = false d.active = false
e <- err e <- err
@ -467,10 +467,12 @@ func (d *FTPDownloader) Download(dst *os.File, src *url.URL) error {
// spin until it's done // spin until it's done
err = <-errch err = <-errch
progressBar.Finish()
reader.Close() reader.Close()
if err == nil && d.progress != d.total { if err == nil && d.current != d.total {
err = fmt.Errorf("FTP total transfer size was %d when %d was expected", d.progress, d.total) err = fmt.Errorf("FTP total transfer size was %d when %d was expected", d.current, d.total)
} }
// log out // log out
@ -484,12 +486,14 @@ type FileDownloader struct {
bufferSize *uint bufferSize *uint
active bool active bool
progress uint64 current uint64
total uint64 total uint64
progress *pb.ProgressBar
} }
func (d *FileDownloader) Progress() uint64 { func (d *FileDownloader) Progress() uint64 {
return d.progress return d.current
} }
func (d *FileDownloader) Total() uint64 { func (d *FileDownloader) Total() uint64 {
@ -549,7 +553,7 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
} }
/* download the file using the operating system's facilities */ /* download the file using the operating system's facilities */
d.progress = 0 d.current = 0
d.active = true d.active = true
f, err := os.Open(realpath) f, err := os.Open(realpath)
@ -564,7 +568,9 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
return err return err
} }
d.total = uint64(fi.Size()) d.total = uint64(fi.Size())
progressBar := pb.New64(int64(d.Total())).Start()
d.progress.Total = int64(d.total)
progressBar := d.progress.Start()
// no bufferSize specified, so copy synchronously. // no bufferSize specified, so copy synchronously.
if d.bufferSize == nil { if d.bufferSize == nil {
@ -572,8 +578,8 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
n, err = io.Copy(dst, f) n, err = io.Copy(dst, f)
d.active = false d.active = false
d.progress += uint64(n) d.current += uint64(n)
progressBar.Set64(int64(d.Progress())) progressBar.Set64(int64(d.current))
// use a goro in case someone else wants to enable cancel/resume // use a goro in case someone else wants to enable cancel/resume
} else { } else {
@ -585,8 +591,8 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
break break
} }
d.progress += uint64(n) d.current += uint64(n)
progressBar.Set64(int64(d.Progress())) progressBar.Set64(int64(d.current))
} }
d.active = false d.active = false
e <- err e <- err
@ -595,6 +601,7 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
// ...and we spin until it's done // ...and we spin until it's done
err = <-errch err = <-errch
} }
progressBar.Finish()
f.Close() f.Close()
return err return err
} }
@ -605,12 +612,14 @@ type SMBDownloader struct {
bufferSize *uint bufferSize *uint
active bool active bool
progress uint64 current uint64
total uint64 total uint64
progress *pb.ProgressBar
} }
func (d *SMBDownloader) Progress() uint64 { func (d *SMBDownloader) Progress() uint64 {
return d.progress return d.current
} }
func (d *SMBDownloader) Total() uint64 { func (d *SMBDownloader) Total() uint64 {
@ -663,7 +672,7 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
} }
/* Open up the "\\"-prefixed path using the Windows filesystem */ /* Open up the "\\"-prefixed path using the Windows filesystem */
d.progress = 0 d.current = 0
d.active = true d.active = true
f, err := os.Open(realpath) f, err := os.Open(realpath)
@ -678,7 +687,9 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
return err return err
} }
d.total = uint64(fi.Size()) d.total = uint64(fi.Size())
progressBar := pb.New64(int64(d.Total())).Start()
d.progress.Total = int64(d.total)
progressBar := d.progress.Start()
// no bufferSize specified, so copy synchronously. // no bufferSize specified, so copy synchronously.
if d.bufferSize == nil { if d.bufferSize == nil {
@ -686,8 +697,8 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
n, err = io.Copy(dst, f) n, err = io.Copy(dst, f)
d.active = false d.active = false
d.progress += uint64(n) d.current += uint64(n)
progressBar.Set64(int64(d.Progress())) progressBar.Set64(int64(d.current))
// use a goro in case someone else wants to enable cancel/resume // use a goro in case someone else wants to enable cancel/resume
} else { } else {
@ -699,8 +710,8 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
break break
} }
d.progress += uint64(n) d.current += uint64(n)
progressBar.Set64(int64(d.Progress())) progressBar.Set64(int64(d.current))
} }
d.active = false d.active = false
e <- err e <- err
@ -709,6 +720,7 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
// ...and as usual we spin until it's done // ...and as usual we spin until it's done
err = <-errch err = <-errch
} }
progressBar.Finish()
f.Close() f.Close()
return err return err
} }

View File

@ -9,6 +9,8 @@ import (
"github.com/hashicorp/packer/packer" "github.com/hashicorp/packer/packer"
"github.com/mitchellh/multistep" "github.com/mitchellh/multistep"
"gopkg.in/cheggaaa/pb.v1"
) )
// StepDownload downloads a remote file using the download client within // StepDownload downloads a remote file using the download client within
@ -139,7 +141,23 @@ func (s *StepDownload) Cleanup(multistep.StateBag) {}
func (s *StepDownload) download(config *DownloadConfig, state multistep.StateBag) (string, error, bool) { func (s *StepDownload) download(config *DownloadConfig, state multistep.StateBag) (string, error, bool) {
var path string var path string
ui := state.Get("ui").(packer.Ui) ui := state.Get("ui").(packer.Ui)
download := NewDownloadClient(config)
// design the appearance of the progress bar
bar := pb.New64(0)
bar.ShowPercent = true
bar.ShowCounters = true
bar.ShowSpeed = false
bar.ShowBar = true
bar.ShowTimeLeft = false
bar.ShowFinalTime = false
bar.SetUnits(pb.U_BYTES)
bar.Format("[=>-]")
bar.SetRefreshRate(1 * time.Second)
bar.SetWidth(25)
bar.Callback = ui.Message
// create download client with config and progress bar
download := NewDownloadClient(config, bar)
downloadCompleteCh := make(chan error, 1) downloadCompleteCh := make(chan error, 1)
go func() { go func() {
@ -148,24 +166,19 @@ func (s *StepDownload) download(config *DownloadConfig, state multistep.StateBag
downloadCompleteCh <- err downloadCompleteCh <- err
}() }()
progressTicker := time.NewTicker(5 * time.Second)
defer progressTicker.Stop()
for { for {
select { select {
case err := <-downloadCompleteCh: case err := <-downloadCompleteCh:
bar.Finish()
if err != nil { if err != nil {
return "", err, true return "", err, true
} }
return path, nil, true return path, nil, true
case <-progressTicker.C:
progress := download.PercentProgress()
if progress >= 0 {
ui.Message(fmt.Sprintf("Download progress: %d%%", progress))
}
case <-time.After(1 * time.Second): case <-time.After(1 * time.Second):
if _, ok := state.GetOk(multistep.StateCancelled); ok { if _, ok := state.GetOk(multistep.StateCancelled); ok {
bar.Finish()
ui.Say("Interrupt received. Cancelling download...") ui.Say("Interrupt received. Cancelling download...")
return "", nil, false return "", nil, false
} }