github.com/mvdan/u-root-coreutils@v0.0.0-20230122170626-c2eef2898555/pkg/progress/progress.go (about)

     1  // Copyright 2021 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // The progress package starts printing the status of a progress every second
     6  // It takes a pointer to an int64 which holds the amount of
     7  // bytes copied - and prints it.
     8  
     9  package progress
    10  
    11  import (
    12  	"fmt"
    13  	"io"
    14  	"os"
    15  	"sync"
    16  	"sync/atomic"
    17  	"time"
    18  )
    19  
    20  type progressData struct {
    21  	mode         string // one of: none, xfer, progress
    22  	start        time.Time
    23  	end          time.Time
    24  	endTimeMutex sync.Mutex
    25  	variable     *int64 // must be aligned for atomic operations
    26  	quit         chan struct{}
    27  }
    28  
    29  // Begin - start a progress routine
    30  //
    31  // mode describes in which mode it runs, none, progress or xfer
    32  // variable holds the amount of bytes written
    33  func Begin(mode string, variable *int64) (ProgressData *progressData) {
    34  	p := &progressData{
    35  		mode:         mode,
    36  		start:        time.Now(),
    37  		endTimeMutex: sync.Mutex{},
    38  		variable:     variable,
    39  	}
    40  	if p.mode == "progress" {
    41  		p.print(os.Stderr)
    42  		// Print progress in a separate goroutine.
    43  		p.quit = make(chan struct{}, 1)
    44  		go func() {
    45  			ticker := time.NewTicker(1 * time.Second)
    46  			defer ticker.Stop()
    47  			for {
    48  				select {
    49  				case <-ticker.C:
    50  					p.print(os.Stderr)
    51  				case <-p.quit:
    52  					p.endTimeMutex.Lock()
    53  					p.end = time.Now()
    54  					p.endTimeMutex.Unlock()
    55  					return
    56  				}
    57  			}
    58  		}()
    59  	}
    60  	return p
    61  }
    62  
    63  // End - Ends the progress and send quit signal to the channel
    64  func (p *progressData) End() {
    65  	if p.mode == "progress" {
    66  		// Properly synchronize goroutine.
    67  		p.quit <- struct{}{}
    68  		p.quit <- struct{}{}
    69  	} else {
    70  		p.endTimeMutex.Lock()
    71  		p.end = time.Now()
    72  		p.endTimeMutex.Unlock()
    73  	}
    74  	if p.mode == "progress" || p.mode == "xfer" {
    75  		// Print grand total.
    76  		p.print(os.Stderr)
    77  		fmt.Fprint(os.Stderr, "\n")
    78  	}
    79  }
    80  
    81  // With "status=progress", this is called from 3 places:
    82  // - Once at the beginning to appear responsive
    83  // - Every 1s afterwards
    84  // - Once at the end so the final value is accurate
    85  func (p *progressData) print(out io.Writer) {
    86  	elapse := time.Since(p.start)
    87  	n := atomic.LoadInt64(p.variable)
    88  	d := float64(n)
    89  	const mib = 1024 * 1024
    90  	const mb = 1000 * 1000
    91  	// The ANSI escape may be undesirable to some eyes.
    92  	if p.mode == "progress" {
    93  		out.Write([]byte("\033[2K\r"))
    94  	}
    95  	fmt.Fprintf(out, "%d bytes (%.3f MB, %.3f MiB) copied, %.3f s, %.3f MB/s",
    96  		n, d/mb, d/mib, elapse.Seconds(), float64(d)/elapse.Seconds()/mb)
    97  }