print download progress and speed

Author: cenkalti

download.go

@@ -8,7 +8,11 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 	"time"
+
+	"github.com/cenkalti/log"
+	"github.com/paulbellamy/ratecounter"
 )
 
 type Download struct {
@@ -99,8 +103,11 @@ func (d *Download) Run(ctx context.Context) error {
 	trw := &TimerResetWriter{timer: time.AfterFunc(defaultTimeout, cancel)}
 	tr := io.TeeReader(rc, trw)
 
+	pr := NewProgressReader(tr, d.state.Offset, d.state.Size, d.String())
+	go pr.Run()
 	remaining := d.state.Size - d.state.Offset
-	n, copyErr := io.CopyN(wc, tr, remaining)
+	n, copyErr := io.CopyN(wc, pr, remaining)
+	pr.Close()
 
 	err = wc.Close()
 	if err != nil {
@@ -167,3 +174,48 @@ func (w *TimerResetWriter) Write(p []byte) (int, error) {
 	w.timer.Reset(defaultTimeout)
 	return len(p), nil
 }
+
+type ProgressReader struct {
+	r       io.Reader
+	offset  int64
+	size    int64
+	prefix  string
+	counter *ratecounter.RateCounter
+	closeC  chan struct{}
+}
+
+func NewProgressReader(r io.Reader, offset, size int64, prefix string) *ProgressReader {
+	return &ProgressReader{
+		r:       r,
+		offset:  offset,
+		size:    size,
+		prefix:  prefix,
+		counter: ratecounter.NewRateCounter(time.Second),
+		closeC:  make(chan struct{}),
+	}
+}
+
+func (r *ProgressReader) Read(p []byte) (int, error) {
+	n, err := r.r.Read(p)
+	r.counter.Incr(int64(n))
+	atomic.AddInt64(&r.offset, int64(n))
+	return n, err
+}
+
+func (r *ProgressReader) Run() {
+	for {
+		select {
+		case <-time.After(time.Second): // TODO use time.Ticker
+			offset := atomic.LoadInt64(&r.offset)
+			progress := (offset * 100) / r.size
+			speed := r.counter.Rate() / 1024
+			log.Infof("%s %d%% %dKB/s", r.prefix, progress, speed)
+		case <-r.closeC:
+			return
+		}
+	}
+}
+
+func (r *ProgressReader) Close() {
+	close(r.closeC)
+}

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/BurntSushi/toml v0.3.1
 	github.com/adrg/xdg v0.2.1
 	github.com/cenkalti/log v1.0.0
+	github.com/paulbellamy/ratecounter v0.2.0
 	github.com/putdotio/go-putio v1.3.2
 	github.com/stretchr/testify v1.6.1 // indirect
 	go.etcd.io/bbolt v1.3.5

go.sum

@@ -15,6 +15,8 @@ github.com/hashicorp/go-multierror v1.1.0 h1:B9UzwGQJehnUY1yNrnwREHc3fGbC2xefo8g
 github.com/hashicorp/go-multierror v1.1.0/go.mod h1:spPvp8C1qA32ftKqdAHm4hHTbPw+vmowP0z+KUhOZdA=
 github.com/mattn/go-isatty v0.0.12 h1:wuysRhFDzyxgEmMf5xjvJ2M9dZoWAXNNr5LSBS7uHXY=
 github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
+github.com/paulbellamy/ratecounter v0.2.0 h1:2L/RhJq+HA8gBQImDXtLPrDXK5qAj6ozWVK/zFXVJGs=
+github.com/paulbellamy/ratecounter v0.2.0/go.mod h1:Hfx1hDpSGoqxkVVpBi/IlYD7kChlfo5C6hzIHwPqfFE=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/putdotio/go-putio v1.3.2 h1:PZeVe9ckULim99J+zMDlQ6Lq+iv8/vJM4EMGrXfPAY4=

main.go

@@ -20,7 +20,6 @@ import (
 var Version = "0.0.0"
 
 // TODO Reconciliation tests
-// TODO Measure download progress
 // TODO Measure upload progress
 // TODO Daemon mode
 // TODO HTTP API