github.com/hashicorp/packer@v1.14.3/packer/progressbar.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: BUSL-1.1
     3  
     4  //go:build !solaris
     5  // +build !solaris
     6  
     7  package packer
     8  
     9  import (
    10  	"io"
    11  	"path/filepath"
    12  	"sync"
    13  
    14  	pb "github.com/cheggaaa/pb"
    15  )
    16  
    17  func ProgressBarConfig(bar *pb.ProgressBar, prefix string) {
    18  	bar.SetUnits(pb.U_BYTES)
    19  	bar.Prefix(prefix)
    20  }
    21  
    22  // UiProgressBar is a progress bar compatible with go-getter used in our
    23  // UI structs.
    24  type UiProgressBar struct {
    25  	lock sync.Mutex
    26  	pool *pb.Pool
    27  	pbs  int
    28  }
    29  
    30  func (p *UiProgressBar) TrackProgress(src string, currentSize, totalSize int64, stream io.ReadCloser) io.ReadCloser {
    31  	if p == nil {
    32  		return stream
    33  	}
    34  	p.lock.Lock()
    35  	defer p.lock.Unlock()
    36  
    37  	newPb := pb.New64(totalSize)
    38  	newPb.Set64(currentSize)
    39  	ProgressBarConfig(newPb, filepath.Base(src))
    40  
    41  	if p.pool == nil {
    42  		pool := pb.NewPool()
    43  		err := pool.Start()
    44  		if err != nil {
    45  			// here, we probably cannot lock
    46  			// stdout, so let's just return
    47  			// stream to avoid any error.
    48  			return stream
    49  		}
    50  		p.pool = pool
    51  	}
    52  	p.pool.Add(newPb)
    53  	reader := newPb.NewProxyReader(stream)
    54  
    55  	p.pbs++
    56  	return &readCloser{
    57  		Reader: reader,
    58  		close: func() error {
    59  			p.lock.Lock()
    60  			defer p.lock.Unlock()
    61  
    62  			newPb.Finish()
    63  			p.pbs--
    64  			if p.pbs <= 0 {
    65  				_ = p.pool.Stop()
    66  				p.pool = nil
    67  			}
    68  			return nil
    69  		},
    70  	}
    71  }
    72  
    73  type readCloser struct {
    74  	io.Reader
    75  	close func() error
    76  }
    77  
    78  func (c *readCloser) Close() error { return c.close() }