Skip to content

Commit

Permalink
Reverted removal of progress-bar that was done by commit 5d97b10.
Browse files Browse the repository at this point in the history
  • Loading branch information
@arizvisa
arizvisa committed Feb 5, 2018
1 parent 7d5d62d commit cfabf1c
Show file tree
Hide file tree
Showing 24 changed files with 2,724 additions and 37 deletions.
5 changes: 3 additions & 2 deletions common/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package common

import (
"fmt"
"github.com/cheggaaa/pb"
"net/url"
"os"
"path/filepath"
Expand Down Expand Up @@ -54,7 +55,7 @@ func SupportedProtocol(u *url.URL) bool {

// build a dummy NewDownloadClient since this is the only place that valid
// protocols are actually exposed.
cli := NewDownloadClient(&DownloadConfig{})
cli := NewDownloadClient(&DownloadConfig{}, *pb.New(0))

// Iterate through each downloader to see if a protocol was found.
ok := false
Expand Down Expand Up @@ -186,7 +187,7 @@ func FileExistsLocally(original string) bool {

// First create a dummy downloader so we can figure out which
// protocol to use.
cli := NewDownloadClient(&DownloadConfig{})
cli := NewDownloadClient(&DownloadConfig{}, *pb.New(0))
d, ok := cli.config.DownloaderMap[u.Scheme]
if !ok {
return false
Expand Down
48 changes: 34 additions & 14 deletions common/download.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ import (
"strings"
)

// required import for progress-bar
import "github.com/cheggaaa/pb"

// imports related to each Downloader implementation
import (
"io"
Expand Down Expand Up @@ -81,22 +84,23 @@ func HashForType(t string) hash.Hash {

// NewDownloadClient returns a new DownloadClient for the given
// configuration.
func NewDownloadClient(c *DownloadConfig) *DownloadClient {
func NewDownloadClient(c *DownloadConfig, bar pb.ProgressBar) *DownloadClient {
const mtu = 1500 /* ethernet */ - 20 /* ipv4 */ - 20 /* tcp */

// Create downloader map if it hasn't been specified already.
if c.DownloaderMap == nil {
c.DownloaderMap = map[string]Downloader{
"file": &FileDownloader{bufferSize: nil},
"http": &HTTPDownloader{userAgent: c.UserAgent},
"https": &HTTPDownloader{userAgent: c.UserAgent},
"smb": &SMBDownloader{bufferSize: nil},
"file": &FileDownloader{progress: &bar, bufferSize: nil},
"http": &HTTPDownloader{progress: &bar, userAgent: c.UserAgent},
"https": &HTTPDownloader{progress: &bar, userAgent: c.UserAgent},
"smb": &SMBDownloader{progress: &bar, bufferSize: nil},
}
}
return &DownloadClient{config: c}
}

// A downloader implements the ability to transfer, cancel, or resume a file.
// A downloader implements the ability to transfer a file, and cancel or resume
// it.
type Downloader interface {
Resume()
Cancel()
Expand Down Expand Up @@ -205,14 +209,6 @@ func (d *DownloadClient) Get() (string, error) {
return finalPath, err
}

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
// download.
func (d *DownloadClient) VerifyChecksum(path string) (bool, error) {
Expand All @@ -238,6 +234,8 @@ type HTTPDownloader struct {
current uint64
total uint64
userAgent string

progress *pb.ProgressBar
}

func (d *HTTPDownloader) Cancel() {
Expand Down Expand Up @@ -302,6 +300,10 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {

d.total = d.current + uint64(resp.ContentLength)

d.progress.Total = int64(d.total)
progressBar := d.progress.Start()
progressBar.Set64(int64(d.current))

var buffer [4096]byte
for {
n, err := resp.Body.Read(buffer[:])
Expand All @@ -310,6 +312,7 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
}

d.current += uint64(n)
progressBar.Set64(int64(d.current))

if _, werr := dst.Write(buffer[:n]); werr != nil {
return werr
Expand All @@ -319,6 +322,7 @@ func (d *HTTPDownloader) Download(dst *os.File, src *url.URL) error {
break
}
}
progressBar.Finish()
return nil
}

Expand All @@ -338,6 +342,8 @@ type FileDownloader struct {
active bool
current uint64
total uint64

progress *pb.ProgressBar
}

func (d *FileDownloader) Progress() uint64 {
Expand Down Expand Up @@ -428,13 +434,17 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
}
d.total = uint64(fi.Size())

d.progress.Total = int64(d.total)
progressBar := d.progress.Start()

// no bufferSize specified, so copy synchronously.
if d.bufferSize == nil {
var n int64
n, err = io.Copy(dst, f)
d.active = false

d.current += uint64(n)
progressBar.Set64(int64(d.current))

// use a goro in case someone else wants to enable cancel/resume
} else {
Expand All @@ -447,6 +457,7 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
}

d.current += uint64(n)
progressBar.Set64(int64(d.current))
}
d.active = false
e <- err
Expand All @@ -455,6 +466,7 @@ func (d *FileDownloader) Download(dst *os.File, src *url.URL) error {
// ...and we spin until it's done
err = <-errch
}
progressBar.Finish()
f.Close()
return err
}
Expand All @@ -467,6 +479,8 @@ type SMBDownloader struct {
active bool
current uint64
total uint64

progress *pb.ProgressBar
}

func (d *SMBDownloader) Progress() uint64 {
Expand Down Expand Up @@ -539,13 +553,17 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
}
d.total = uint64(fi.Size())

d.progress.Total = int64(d.total)
progressBar := d.progress.Start()

// no bufferSize specified, so copy synchronously.
if d.bufferSize == nil {
var n int64
n, err = io.Copy(dst, f)
d.active = false

d.current += uint64(n)
progressBar.Set64(int64(d.current))

// use a goro in case someone else wants to enable cancel/resume
} else {
Expand All @@ -558,6 +576,7 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
}

d.current += uint64(n)
progressBar.Set64(int64(d.current))
}
d.active = false
e <- err
Expand All @@ -566,6 +585,7 @@ func (d *SMBDownloader) Download(dst *os.File, src *url.URL) error {
// ...and as usual we spin until it's done
err = <-errch
}
progressBar.Finish()
f.Close()
return err
}
22 changes: 12 additions & 10 deletions common/download_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ import (
"runtime"
"strings"
"testing"

"github.com/cheggaaa/pb"
)

func TestDownloadClientVerifyChecksum(t *testing.T) {
Expand All @@ -36,7 +38,7 @@ func TestDownloadClientVerifyChecksum(t *testing.T) {
Checksum: checksum,
}

d := NewDownloadClient(config)
d := NewDownloadClient(config, *pb.New64(0))
result, err := d.VerifyChecksum(tf.Name())
if err != nil {
t.Fatalf("Verify err: %s", err)
Expand All @@ -59,7 +61,7 @@ func TestDownloadClient_basic(t *testing.T) {
Url: ts.URL + "/basic.txt",
TargetPath: tf.Name(),
CopyFile: true,
})
}, *pb.New64(0))

path, err := client.Get()
if err != nil {
Expand Down Expand Up @@ -95,7 +97,7 @@ func TestDownloadClient_checksumBad(t *testing.T) {
Hash: HashForType("md5"),
Checksum: checksum,
CopyFile: true,
})
}, *pb.New64(0))
if _, err := client.Get(); err == nil {
t.Fatal("should error")
}
Expand All @@ -120,7 +122,7 @@ func TestDownloadClient_checksumGood(t *testing.T) {
Hash: HashForType("md5"),
Checksum: checksum,
CopyFile: true,
})
}, *pb.New64(0))
path, err := client.Get()
if err != nil {
t.Fatalf("err: %s", err)
Expand Down Expand Up @@ -151,7 +153,7 @@ func TestDownloadClient_checksumNoDownload(t *testing.T) {
Hash: HashForType("md5"),
Checksum: checksum,
CopyFile: true,
})
}, *pb.New64(0))
path, err := client.Get()
if err != nil {
t.Fatalf("err: %s", err)
Expand Down Expand Up @@ -190,7 +192,7 @@ func TestDownloadClient_resume(t *testing.T) {
Url: ts.URL,
TargetPath: tf.Name(),
CopyFile: true,
})
}, *pb.New64(0))
path, err := client.Get()
if err != nil {
t.Fatalf("err: %s", err)
Expand Down Expand Up @@ -250,7 +252,7 @@ func TestDownloadClient_usesDefaultUserAgent(t *testing.T) {
CopyFile: true,
}

client := NewDownloadClient(config)
client := NewDownloadClient(config, *pb.New64(0))
_, err = client.Get()
if err != nil {
t.Fatal(err)
Expand Down Expand Up @@ -282,7 +284,7 @@ func TestDownloadClient_setsUserAgent(t *testing.T) {
CopyFile: true,
}

client := NewDownloadClient(config)
client := NewDownloadClient(config, *pb.New64(0))
_, err = client.Get()
if err != nil {
t.Fatal(err)
Expand Down Expand Up @@ -381,7 +383,7 @@ func TestDownloadFileUrl(t *testing.T) {
CopyFile: false,
}

client := NewDownloadClient(config)
client := NewDownloadClient(config, *pb.New64(0))

// Verify that we fail to match the checksum
_, err = client.Get()
Expand Down Expand Up @@ -412,7 +414,7 @@ func SimulateFileUriDownload(t *testing.T, uri string) (string, error) {
}

// go go go
client := NewDownloadClient(config)
client := NewDownloadClient(config, *pb.New64(0))
path, err := client.Get()

// ignore any non-important checksum errors if it's not a unc path
Expand Down
45 changes: 34 additions & 11 deletions common/step_download.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"log"
"time"

"github.com/cheggaaa/pb"
"github.com/hashicorp/packer/helper/multistep"
"github.com/hashicorp/packer/packer"
)
Expand Down Expand Up @@ -46,6 +47,23 @@ type StepDownload struct {
Extension string
}

// Default progress bar appearance
func GetDefaultProgressBar() pb.ProgressBar {
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)

return *bar
}

func (s *StepDownload) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
cache := state.Get("cache").(packer.Cache)
ui := state.Get("ui").(packer.Ui)
Expand All @@ -62,6 +80,10 @@ func (s *StepDownload) Run(_ context.Context, state multistep.StateBag) multiste

ui.Say(fmt.Sprintf("Downloading or copying %s", s.Description))

// Get a default-looking progress bar and connect it to the ui.
bar := GetDefaultProgressBar()
bar.Callback = ui.Message

// First try to use any already downloaded file
// If it fails, proceed to regualar download logic

Expand Down Expand Up @@ -95,7 +117,7 @@ func (s *StepDownload) Run(_ context.Context, state multistep.StateBag) multiste
}
downloadConfigs[i] = config

