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>

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