Merge pull request #421 from shaheerkootteeri/Issue440-progress-imgpk…

…gpush

Adds progress bar for image push

pkg/imgpkg/bundle/contents.go

@@ -13,7 +13,9 @@ import (
 	regname "github.com/google/go-containerregistry/pkg/name"
 	regv1 "github.com/google/go-containerregistry/pkg/v1"
 	regremote "github.com/google/go-containerregistry/pkg/v1/remote"
+	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/internal/util"
 	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/plainimage"
+	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/registry"
 )
 
 const (
@@ -30,8 +32,9 @@ type Contents struct {
 //go:generate go run github.com/maxbrunsfeld/counterfeiter/v6 . ImagesMetadataWriter
 type ImagesMetadataWriter interface {
 	ImagesMetadata
-	WriteImage(regname.Reference, regv1.Image) error
+	WriteImage(regname.Reference, regv1.Image, chan regv1.Update) error
 	WriteTag(ref regname.Tag, taggagle regremote.Taggable) error
+	CloneWithLogger(logger util.ProgressLogger) registry.Registry
 }
 
 func NewContents(paths []string, excludedPaths []string) Contents {

pkg/imgpkg/bundle/locations_configs.go

@@ -132,7 +132,7 @@ func (r LocationsConfigs) Save(reg ImagesMetadataWriter, bundleRef name.Digest,
 
 	r.ui.Tracef("Pushing image\n")
 
-	_, err = plainimage.NewContents([]string{tmpDir}, nil).Push(locRef, nil, reg, ui)
+	_, err = plainimage.NewContents([]string{tmpDir}, nil).Push(locRef, nil, reg.CloneWithLogger(util.NewNoopProgressBar()), ui)
 	if err != nil {
 		// Immutable tag errors within registries are not standardized.
 		// Assume word "immutable" would be present in most cases.

pkg/imgpkg/internal/util/progress_logger.go

@@ -14,32 +14,40 @@ import (
 	"github.com/mattn/go-isatty"
 )
 
+// ProgressLogger Progress bar
 type ProgressLogger interface {
-	Start(progress <-chan regv1.Update)
+	Start(ctx context.Context, progress <-chan regv1.Update)
 	End()
 }
 
 // NewProgressBar constructor to build a ProgressLogger responsible for printing out a progress bar using updates when
 // writing to a registry via ggcr
 func NewProgressBar(ui goui.UI, finalMessage, errorMessagePrefix string) ProgressLogger {
-	ctx, cancel := context.WithCancel(context.Background())
 	if isatty.IsTerminal(os.Stdout.Fd()) {
-		return &ProgressBarLogger{ctx: ctx, cancelFunc: cancel, ui: ui, finalMessage: finalMessage, errorMessagePrefix: errorMessagePrefix}
+		return &ProgressBarLogger{ui: ui, finalMessage: finalMessage, errorMessagePrefix: errorMessagePrefix}
 	}
 
-	return &ProgressBarNoTTYLogger{ui: ui, ctx: ctx, cancelFunc: cancel, finalMessage: finalMessage}
+	return &ProgressBarNoTTYLogger{ui: ui, finalMessage: finalMessage}
 }
 
+// NewNoopProgressBar constructs a Noop Progress bar that will not display anything
+func NewNoopProgressBar() ProgressLogger {
+	return &ProgressBarNoTTYLogger{}
+}
+
+// ProgressBarLogger display progress bar on output
 type ProgressBarLogger struct {
-	ctx                context.Context
 	cancelFunc         context.CancelFunc
 	bar                *pb.ProgressBar
 	ui                 goui.UI
 	finalMessage       string
 	errorMessagePrefix string
 }
 
-func (l *ProgressBarLogger) Start(progressChan <-chan regv1.Update) {
+// Start the display of the Progress Bar
+func (l *ProgressBarLogger) Start(ctx context.Context, progressChan <-chan regv1.Update) {
+	ctx, cancelFunc := context.WithCancel(ctx)
+	l.cancelFunc = cancelFunc
 	// Add a new empty line to separate the progress bar from prior output
 	fmt.Println()
 	l.bar = pb.New64(0)
@@ -48,7 +56,7 @@ func (l *ProgressBarLogger) Start(progressChan <-chan regv1.Update) {
 	go func() {
 		for {
 			select {
-			case <-l.ctx.Done():
+			case <-ctx.Done():
 				return
 			case update := <-progressChan:
 				if update.Error != nil {
@@ -70,32 +78,43 @@ func (l *ProgressBarLogger) Start(progressChan <-chan regv1.Update) {
 	}()
 }
 
+// End stops the progress bar and writes the final message
 func (l *ProgressBarLogger) End() {
-	l.cancelFunc()
+	if l.cancelFunc != nil {
+		l.cancelFunc()
+	}
 	l.bar.Finish()
 	l.ui.BeginLinef("\n%s", l.finalMessage)
 }
 
+// ProgressBarNoTTYLogger does not display the progress bar
 type ProgressBarNoTTYLogger struct {
-	ctx          context.Context
 	cancelFunc   context.CancelFunc
 	ui           goui.UI
 	finalMessage string
 }
 
-func (l *ProgressBarNoTTYLogger) Start(progressChan <-chan regv1.Update) {
+// Start consuming the progress channel but does not display anything
+func (l *ProgressBarNoTTYLogger) Start(ctx context.Context, progressChan <-chan regv1.Update) {
+	ctx, cancelFunc := context.WithCancel(ctx)
+	l.cancelFunc = cancelFunc
 	go func() {
 		for {
 			select {
-			case <-l.ctx.Done():
+			case <-ctx.Done():
 				return
 			case <-progressChan:
 			}
 		}
 	}()
 }
 
+// End Write the final message
 func (l *ProgressBarNoTTYLogger) End() {
-	l.cancelFunc()
-	l.ui.BeginLinef(l.finalMessage)
+	if l.cancelFunc != nil {
+		l.cancelFunc()
+	}
+	if l.ui != nil && l.finalMessage != "" {
+		l.ui.BeginLinef(l.finalMessage)
+	}
 }

pkg/imgpkg/plainimage/contents.go

@@ -23,7 +23,7 @@ type Contents struct {
 }
 
 type ImagesWriter interface {
-	WriteImage(regname.Reference, regv1.Image) error
+	WriteImage(regname.Reference, regv1.Image, chan regv1.Update) error
 	WriteTag(ref regname.Tag, taggagle regremote.Taggable) error
 }
 
@@ -46,7 +46,8 @@ func (i Contents) Push(uploadRef regname.Tag, labels map[string]string, writer I
 
 	defer img.Remove()
 
-	err = writer.WriteImage(uploadRef, img)
+	err = writer.WriteImage(uploadRef, img, nil)
+
 	if err != nil {
 		return "", fmt.Errorf("Writing '%s': %s", uploadRef.Name(), err)
 	}

pkg/imgpkg/registry/registry.go

@@ -20,6 +20,7 @@ import (
 	regv1 "github.com/google/go-containerregistry/pkg/v1"
 	regremote "github.com/google/go-containerregistry/pkg/v1/remote"
 	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
+	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/internal/util"
 	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/registry/auth"
 )
 
@@ -50,13 +51,14 @@ type Registry interface {
 	FirstImageExists(digests []string) (string, error)
 
 	MultiWrite(imageOrIndexesToUpload map[regname.Reference]regremote.Taggable, concurrency int, updatesCh chan regv1.Update) error
-	WriteImage(reference regname.Reference, image regv1.Image) error
+	WriteImage(regname.Reference, regv1.Image, chan regv1.Update) error
 	WriteIndex(reference regname.Reference, index regv1.ImageIndex) error
 	WriteTag(tag regname.Tag, taggable regremote.Taggable) error
 
 	ListTags(repo regname.Repository) ([]string, error)
 
 	CloneWithSingleAuth(imageRef regname.Tag) (Registry, error)
+	CloneWithLogger(logger util.ProgressLogger) Registry
 }
 
 // ImagesReader Interface for Reading Images
@@ -74,11 +76,12 @@ type ImagesReader interface {
 type ImagesReaderWriter interface {
 	ImagesReader
 	MultiWrite(imageOrIndexesToUpload map[regname.Reference]regremote.Taggable, concurrency int, updatesCh chan regv1.Update) error
-	WriteImage(regname.Reference, regv1.Image) error
+	WriteImage(regname.Reference, regv1.Image, chan regv1.Update) error
 	WriteIndex(regname.Reference, regv1.ImageIndex) error
 	WriteTag(regname.Tag, regremote.Taggable) error
 
 	CloneWithSingleAuth(imageRef regname.Tag) (Registry, error)
+	CloneWithLogger(logger util.ProgressLogger) Registry
 }
 
 var _ Registry = &SimpleRegistry{}
@@ -190,6 +193,18 @@ func (r SimpleRegistry) CloneWithSingleAuth(imageRef regname.Tag) (Registry, err
 	}, nil
 }
 
+// CloneWithLogger Clones the provided registry updating the progress logger to NoTTYLogger
+// that does not display the progress bar
+func (r SimpleRegistry) CloneWithLogger(_ util.ProgressLogger) Registry {
+	return &SimpleRegistry{
+		remoteOpts:      r.remoteOpts,
+		refOpts:         r.refOpts,
+		keychain:        r.keychain,
+		roundTrippers:   r.roundTrippers,
+		transportAccess: &sync.Mutex{},
+	}
+}
+
 // readOpts Returns the readOpts + the keychain
 func (r *SimpleRegistry) readOpts(ref regname.Reference) ([]regremote.Option, error) {
 	rt, err := r.transport(ref, ref.Scope(transport.PullScope))
@@ -319,7 +334,7 @@ func (r *SimpleRegistry) MultiWrite(imageOrIndexesToUpload map[regname.Reference
 }
 
 // WriteImage Upload Image to registry
-func (r *SimpleRegistry) WriteImage(ref regname.Reference, img regv1.Image) error {
+func (r *SimpleRegistry) WriteImage(ref regname.Reference, img regv1.Image, updatesCh chan regv1.Update) error {
 	if err := r.validateRef(ref); err != nil {
 		return err
 	}
@@ -332,7 +347,9 @@ func (r *SimpleRegistry) WriteImage(ref regname.Reference, img regv1.Image) erro
 	if err != nil {
 		return err
 	}
-
+	if updatesCh != nil {
+		opts = append(opts, regremote.WithProgress(updatesCh))
+	}
 	err = regremote.Write(overriddenRef, img, opts...)
 	if err != nil {
 		return fmt.Errorf("Writing image: %s", err)