if match, _ := NewDownloadClient(config).VerifyChecksum(config.TargetPath); match {
if match, _ := NewDownloadClient(config, bar).VerifyChecksum(config.TargetPath); match {
ui.Message(fmt.Sprintf("Found already downloaded, initial checksum matched, no download needed: %s", url))
finalPath = config.TargetPath
break
Expand Down Expand Up @@ -140,7 +162,13 @@ func (s *StepDownload) Cleanup(multistep.StateBag) {}
func (s *StepDownload) download(config *DownloadConfig, state multistep.StateBag) (string, error, bool) {
var path string
ui := state.Get("ui").(packer.Ui)
download := NewDownloadClient(config)

// Get a default looking progress bar and connect it to the ui.
bar := GetDefaultProgressBar()
bar.Callback = ui.Message

// Create download client with config and progress bar
download := NewDownloadClient(config, bar)

downloadCompleteCh := make(chan error, 1)
go func() {
Expand All @@ -149,24 +177,19 @@ func (s *StepDownload) download(config *DownloadConfig, state multistep.StateBag
downloadCompleteCh <- err
}()

progressTicker := time.NewTicker(5 * time.Second)
defer progressTicker.Stop()

for {
select {
case err := <-downloadCompleteCh:
bar.Finish()

if err != nil {
return "", err, 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):
if _, ok := state.GetOk(multistep.StateCancelled); ok {
bar.Finish()
ui.Say("Interrupt received. Cancelling download...")
return "", nil, false
}
Expand Down
Loading

0 comments on commit cfabf1c

Please sign in to comment.