aboutsummaryrefslogtreecommitdiffstats
path: root/drive/progress.go
blob: 9d4eb5a7abf6eb675fabcf6d8c0b5d070b445720 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
package drive

import (
    "io"
    "io/ioutil"
    "fmt"
    "time"
)

const MaxDrawInterval = time.Second * 1
const MaxRateInterval = time.Second * 3

func getProgressReader(r io.Reader, w io.Writer, size int64) io.Reader {
    // Don't wrap reader if output is discarded or size is too small
    if w == ioutil.Discard || size < 1024 * 1024 {
        return r
    }

    return &Progress{
        Reader: r,
        Writer: w,
        Size: size,
    }
}

type Progress struct {
    Writer io.Writer
    Reader io.Reader
    Size int64
    progress int64
    rate int64
    rateProgress int64
    rateUpdated time.Time
    updated time.Time
    done bool
}

func (self *Progress) Read(p []byte) (int, error) {
    // Read
    n, err := self.Reader.Read(p)

    now := time.Now()
    isLast := err != nil

    // Increment progress
    newProgress := self.progress + int64(n)
    self.progress = newProgress

    // Initialize rate state
    if self.rateUpdated.IsZero() {
        self.rateUpdated = now
        self.rateProgress = newProgress
    }

    // Update rate every x seconds
    if self.rateUpdated.Add(MaxRateInterval).Before(now) {
        self.rate = calcRate(newProgress - self.rateProgress, self.rateUpdated, now)
        self.rateUpdated = now
        self.rateProgress = newProgress
    }

    // Draw progress every x seconds
    if self.updated.Add(MaxDrawInterval).Before(now) || isLast {
        self.draw(isLast)
        self.updated = now
    }

    // Mark as done if error occurs
    self.done = isLast

    return n, err
}

func (self *Progress) draw(isLast bool) {
    if self.done {
        return
    }

    self.clear()

    // Print progress
    fmt.Fprintf(self.Writer, "%s/%s", formatSize(self.progress, false), formatSize(self.Size, false))

    // Print rate
    if self.rate > 0 {
        fmt.Fprintf(self.Writer, ", Rate: %s/s", formatSize(self.rate, false))
    }

    if isLast {
        self.clear()
    }
}

func (self *Progress) clear() {
    fmt.Fprintf(self.Writer, "\r%50s\r", "")
}