gopkg.in/mitchellh/packer.v1@v1.3.2/packer/progressbar.go (about)

     1  package packer
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/cheggaaa/pb"
     9  )
    10  
    11  // ProgressBar allows to graphically display
    12  // a self refreshing progress bar.
    13  type ProgressBar interface {
    14  	Start(total int64)
    15  	Add(current int64)
    16  	NewProxyReader(r io.Reader) (proxy io.Reader)
    17  	Finish()
    18  }
    19  
    20  // StackableProgressBar is a progress bar that
    21  // allows to track multiple downloads at once.
    22  // Every call to Start increments a counter that
    23  // will display the number of current loadings.
    24  // Every call to Start will add total to an internal
    25  // total that is the total displayed.
    26  // First call to Start will start a goroutine
    27  // that is waiting for every download to be finished.
    28  // Last call to Finish triggers a cleanup.
    29  // When all active downloads are finished
    30  // StackableProgressBar will clean itself to a default
    31  // state.
    32  type StackableProgressBar struct {
    33  	mtx   sync.Mutex // locks in Start, Finish, Add & NewProxyReader
    34  	Bar   BasicProgressBar
    35  	items int32
    36  	total int64
    37  
    38  	started             bool
    39  	ConfigProgressbarFN func(*pb.ProgressBar)
    40  }
    41  
    42  var _ ProgressBar = new(StackableProgressBar)
    43  
    44  func defaultProgressbarConfigFn(bar *pb.ProgressBar) {
    45  	bar.SetUnits(pb.U_BYTES)
    46  }
    47  
    48  func (spb *StackableProgressBar) start() {
    49  	bar := pb.New(0)
    50  	if spb.ConfigProgressbarFN == nil {
    51  		spb.ConfigProgressbarFN = defaultProgressbarConfigFn
    52  	}
    53  	spb.ConfigProgressbarFN(bar)
    54  
    55  	bar.Start()
    56  	spb.Bar.ProgressBar = bar
    57  	spb.started = true
    58  }
    59  
    60  func (spb *StackableProgressBar) Start(total int64) {
    61  	spb.mtx.Lock()
    62  
    63  	spb.total += total
    64  	spb.items++
    65  
    66  	if !spb.started {
    67  		spb.start()
    68  	}
    69  	spb.Bar.SetTotal64(spb.total)
    70  	spb.prefix()
    71  	spb.mtx.Unlock()
    72  }
    73  
    74  func (spb *StackableProgressBar) Add(total int64) {
    75  	spb.mtx.Lock()
    76  	defer spb.mtx.Unlock()
    77  	if spb.Bar.ProgressBar != nil {
    78  		spb.Bar.Add(total)
    79  	}
    80  }
    81  
    82  func (spb *StackableProgressBar) NewProxyReader(r io.Reader) io.Reader {
    83  	spb.mtx.Lock()
    84  	defer spb.mtx.Unlock()
    85  	return spb.Bar.NewProxyReader(r)
    86  }
    87  
    88  func (spb *StackableProgressBar) prefix() {
    89  	spb.Bar.ProgressBar.Prefix(fmt.Sprintf("%d items: ", spb.items))
    90  }
    91  
    92  func (spb *StackableProgressBar) Finish() {
    93  	spb.mtx.Lock()
    94  	defer spb.mtx.Unlock()
    95  
    96  	if spb.items > 0 {
    97  		spb.items--
    98  	}
    99  	if spb.items == 0 && spb.Bar.ProgressBar != nil {
   100  		// slef cleanup
   101  		spb.Bar.ProgressBar.Finish()
   102  		spb.Bar.ProgressBar = nil
   103  		spb.started = false
   104  		spb.total = 0
   105  		return
   106  	}
   107  }
   108  
   109  // BasicProgressBar is packer's basic progress bar.
   110  // Current implementation will always try to keep
   111  // itself at the bottom of a terminal.
   112  type BasicProgressBar struct {
   113  	*pb.ProgressBar
   114  }
   115  
   116  var _ ProgressBar = new(BasicProgressBar)
   117  
   118  func (bpb *BasicProgressBar) Start(total int64) {
   119  	bpb.SetTotal64(total)
   120  	bpb.ProgressBar.Start()
   121  }
   122  
   123  func (bpb *BasicProgressBar) Add(current int64) {
   124  	bpb.ProgressBar.Add64(current)
   125  }
   126  func (bpb *BasicProgressBar) NewProxyReader(r io.Reader) io.Reader {
   127  	return &ProxyReader{
   128  		Reader:      r,
   129  		ProgressBar: bpb,
   130  	}
   131  }
   132  func (bpb *BasicProgressBar) NewProxyReadCloser(r io.ReadCloser) io.ReadCloser {
   133  	return &ProxyReader{
   134  		Reader:      r,
   135  		ProgressBar: bpb,
   136  	}
   137  }
   138  
   139  // NoopProgressBar is a silent progress bar.
   140  type NoopProgressBar struct {
   141  }
   142  
   143  var _ ProgressBar = new(NoopProgressBar)
   144  
   145  func (npb *NoopProgressBar) Start(int64)                                      {}
   146  func (npb *NoopProgressBar) Add(int64)                                        {}
   147  func (npb *NoopProgressBar) Finish()                                          {}
   148  func (npb *NoopProgressBar) NewProxyReader(r io.Reader) io.Reader             { return r }
   149  func (npb *NoopProgressBar) NewProxyReadCloser(r io.ReadCloser) io.ReadCloser { return r }
   150  
   151  // ProxyReader implements io.ReadCloser but sends
   152  // count of read bytes to a progress bar
   153  type ProxyReader struct {
   154  	io.Reader
   155  	ProgressBar
   156  }
   157  
   158  func (r *ProxyReader) Read(p []byte) (n int, err error) {
   159  	n, err = r.Reader.Read(p)
   160  	r.ProgressBar.Add(int64(n))
   161  	return
   162  }
   163  
   164  // Close the reader if it implements io.Closer
   165  func (r *ProxyReader) Close() (err error) {
   166  	if closer, ok := r.Reader.(io.Closer); ok {
   167  		return closer.Close()
   168  	}
   169  	return
   170  }