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 <bradfitz@tailscale.com>

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
 	}

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 {

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}
+	cw := &progresstracking.CountingWriter{W: f}
-	stop := startProgress(zf.Name, total, &cw.count)
+	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

@@ -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

@@ -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

@@ -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{})
+	return progresstracking.Ticker(done.Load, total, func(d, t int64) {
-	finished := make(chan struct{})
+		pct := 0.0
-	go func() {
+		if t > 0 {
-		defer close(finished)
+			pct = float64(d) * 100 / float64(t)
-		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)
 		}
-		for {
+		fmt.Fprintf(os.Stderr, "  ext4 perm: %s / %s (%.1f%%)\n",
-			select {
+			humanBytes(d), humanBytes(t), pct)
-			case <-stopCh:
+	})
-				report()
-				return
-			case <-t.C:
-				report()
-			}
-		}
-	}()
-	return func() {
-		close(stopCh)
-		<-finished
-	}
 }
 
 func humanBytes(n int64) string {

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
+}