webnet/tailscale
1
0
Fork 0
Code Issues 1 Pull Requests 1 Packages Projects Releases Wiki Activity

util/progresstracking: add Ticker, NewWriter, and CountingWriter

Browse Source 
Add three new helpers to the existing progresstracking package:

  - Ticker: spawns a 1 Hz goroutine that calls a report function with
    the current value of an atomic counter and a total. Returns a stop
    function (safe to call multiple times via sync.OnceFunc) that fires
    one final report and blocks until the goroutine exits.

  - NewWriter: wraps an io.Writer and calls onProgress at most once per
    interval with the cumulative byte count.

  - CountingWriter: an io.Writer that atomically counts bytes written,
    for use with Ticker.

These will be used by the appliance flash and OTA update code in
subsequent commits.

Updates #1866

Change-Id: If353cea6506f5351b6fb19bfdb7bc9b78fe7855e
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
This commit is contained in:
Brad Fitzpatrick
2026-07-01 10:02:05 -07:00
committed by Brad Fitzpatrick
parent d0fcb668d5
commit a8f3c861a4
7 changed files with 122 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clientupdate/clientupdate_gokrazy.go
+15 -3
View File
@@ -15,9 +15,11 @@ import (
"net/http"
"os"
"strings"
"time"
"tailscale.com/clientupdate/distsign"
"tailscale.com/types/logger"
"tailscale.com/util/progresstracking"
)
const (
@@ -50,8 +52,9 @@ func gokrazyUpdateFromURL(ctx context.Context, args GokrazyUpdateArgs) error {
tmp.Close()
defer os.Remove(tmpName)
logf("downloading %s", args.URL)
if args.AllowUnsigned {
if err := downloadUnverified(ctx, args.URL, tmpName); err != nil {
if err := downloadUnverified(ctx, logf, args.URL, tmpName); err != nil {
return err
}
} else {
@@ -66,6 +69,8 @@ func gokrazyUpdateFromURL(ctx context.Context, args GokrazyUpdateArgs) error {
}
defer zr.Close()
logf("download complete")
gokClient := gokrazyHTTPClient()
for _, part := range []struct {
name string
@@ -75,6 +80,7 @@ func gokrazyUpdateFromURL(ctx context.Context, args GokrazyUpdateArgs) error {
{"boot.img", "/update/boot"},
{"mbr.img", "/update/mbr"},
} {
logf("writing %s...", part.name)
if err := putGokrazyGAFMember(ctx, gokClient, zr.File, part.name, part.path); err != nil {
return err
}
@@ -94,7 +100,7 @@ func gokrazyUpdateFromURL(ctx context.Context, args GokrazyUpdateArgs) error {
// downloadUnverified saves the GAF at srcURL to dstPath without verifying
// a signature. It is used only when args.AllowUnsigned is set, for tests
// that serve the GAF from a fileserver that does not publish distsign.pub.
func downloadUnverified(ctx context.Context, srcURL, dstPath string) error {
func downloadUnverified(ctx context.Context, logf logger.Logf, srcURL, dstPath string) error {
req, err := http.NewRequestWithContext(ctx, "GET", srcURL, nil)
if err != nil {
return err
@@ -111,7 +117,13 @@ func downloadUnverified(ctx context.Context, srcURL, dstPath string) error {
if err != nil {
return err
}
if _, err := io.Copy(f, res.Body); err != nil {
total := res.ContentLength
pw := progresstracking.NewWriter(io.Discard, total, time.Second, func(done int64) {
if total > 0 {
logf("downloading: %d / %d MB (%.0f%%)", done>>20, total>>20, float64(done)/float64(total)*100)
}
})
if _, err := io.Copy(f, io.TeeReader(res.Body, pw)); err != nil {
f.Close()
return err
}
clientupdate/distsign/distsign.go
+5 -22
View File
@@ -60,6 +60,7 @@ import (
"tailscale.com/types/logger"
"tailscale.com/util/httpm"
"tailscale.com/util/must"
"tailscale.com/util/progresstracking"
)
const (
@@ -373,7 +374,10 @@ func (c *Client) download(ctx context.Context, url, dst string, limit int64) ([]
return nil, 0, err
}
defer of.Close()
pw := &progressWriter{total: res.ContentLength, logf: c.logf}
total := res.ContentLength
pw := progresstracking.NewWriter(io.Discard, total, 2*time.Second, func(done int64) {
c.logf("Downloaded %v/%v (%.1f%%)", done, total, float64(done)/float64(total)*100)
})
h := NewPackageHash()
n, err := io.Copy(io.MultiWriter(of, h, pw), io.LimitReader(dlRes.Body, limit))
if err != nil {
@@ -388,31 +392,10 @@ func (c *Client) download(ctx context.Context, url, dst string, limit int64) ([]
if err := of.Close(); err != nil {
return nil, n, err
}
pw.print()
return h.Sum(nil), h.Len(), nil
}
type progressWriter struct {
done int64
total int64
lastPrint time.Time
logf logger.Logf
}
func (pw *progressWriter) Write(p []byte) (n int, err error) {
pw.done += int64(len(p))
if time.Since(pw.lastPrint) > 2*time.Second {
pw.print()
}
return len(p), nil
}
func (pw *progressWriter) print() {
pw.lastPrint = time.Now()
pw.logf("Downloaded %v/%v (%.1f%%)", pw.done, pw.total, float64(pw.done)/float64(pw.total)*100)
}
func parsePrivateKey(data []byte, typeTag string) (ed25519.PrivateKey, error) {
b, rest := pem.Decode(data)
if b == nil {
cmd/tailscale/cli/configure-flash-appliance.go
+9 -53
View File
@@ -18,14 +18,13 @@ import (
"slices"
"sort"
"strings"
"sync/atomic"
"time"
"github.com/bradfitz/monogok/disklayout"
"github.com/peterbourgon/ff/v3/ffcli"
"tailscale.com/clientupdate"
"tailscale.com/clientupdate/distsign"
"tailscale.com/gokrazy/mkfs"
"tailscale.com/util/progresstracking"
"tailscale.com/util/prompt"
)
@@ -442,62 +441,19 @@ func writeZipMemberAt(f *os.File, zf *zip.File, offset int64) error {
return err
}
total := int64(zf.UncompressedSize64)
cw := &countingWriter{w: f}
stop := startProgress(zf.Name, total, &cw.count)
cw := &progresstracking.CountingWriter{W: f}
stop := progresstracking.Ticker(cw.Count, total, func(d, t int64) {
pct := 0.0
if t > 0 {
pct = float64(d) * 100 / float64(t)
}
fmt.Fprintf(Stderr, " %s: %s / %s (%.1f%%)\n", zf.Name, humanBytes(d), humanBytes(t), pct)
})
defer stop()
_, err = io.Copy(cw, rc)
return err
}
// countingWriter wraps an io.Writer and tracks total bytes written so
// the progress goroutine can report it.
type countingWriter struct {
w io.Writer
count atomic.Int64
}
func (c *countingWriter) Write(b []byte) (int, error) {
n, err := c.w.Write(b)
if n > 0 {
c.count.Add(int64(n))
}
return n, err
}
// startProgress spawns a 1 Hz goroutine that prints "<name>: <done> / <total>"
// to Stderr until the returned stop function is called. The final tick on
// stop reports the final state, so the caller doesn't need to repeat it.
func startProgress(name string, total int64, done *atomic.Int64) func() {
stop := make(chan struct{})
finished := make(chan struct{})
go func() {
defer close(finished)
t := time.NewTicker(time.Second)
defer t.Stop()
report := func() {
d := done.Load()
pct := 0.0
if total > 0 {
pct = float64(d) * 100 / float64(total)
}
fmt.Fprintf(Stderr, " %s: %s / %s (%.1f%%)\n", name, humanBytes(d), humanBytes(total), pct)
}
for {
select {
case <-stop:
report()
return
case <-t.C:
report()
}
}
}()
return func() {
close(stop)
<-finished
}
}
// humanBytes returns a friendly approximation of n bytes, e.g. "62.5 GB".
func humanBytes(n int64) string {
const (
cmd/tailscale/depaware.txt
+1
View File
@@ -292,6 +292,7 @@ tailscale.com/cmd/tailscale dependencies: (generated by github.com/tailscale/dep
tailscale.com/util/mak from tailscale.com/cmd/tailscale/cli+
tailscale.com/util/must from tailscale.com/clientupdate/distsign+
tailscale.com/util/nocasemaps from tailscale.com/types/ipproto
tailscale.com/util/progresstracking from tailscale.com/clientupdate/distsign+
tailscale.com/util/prompt from tailscale.com/cmd/tailscale/cli
💣 tailscale.com/util/qrcodes from tailscale.com/cmd/tailscale/cli
tailscale.com/util/quarantine from tailscale.com/cmd/tailscale/cli
cmd/tailscaled/depaware.txt
+1 -1
View File
@@ -469,7 +469,7 @@ tailscale.com/cmd/tailscaled dependencies: (generated by github.com/tailscale/de
W 💣 tailscale.com/util/osdiag/internal/wsc from tailscale.com/util/osdiag
tailscale.com/util/osshare from tailscale.com/cmd/tailscaled+
tailscale.com/util/osuser from tailscale.com/ipn/ipnlocal+
tailscale.com/util/progresstracking from tailscale.com/feature/taildrop
tailscale.com/util/progresstracking from tailscale.com/feature/taildrop+
tailscale.com/util/race from tailscale.com/net/dns/resolver
tailscale.com/util/racebuild from tailscale.com/logpolicy
tailscale.com/util/rands from tailscale.com/ipn/ipnlocal+
gokrazy/mkfs/mkfs.go
+8 -31
View File
@@ -24,6 +24,7 @@ import (
"github.com/bradfitz/monogok/disklayout"
"github.com/diskfs/go-diskfs/backend"
"github.com/diskfs/go-diskfs/filesystem/ext4"
"tailscale.com/util/progresstracking"
)
// gptSecondaryReservedSectors is the number of 512-byte sectors that
@@ -288,39 +289,15 @@ func (m *memBackend) flushTo(f io.WriterAt, baseOffset int64) error {
return nil
}
// startExt4FlushProgress prints "ext4 perm: N MB / M MB (X%)" to
// os.Stderr roughly once a second, plus a final tick on stop. Returns
// a function the caller must invoke when the flush is done.
func startExt4FlushProgress(done *atomic.Int64, total int64) func() {
stopCh := make(chan struct{})
finished := make(chan struct{})
go func() {
defer close(finished)
t := time.NewTicker(time.Second)
defer t.Stop()
report := func() {
d := done.Load()
pct := 0.0
if total > 0 {
pct = float64(d) * 100 / float64(total)
}
fmt.Fprintf(os.Stderr, " ext4 perm: %s / %s (%.1f%%)\n",
humanBytes(d), humanBytes(total), pct)
return progresstracking.Ticker(done.Load, total, func(d, t int64) {
pct := 0.0
if t > 0 {
pct = float64(d) * 100 / float64(t)
}
for {
select {
case <-stopCh:
report()
return
case <-t.C:
report()
}
}
}()
return func() {
close(stopCh)
<-finished
}
fmt.Fprintf(os.Stderr, " ext4 perm: %s / %s (%.1f%%)\n",
humanBytes(d), humanBytes(t), pct)
})
}
func humanBytes(n int64) string {
util/progresstracking/progresstracking.go
+83 -1
View File
@@ -2,11 +2,14 @@
// SPDX-License-Identifier: BSD-3-Clause
// Package progresstracking provides wrappers around io.Reader and io.Writer
// that track progress.
// that track progress, and a Ticker for reporting progress from an atomic
// counter on a regular interval.
package progresstracking
import (
"io"
"sync"
"sync/atomic"
"time"
)
@@ -37,3 +40,82 @@ func (r *reader) Read(p []byte) (int, error) {
}
return n, err
}
// NewWriter wraps w with a writer that calls onProgress after every write
// that brings the total past the next interval threshold. onProgress receives
// the cumulative byte count. If expectedTotal > 0, a final onProgress call is
// guaranteed when the cumulative count reaches or exceeds it, even if the
// interval hasn't elapsed.
func NewWriter(w io.Writer, expectedTotal int64, interval time.Duration, onProgress func(totalWritten int64)) io.Writer {
return &writer{w: w, expectedTotal: expectedTotal, interval: interval, onProgress: onProgress}
}
type writer struct {
w io.Writer
expectedTotal int64 // non-zero if known
interval time.Duration
onProgress func(int64)
lastTracked time.Time
total int64
reachedTotal bool
}
func (pw *writer) Write(p []byte) (int, error) {
n, err := pw.w.Write(p)
pw.total += int64(n)
if !pw.reachedTotal && pw.expectedTotal > 0 && pw.total >= pw.expectedTotal {
pw.onProgress(pw.total)
pw.reachedTotal = true
} else if time.Since(pw.lastTracked) > pw.interval {
pw.onProgress(pw.total)
pw.lastTracked = time.Now()
}
return n, err
}
// Ticker reports progress on a regular interval by polling a counter function.
// It spawns a background goroutine that calls report approximately every
// second. Call the returned stop function when the operation is complete;
// stop calls report one final time and blocks until the goroutine exits.
// The stop function is safe to call multiple times, but will only call
// report the first time it is invoked.
func Ticker(done func() int64, total int64, report func(done, total int64)) (stop func()) {
stopCh := make(chan struct{})
finished := make(chan struct{})
go func() {
defer close(finished)
t := time.NewTicker(time.Second)
defer t.Stop()
for {
select {
case <-stopCh:
report(done(), total)
return
case <-t.C:
report(done(), total)
}
}
}()
return sync.OnceFunc(func() {
close(stopCh)
<-finished
})
}
// CountingWriter wraps an io.Writer and atomically tracks total bytes
// written, suitable for use with Ticker.
type CountingWriter struct {
W io.Writer
count atomic.Int64
}
// Count returns the total number of bytes written.
func (c *CountingWriter) Count() int64 { return c.count.Load() }
func (c *CountingWriter) Write(b []byte) (int, error) {
n, err := c.W.Write(b)
if n > 0 {
c.count.Add(int64(n))
}
return n, err
}