Skip to content

Commit 96285e3

Browse files
feat(feeds): upload progress bar
1 parent da565dd commit 96285e3

2 files changed

Lines changed: 147 additions & 3 deletions

File tree

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ require (
4444
github.com/stretchr/testify v1.9.0
4545
github.com/subosito/gotenv v1.2.0 // indirect
4646
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f // indirect
47-
golang.org/x/term v0.0.0-20220411215600-e5f449aeb171 // indirect
47+
golang.org/x/term v0.0.0-20220411215600-e5f449aeb171
4848
golang.org/x/text v0.3.8 // indirect
4949
gopkg.in/ini.v1 v1.66.4 // indirect
5050
gopkg.in/yaml.v2 v2.4.0 // indirect

pkg/utils/upload.go

Lines changed: 146 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,19 @@ package utils
22

33
import (
44
"context"
5+
"fmt"
6+
"io"
57
"net/http"
68
"os"
9+
"path/filepath"
10+
"sync/atomic"
11+
"time"
12+
13+
"golang.org/x/term"
714
)
815

916
// UploadToS3 uploads a local file to the given S3 presigned URL via HTTP PUT.
17+
// When stderr is a TTY, a progress bar is rendered so large uploads are not opaque.
1018
func UploadToS3(ctx context.Context, filename, url string) error {
1119
data, err := os.Open(filename)
1220
if err != nil {
@@ -19,18 +27,154 @@ func UploadToS3(ctx context.Context, filename, url string) error {
1927
return err
2028
}
2129

22-
req, err := http.NewRequestWithContext(ctx, "PUT", url, data)
30+
size := stat.Size()
31+
body := io.Reader(data)
32+
var pr *progressReader
33+
if term.IsTerminal(int(os.Stderr.Fd())) {
34+
pr = newProgressReader(data, size, filepath.Base(filename), os.Stderr)
35+
body = pr
36+
}
37+
38+
req, err := http.NewRequestWithContext(ctx, "PUT", url, body)
2339
if err != nil {
2440
return err
2541
}
2642
req.Header.Set("Content-Type", "application/json")
27-
req.ContentLength = stat.Size()
43+
req.ContentLength = size
2844

2945
resp, err := http.DefaultClient.Do(req)
3046
if err != nil {
47+
if pr != nil {
48+
pr.abort()
49+
}
3150
return err
3251
}
3352
defer resp.Body.Close()
3453

54+
if pr != nil {
55+
pr.finish()
56+
}
3557
return nil
3658
}
59+
60+
// progressReader wraps an io.Reader and renders a single-line progress bar
61+
// (carriage-return overwritten) to the given writer. Updates are throttled.
62+
type progressReader struct {
63+
r io.Reader
64+
total int64
65+
read int64
66+
name string
67+
out io.Writer
68+
start time.Time
69+
lastPrint time.Time
70+
barWidth int
71+
done int32
72+
}
73+
74+
func newProgressReader(r io.Reader, total int64, name string, out io.Writer) *progressReader {
75+
return &progressReader{
76+
r: r,
77+
total: total,
78+
name: name,
79+
out: out,
80+
start: time.Now(),
81+
barWidth: 30,
82+
}
83+
}
84+
85+
func (p *progressReader) Read(b []byte) (int, error) {
86+
n, err := p.r.Read(b)
87+
if n > 0 {
88+
atomic.AddInt64(&p.read, int64(n))
89+
p.maybeRender(false)
90+
}
91+
return n, err
92+
}
93+
94+
func (p *progressReader) maybeRender(force bool) {
95+
now := time.Now()
96+
if !force && now.Sub(p.lastPrint) < 100*time.Millisecond {
97+
return
98+
}
99+
p.lastPrint = now
100+
p.render()
101+
}
102+
103+
func (p *progressReader) render() {
104+
read := atomic.LoadInt64(&p.read)
105+
elapsed := time.Since(p.start).Seconds()
106+
var rate float64
107+
if elapsed > 0 {
108+
rate = float64(read) / elapsed
109+
}
110+
111+
var pct float64
112+
var bar string
113+
var eta string
114+
if p.total > 0 {
115+
pct = float64(read) / float64(p.total) * 100
116+
filled := min(int(float64(p.barWidth)*float64(read)/float64(p.total)), p.barWidth)
117+
b := make([]byte, p.barWidth)
118+
for i := 0; i < p.barWidth; i++ {
119+
if i < filled {
120+
b[i] = '='
121+
} else {
122+
b[i] = ' '
123+
}
124+
}
125+
bar = string(b)
126+
if rate > 0 && read < p.total {
127+
remaining := time.Duration(float64(p.total-read)/rate) * time.Second
128+
eta = " ETA " + formatDuration(remaining)
129+
}
130+
fmt.Fprintf(p.out, "\rUploading %s [%s] %5.1f%% %s/%s %s/s%s",
131+
p.name, bar, pct, humanBytes(read), humanBytes(p.total), humanBytes(int64(rate)), eta)
132+
} else {
133+
fmt.Fprintf(p.out, "\rUploading %s %s %s/s", p.name, humanBytes(read), humanBytes(int64(rate)))
134+
}
135+
}
136+
137+
func (p *progressReader) finish() {
138+
if !atomic.CompareAndSwapInt32(&p.done, 0, 1) {
139+
return
140+
}
141+
p.render()
142+
fmt.Fprintln(p.out)
143+
}
144+
145+
func (p *progressReader) abort() {
146+
if !atomic.CompareAndSwapInt32(&p.done, 0, 1) {
147+
return
148+
}
149+
fmt.Fprintln(p.out)
150+
}
151+
152+
func humanBytes(n int64) string {
153+
const unit = 1024
154+
if n < unit {
155+
return fmt.Sprintf("%d B", n)
156+
}
157+
div, exp := int64(unit), 0
158+
for n/div >= unit {
159+
div *= unit
160+
exp++
161+
}
162+
return fmt.Sprintf("%.1f %cB", float64(n)/float64(div), "KMGTPE"[exp])
163+
}
164+
165+
func formatDuration(d time.Duration) string {
166+
if d < time.Second {
167+
return "<1s"
168+
}
169+
d = d.Round(time.Second)
170+
h := int(d / time.Hour)
171+
m := int((d % time.Hour) / time.Minute)
172+
s := int((d % time.Minute) / time.Second)
173+
if h > 0 {
174+
return fmt.Sprintf("%dh%02dm%02ds", h, m, s)
175+
}
176+
if m > 0 {
177+
return fmt.Sprintf("%dm%02ds", m, s)
178+
}
179+
return fmt.Sprintf("%ds", s)
180+
}

0 commit comments

Comments
 (0)