github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/local/progress.go (about)

     1  package local
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  	"time"
     7  
     8  	"github.com/jedib0t/go-pretty/v6/progress"
     9  )
    10  
    11  const (
    12  	progressTrackerLength          = 20
    13  	progressTrackerWidth           = 40
    14  	progressTrackerDone            = 100
    15  	progressTrackerUpdateFrequency = 50 * time.Millisecond
    16  )
    17  
    18  type ProgressUpdaterReader struct {
    19  	r io.Reader
    20  	t *progress.Tracker
    21  }
    22  
    23  func (pu *ProgressUpdaterReader) Read(p []byte) (n int, err error) {
    24  	n, err = pu.r.Read(p)
    25  	pu.t.Increment(int64(n))
    26  	if err == io.EOF {
    27  		pu.t.MarkAsDone()
    28  	} else if err != nil {
    29  		pu.t.IncrementWithError(int64(n))
    30  	}
    31  	return
    32  }
    33  
    34  type ProgressUpdater struct {
    35  	t *progress.Tracker
    36  }
    37  
    38  func (p *ProgressUpdater) Reader(reader io.Reader) io.Reader {
    39  	return &ProgressUpdaterReader{
    40  		r: reader,
    41  		t: p.t,
    42  	}
    43  }
    44  
    45  func (p *ProgressUpdater) Done() {
    46  	p.t.MarkAsDone()
    47  }
    48  
    49  func (p *ProgressUpdater) Error() {
    50  	p.t.MarkAsErrored()
    51  }
    52  
    53  type ProgressSpinner struct {
    54  	t *progress.Tracker
    55  }
    56  
    57  func (p *ProgressSpinner) Error() {
    58  	p.t.MarkAsErrored()
    59  }
    60  
    61  func (p *ProgressSpinner) Done() {
    62  	p.t.MarkAsDone()
    63  }
    64  
    65  type ProgressPool struct {
    66  	pw   progress.Writer
    67  	done chan bool
    68  }
    69  
    70  func (p *ProgressPool) Start() {
    71  	const tickerDuration = 25 * time.Millisecond
    72  	go p.pw.Render()
    73  	go func() {
    74  		t := time.NewTicker(tickerDuration)
    75  		for {
    76  			select {
    77  			case <-p.done:
    78  				t.Stop()
    79  				return
    80  			case <-t.C:
    81  			}
    82  		}
    83  	}()
    84  }
    85  func (p *ProgressPool) Stop() {
    86  	p.pw.Stop()
    87  	p.done <- true
    88  	// according to examples, give enough time for the last render loop to complete.
    89  	for p.pw.IsRenderInProgress() {
    90  		const renderSleep = 5 * time.Millisecond
    91  		time.Sleep(renderSleep)
    92  	}
    93  }
    94  
    95  func (p *ProgressPool) AddReader(name string, sizeBytes int64) *ProgressUpdater {
    96  	tracker := &progress.Tracker{
    97  		Message: name,
    98  		Total:   sizeBytes,
    99  		Units:   progress.UnitsBytes,
   100  	}
   101  	p.pw.AppendTracker(tracker)
   102  	return &ProgressUpdater{t: tracker}
   103  }
   104  
   105  func (p *ProgressPool) AddSpinner(name string) *ProgressSpinner {
   106  	tracker := &progress.Tracker{
   107  		Message: name,
   108  		Total:   progressTrackerDone,
   109  		Units: progress.Units{
   110  			Notation:  "%",
   111  			Formatter: progress.FormatNumber,
   112  		},
   113  	}
   114  	p.pw.AppendTracker(tracker)
   115  	return &ProgressSpinner{t: tracker}
   116  }
   117  
   118  func NewProgressPool() *ProgressPool {
   119  	pw := progress.NewWriter()
   120  	pw.SetAutoStop(false) // important
   121  	pw.SetTrackerLength(progressTrackerLength)
   122  	pw.SetMessageWidth(progressTrackerWidth)
   123  	pw.SetSortBy(progress.SortByValue)
   124  	pw.SetStyle(progress.StyleDefault)
   125  	pw.SetTrackerPosition(progress.PositionRight)
   126  	pw.SetUpdateFrequency(progressTrackerUpdateFrequency)
   127  	pw.Style().Colors = progress.StyleColorsExample
   128  	pw.Style().Options.PercentFormat = "%4.1f%%"
   129  
   130  	return &ProgressPool{
   131  		pw:   pw,
   132  		done: make(chan bool),
   133  	}
   134  }
   135  
   136  type fileWrapper struct {
   137  	file   *os.File
   138  	reader io.Reader
   139  }
   140  
   141  func (f fileWrapper) Read(p []byte) (n int, err error) {
   142  	return f.reader.Read(p)
   143  }
   144  
   145  func (f fileWrapper) Seek(offset int64, whence int) (int64, error) {
   146  	return f.file.Seek(offset, whence)
   147  }