github.com/jfrog/jfrog-cli-core/v2@v2.51.0/common/progressbar/readerprogressbar.go (about)

     1  package progressbar
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  type ReaderProgressBar struct {
     8  	*progressBarUnit
     9  	Id int
    10  }
    11  
    12  // Wraps an io.Reader for bytes reading tracking
    13  type proxyReader struct {
    14  	unit *progressBarUnit
    15  	io.ReadCloser
    16  }
    17  
    18  // Used to update the progress bar progress.
    19  func (p *ReaderProgressBar) ActionWithProgress(reader io.Reader) (results io.Reader) {
    20  	return p.readWithProgress(reader)
    21  }
    22  
    23  func (p *ReaderProgressBar) SetProgress(progress int64) {
    24  	p.bar.SetCurrent(progress)
    25  }
    26  
    27  // Abort aborts a progress indicator. Called on both successful and unsuccessful operations
    28  func (p *ReaderProgressBar) Abort() {
    29  	close(p.incrChannel)
    30  	p.bar.Abort(true)
    31  }
    32  
    33  // GetId Returns the ProgressBar ID
    34  //
    35  //nolint:gocritic
    36  func (p *ReaderProgressBar) GetId() (Id int) {
    37  	return p.Id
    38  }
    39  
    40  func (p *ReaderProgressBar) getProgressBarUnit() (unit *progressBarUnit) {
    41  	return p.progressBarUnit
    42  }
    43  
    44  // Wraps a body of a response (io.Reader) and increments bar accordingly
    45  func (p *ReaderProgressBar) readWithProgress(reader io.Reader) (wrappedReader io.Reader) {
    46  	wrappedReader = initProxyReader(p.progressBarUnit, reader)
    47  	return wrappedReader
    48  }
    49  
    50  func initProxyReader(unit *progressBarUnit, reader io.Reader) io.ReadCloser {
    51  	if reader == nil {
    52  		return nil
    53  	}
    54  	rc, ok := reader.(io.ReadCloser)
    55  	if !ok {
    56  		rc = io.NopCloser(reader)
    57  	}
    58  	return &proxyReader{unit, rc}
    59  }
    60  
    61  // Overrides the Read method of the original io.Reader.
    62  func (pr *proxyReader) Read(p []byte) (n int, err error) {
    63  	n, err = pr.ReadCloser.Read(p)
    64  	if n > 0 && (err == nil || err == io.EOF) {
    65  		pr.incrChannel(n)
    66  	}
    67  	return
    68  }
    69  
    70  func (pr *proxyReader) incrChannel(n int) {
    71  	// When an upload / download error occurs (for example, a bad HTTP error code),
    72  	// The progress bar's Abort method is invoked and closes the channel.
    73  	// Therefore, the channel may be already closed at this stage, which leads to a panic.
    74  	// We therefore need to recover if that happens.
    75  	defer func() {
    76  		_ = recover()
    77  	}()
    78  	pr.unit.incrChannel <- n
    79  }