fmt, clean code, use atomic in progress.

Author: melbahja

cmd/got/main.go

@@ -3,10 +3,10 @@ package main
 import (
 	"flag"
 	"fmt"
-	"log"
-
 	"github.com/dustin/go-humanize"
 	"github.com/melbahja/got"
+	"log"
+	"time"
 )
 
 var (
@@ -36,6 +36,7 @@ func main() {
 	if url = flag.Arg(0); url == "" {
 		log.Fatal("Empty download url.")
 	}
+
 	if !(url[:7] == "http://" || url[:8] == "https://") {
 		url = "https://" + url
 	}
@@ -47,15 +48,6 @@ func main() {
 	d := got.Download{
 		URL:  url,
 		Dest: *dest,
-		ProgressFunc: func(i int64, t int64, d *got.Download) {
-			fmt.Printf(
-				"\r\r\b Total Size: %s | Chunk Size: %s | Concurrency: %d | Progress: %s ",
-				humanize.Bytes(uint64(t)),
-				humanize.Bytes(uint64(d.ChunkSize)),
-				d.Concurrency,
-				humanize.Bytes(uint64(i)),
-			)
-		},
 		ChunkSize:   int64(*chunkSize),
 		Interval:    100,
 		Concurrency: *concurrency,
@@ -65,9 +57,24 @@ func main() {
 		log.Fatal(err)
 	}
 
+	// Set progress func to update cli output.
+	d.Progress.ProgressFunc = func(p *got.Progress, d *got.Download) {
+
+		fmt.Printf(
+			"\r\r\bTotal: %s | Chunk: %s | Concurrency: %d | Received: %s | Time: %s | Avg: %s/s | Speed: %s/s",
+			humanize.Bytes(uint64(p.TotalSize)),
+			humanize.Bytes(uint64(d.ChunkSize)),
+			d.Concurrency,
+			humanize.Bytes(uint64(p.Size)),
+			p.TotalCost().Round(time.Second),
+			humanize.Bytes(p.AvgSpeed()),
+			humanize.Bytes(p.Speed()),
+		)
+	}
+
 	if err := d.Start(); err != nil {
 		log.Fatal(err)
 	}
 
-	fmt.Println("| Done!")
+	fmt.Println(" | Done!")
 }

got.go

@@ -34,18 +34,12 @@ type (
 		// Download file info.
 		*Info
 
-		// Progress func.
-		ProgressFunc
-
 		// URL to download.
 		URL string
 
 		// File destination.
 		Dest string
 
-		// Progress interval in ms.
-		Interval int
-
 		// Split file into chunks by ChunkSize in bytes.
 		ChunkSize int64
 
@@ -58,6 +52,12 @@ type (
 		// Max chunks to download at same time.
 		Concurrency int
 
+		// Progress...
+		Progress *Progress
+
+		// Progress interval in ms.
+		Interval int
+
 		// Download file chunks.
 		chunks []*Chunk
 
@@ -66,9 +66,6 @@ type (
 
 		// Is the URL redirected to a different location.
 		redirected bool
-
-		// Progress...
-		progress *Progress
 	}
 )
 
@@ -95,9 +92,6 @@ func (d *Download) Init() error {
 		},
 	}
 
-	// Init progress.
-	d.progress = &Progress{}
-
 	// Set default interval.
 	if d.Interval == 0 {
 		d.Interval = 20
@@ -108,6 +102,16 @@ func (d *Download) Init() error {
 		return err
 	}
 
+	// Set default progress.
+	if d.Progress == nil {
+
+		d.Progress = &Progress{
+			startedAt: time.Now(),
+			Interval:  d.Interval,
+			TotalSize: d.Info.Length,
+		}
+	}
+
 	// Partial content not supported 😢!
 	if d.Info.Rangeable == false || d.Info.Length == 0 {
 		return nil
@@ -177,7 +181,7 @@ func (d *Download) Init() error {
 		d.chunks = append(d.chunks, &Chunk{
 			Start:    startRange,
 			End:      endRange,
-			Progress: d.progress,
+			Progress: d.Progress,
 			Done:     make(chan struct{}),
 		})
 	}
@@ -187,6 +191,7 @@ func (d *Download) Init() error {
 
 // Start downloading.
 func (d *Download) Start() (err error) {
+
 	// Create a new temp dir for this download.
 	temp, err := ioutil.TempDir("", "GotChunks")
 	if err != nil {
@@ -196,8 +201,9 @@ func (d *Download) Start() (err error) {
 
 	ctx, cancel := context.WithCancel(context.TODO())
 	defer cancel()
+
 	// Run progress func.
-	go d.progress.Run(ctx, d)
+	go d.Progress.Run(ctx, d)
 
 	// Partial content not supported,
 	// just download the file in one chunk.
@@ -212,7 +218,7 @@ func (d *Download) Start() (err error) {
 		defer file.Close()
 
 		chunk := &Chunk{
-			Progress: d.progress,
+			Progress: d.Progress,
 		}
 
 		return chunk.Download(d.URL, d.client, file)
@@ -235,9 +241,11 @@ func (d *Download) Start() (err error) {
 		return err
 	}
 
-	if d.ProgressFunc != nil {
-		d.ProgressFunc(d.progress.Size, d.Info.Length, d)
+	// Update progress output after chunks finished.
+	if d.Progress.ProgressFunc != nil {
+		d.Progress.ProgressFunc(d.Progress, d)
 	}
+
 	return nil
 }
 
@@ -283,6 +291,7 @@ func (d *Download) merge(ctx context.Context) error {
 	defer file.Close()
 
 	for i := range d.chunks {
+
 		select {
 		case <-d.chunks[i].Done:
 		case <-ctx.Done():
@@ -309,17 +318,24 @@ func (d *Download) merge(ctx context.Context) error {
 
 // Download chunks
 func (d *Download) dl(ctx context.Context, temp string) error {
+
 	eg, ctx := errgroup.WithContext(ctx)
+
 	// Concurrency limit.
 	max := make(chan int, d.Concurrency)
 
 	for i := 0; i < len(d.chunks); i++ {
+
 		max <- 1
 		i := i
+
 		eg.Go(func() error {
+
 			defer func() {
 				<-max
 			}()
+
+			// Create chunk in temp dir.
 			chunk, err := os.Create(filepath.Join(temp, fmt.Sprintf("chunk-%d", i)))
 
 			if err != nil {

got_test.go

@@ -199,9 +199,9 @@ func downloadTest(t *testing.T, url string, size int64) {
 	defer clean(tmpFile)
 
 	d := &got.Download{
-		URL:          url,
-		Dest:         tmpFile,
-		Concurrency:  2,
+		URL:         url,
+		Dest:        tmpFile,
+		Concurrency: 2,
 	}
 
 	if err := d.Init(); err != nil {

progress.go

@@ -10,30 +10,62 @@ type (
 
 	// Progress can be used to show download progress to the user.
 	Progress struct {
-		Size int64
+		ProgressFunc
+
+		Size, TotalSize int64
+		Interval        int
+
+		lastSize  int64
+		startedAt time.Time
 	}
 
 	// ProgressFunc to show progress state, called based on Download interval.
-	ProgressFunc func(size int64, total int64, d *Download)
+	ProgressFunc func(p *Progress, d *Download)
 )
 
 // Run runs ProgressFunc based on interval if ProgressFunc set.
 func (p *Progress) Run(ctx context.Context, d *Download) {
-	if d.ProgressFunc != nil {
+
+	if p.ProgressFunc != nil {
+
 		for {
-			select {
-			case <-ctx.Done():
+			// Context cancelled
+			if ctx.Err() != nil {
 				return
-			default:
 			}
 
-			d.ProgressFunc(atomic.LoadInt64(&p.Size), d.Info.Length, d)
+			// Run progress func.
+			p.ProgressFunc(p, d)
+
+			// Update last size
+			atomic.StoreInt64(&p.lastSize, atomic.LoadInt64(&p.Size))
 
 			time.Sleep(time.Duration(d.Interval) * time.Millisecond)
 		}
 	}
 }
 
+// Speed returns download speed.
+func (p *Progress) Speed() uint64 {
+	return uint64((atomic.LoadInt64(&p.Size) - atomic.LoadInt64(&p.lastSize)) / int64(p.Interval) * 1000)
+}
+
+// AvgSpeed returns average download speed.
+func (p *Progress) AvgSpeed() uint64 {
+
+	if totalMills := p.TotalCost().Milliseconds(); totalMills > 0 {
+		return uint64(atomic.LoadInt64(&p.Size) / totalMills * 1000)
+	}
+
+	return 0
+}
+
+// TotalCost returns download duration.
+func (p *Progress) TotalCost() time.Duration {
+	return time.Now().Sub(p.startedAt)
+}
+
+// Write updates progress size.
 func (p *Progress) Write(b []byte) (int, error) {
 	n := len(b)
 	atomic.AddInt64(&p.Size, int64(n))