github.com/demonoid81/moby@v0.0.0-20200517203328-62dd8e17c460/pkg/progress/progressreader.go (about)

     1  package progress // import "github.com/demonoid81/moby/pkg/progress"
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  
     7  	"golang.org/x/time/rate"
     8  )
     9  
    10  // Reader is a Reader with progress bar.
    11  type Reader struct {
    12  	in          io.ReadCloser // Stream to read from
    13  	out         Output        // Where to send progress bar to
    14  	size        int64
    15  	current     int64
    16  	lastUpdate  int64
    17  	id          string
    18  	action      string
    19  	rateLimiter *rate.Limiter
    20  }
    21  
    22  // NewProgressReader creates a new ProgressReader.
    23  func NewProgressReader(in io.ReadCloser, out Output, size int64, id, action string) *Reader {
    24  	return &Reader{
    25  		in:          in,
    26  		out:         out,
    27  		size:        size,
    28  		id:          id,
    29  		action:      action,
    30  		rateLimiter: rate.NewLimiter(rate.Every(100*time.Millisecond), 1),
    31  	}
    32  }
    33  
    34  func (p *Reader) Read(buf []byte) (n int, err error) {
    35  	read, err := p.in.Read(buf)
    36  	p.current += int64(read)
    37  	updateEvery := int64(1024 * 512) // 512kB
    38  	if p.size > 0 {
    39  		// Update progress for every 1% read if 1% < 512kB
    40  		if increment := int64(0.01 * float64(p.size)); increment < updateEvery {
    41  			updateEvery = increment
    42  		}
    43  	}
    44  	if p.current-p.lastUpdate > updateEvery || err != nil {
    45  		p.updateProgress(err != nil && read == 0)
    46  		p.lastUpdate = p.current
    47  	}
    48  
    49  	return read, err
    50  }
    51  
    52  // Close closes the progress reader and its underlying reader.
    53  func (p *Reader) Close() error {
    54  	if p.current < p.size {
    55  		// print a full progress bar when closing prematurely
    56  		p.current = p.size
    57  		p.updateProgress(false)
    58  	}
    59  	return p.in.Close()
    60  }
    61  
    62  func (p *Reader) updateProgress(last bool) {
    63  	if last || p.current == p.size || p.rateLimiter.Allow() {
    64  		p.out.WriteProgress(Progress{ID: p.id, Action: p.action, Current: p.current, Total: p.size, LastUpdate: last})
    65  	}
    66  }