Something went wrong. Try again.
forked from tailscale/tailscale
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122// Copyright (c) Tailscale Inc & contributors// SPDX-License-Identifier: BSD-3-Clause
// Package progresstracking provides wrappers around io.Reader and io.Writer// 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")
// NewReader wraps the given Reader with a progress tracking Reader that// reports progress at the following points://// - First read// - Every read spaced at least interval since the prior read// - Last readfunc NewReader(r io.Reader, interval time.Duration, onProgress func(totalRead int, err error)) io.Reader { return &reader{Reader: r, interval: interval, onProgress: onProgress}}
type reader struct { io.Reader interval time.Duration onProgress func(int, error) lastTracked time.Time totalRead int}
func (r *reader) Read(p []byte) (int, error) { n, err := r.Reader.Read(p) r.totalRead += n if time.Since(r.lastTracked) > r.interval || err != nil { r.onProgress(r.totalRead, err) r.lastTracked = time.Now() } 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}