From a8f3c861a4c7420b23de9d6e9d517aa8e862620f Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Wed, 1 Jul 2026 15:00:47 +0000 Subject: [PATCH] util/progresstracking: add Ticker, NewWriter, and CountingWriter 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 --- clientupdate/clientupdate_gokrazy.go | 18 +++- clientupdate/distsign/distsign.go | 27 ++---- .../cli/configure-flash-appliance.go | 62 ++------------ cmd/tailscale/depaware.txt | 1 + cmd/tailscaled/depaware.txt | 2 +- gokrazy/mkfs/mkfs.go | 39 ++------- util/progresstracking/progresstracking.go | 84 ++++++++++++++++++- 7 files changed, 122 insertions(+), 111 deletions(-) diff --git a/clientupdate/clientupdate_gokrazy.go b/clientupdate/clientupdate_gokrazy.go index bad473fac..a78bdf89b 100644 --- a/clientupdate/clientupdate_gokrazy.go +++ b/clientupdate/clientupdate_gokrazy.go @@ -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 } diff --git a/clientupdate/distsign/distsign.go b/clientupdate/distsign/distsign.go index c13ed8497..99ab193f8 100644 --- a/clientupdate/distsign/distsign.go +++ b/clientupdate/distsign/distsign.go @@ -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 { diff --git a/cmd/tailscale/cli/configure-flash-appliance.go b/cmd/tailscale/cli/configure-flash-appliance.go index 113014242..8cf5d8d34 100644 --- a/cmd/tailscale/cli/configure-flash-appliance.go +++ b/cmd/tailscale/cli/configure-flash-appliance.go @@ -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 ": / " -// 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 ( diff --git a/cmd/tailscale/depaware.txt b/cmd/tailscale/depaware.txt index c7a063962..91cbd70b7 100644 --- a/cmd/tailscale/depaware.txt +++ b/cmd/tailscale/depaware.txt @@ -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 diff --git a/cmd/tailscaled/depaware.txt b/cmd/tailscaled/depaware.txt index cb497da84..92e3406b6 100644 --- a/cmd/tailscaled/depaware.txt +++ b/cmd/tailscaled/depaware.txt @@ -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+ diff --git a/gokrazy/mkfs/mkfs.go b/gokrazy/mkfs/mkfs.go index 0e9dedaed..eb9263e19 100644 --- a/gokrazy/mkfs/mkfs.go +++ b/gokrazy/mkfs/mkfs.go @@ -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 { diff --git a/util/progresstracking/progresstracking.go b/util/progresstracking/progresstracking.go index 21cbfa52b..5aad6d9f8 100644 --- a/util/progresstracking/progresstracking.go +++ b/util/progresstracking/progresstracking.go @@ -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 +}