Skip to content

Commit 4a7097d

Browse files
ilopezlunadoringeman
authored andcommitted
Add total and pulled bytes to progress messages (docker#69)
* Adds Total and Pulled fields to ProgressMessage struct to track the total bytes to transfer and bytes transferred so far * Use uint64 type to represent pulled and total
1 parent a1cf455 commit 4a7097d

File tree

2 files changed

+24
-4
lines changed

2 files changed

+24
-4
lines changed

pkg/distribution/distribution/progress.go

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ import (
1313
type ProgressMessage struct {
1414
Type string `json:"type"` // "progress", "success", or "error"
1515
Message string `json:"message"` // Human-readable message
16+
Total uint64 `json:"total"` // Total bytes to transfer
17+
Pulled uint64 `json:"pulled"` // Bytes transferred so far
1618
}
1719

1820
type reporter struct {
@@ -42,6 +44,14 @@ func newProgressReporter(w io.Writer, msgF progressF) *reporter {
4244
}
4345
}
4446

47+
// safeUint64 converts an int64 to uint64, ensuring the value is non-negative
48+
func safeUint64(n int64) uint64 {
49+
if n < 0 {
50+
return 0
51+
}
52+
return uint64(n)
53+
}
54+
4555
// updates returns a channel for receiving progress updates. It is the responsibility of the caller to close
4656
// the channel when they are done sending updates. Should only be called once per reporter instance.
4757
func (r *reporter) updates() chan<- v1.Update {
@@ -60,7 +70,7 @@ func (r *reporter) updates() chan<- v1.Update {
6070
// Only update if enough time has passed or enough bytes downloaded or finished
6171
if now.Sub(lastUpdate) >= updateInterval ||
6272
bytesDownloaded >= minBytesForUpdate {
63-
if err := writeProgress(r.out, r.format(p)); err != nil {
73+
if err := writeProgress(r.out, r.format(p), safeUint64(p.Total), safeUint64(p.Complete)); err != nil {
6474
r.err = err
6575
}
6676
lastUpdate = now
@@ -92,10 +102,12 @@ func writeProgressMessage(w io.Writer, msg ProgressMessage) error {
92102
}
93103

94104
// writeProgress writes a progress update message
95-
func writeProgress(w io.Writer, msg string) error {
105+
func writeProgress(w io.Writer, msg string, total, pulled uint64) error {
96106
return writeProgressMessage(w, ProgressMessage{
97107
Type: "progress",
98108
Message: msg,
109+
Total: total,
110+
Pulled: pulled,
99111
})
100112
}
101113

pkg/distribution/distribution/progress_test.go

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@ import (
1111
func TestProgressMessages(t *testing.T) {
1212
t.Run("writeProgress", func(t *testing.T) {
1313
var buf bytes.Buffer
14-
err := writeProgress(&buf, pullMsg(v1.Update{
14+
update := v1.Update{
15+
Total: 2 * 1024 * 1024,
1516
Complete: 1024 * 1024,
16-
}))
17+
}
18+
err := writeProgress(&buf, pullMsg(update), uint64(update.Total), uint64(update.Complete))
1719
if err != nil {
1820
t.Fatalf("Failed to write progress message: %v", err)
1921
}
@@ -29,6 +31,12 @@ func TestProgressMessages(t *testing.T) {
2931
if msg.Message != "Downloaded: 1.00 MB" {
3032
t.Errorf("Expected message 'Downloaded: 1.00 MB', got '%s'", msg.Message)
3133
}
34+
if msg.Total != uint64(2*1024*1024) {
35+
t.Errorf("Expected total 2MB, got %d", msg.Total)
36+
}
37+
if msg.Pulled != uint64(1024*1024) {
38+
t.Errorf("Expected pulled 1MB, got %d", msg.Pulled)
39+
}
3240
})
3341

3442
t.Run("writeSuccess", func(t *testing.T) {

0 commit comments

Comments
 (0)