bib: integrate upload progress into overall progress

This commit improves the upload progress by integrating it with the
new overall progress flow. After the building is done the upload
will be part of the overall progress bar flow.

Author: mvo5

bib/cmd/bootc-image-builder/cloud.go

@@ -1,13 +1,15 @@
 package main
 
 import (
-	"github.com/cheggaaa/pb/v3"
-	"github.com/osbuild/bootc-image-builder/bib/internal/uploader"
-	"github.com/osbuild/images/pkg/cloud/awscloud"
 	"github.com/spf13/pflag"
+
+	"github.com/osbuild/images/pkg/cloud/awscloud"
+
+	"github.com/osbuild/bootc-image-builder/bib/internal/progress"
+	"github.com/osbuild/bootc-image-builder/bib/internal/uploader"
 )
 
-func uploadAMI(path, targetArch string, flags *pflag.FlagSet) error {
+func uploadAMI(pbar progress.ProgressBar, path, targetArch string, flags *pflag.FlagSet) error {
 	region, err := flags.GetString("aws-region")
 	if err != nil {
 		return err
@@ -20,23 +22,11 @@ func uploadAMI(path, targetArch string, flags *pflag.FlagSet) error {
 	if err != nil {
 		return err
 	}
-	progress, err := flags.GetString("progress")
-	if err != nil {
-		return err
-	}
 
 	client, err := awscloud.NewDefault(region)
 	if err != nil {
 		return err
 	}
 
-	// TODO: extract this as a helper once we add "uploadAzure" or
-	// similar. Eventually we may provide json progress here too.
-	var pbar *pb.ProgressBar
-	switch progress {
-	case "", "plain", "term":
-		pbar = pb.New(0)
-	}
-
 	return uploader.UploadAndRegister(client, path, bucketName, imageName, targetArch, pbar)
 }

bib/cmd/bootc-image-builder/main.go

@@ -487,16 +487,13 @@ func cmdBuild(cmd *cobra.Command, args []string) error {
 
 	pbar.SetMessagef("Build complete!")
 	if upload {
-		// XXX: pass our own progress.ProgressBar here
-		// *for now* just stop our own progress and let the uploadAMI
-		// progress take over - but we really need to fix this in a
-		// followup
-		pbar.Stop()
 		for idx, imgType := range imgTypes {
 			switch imgType {
 			case "ami":
+				pbar.Clear()
+				pbar.SetPulseMsgf("AWS uploading step")
 				diskpath := filepath.Join(outputDir, exports[idx], "disk.raw")
-				if err := uploadAMI(diskpath, targetArch, cmd.Flags()); err != nil {
+				if err := uploadAMI(pbar, diskpath, targetArch, cmd.Flags()); err != nil {
 					return fmt.Errorf("cannot upload AMI: %w", err)
 				}
 			default:

bib/internal/progress/progress.go

@@ -40,7 +40,7 @@ type ProgressBar interface {
 	// SetProgress sets the progress details at the given "level".
 	// Levels should start with "0" and increase as the nesting
 	// gets deeper.
-	SetProgress(level int, msg string, done int, total int) error
+	SetProgress(level int, msg string, done int64, total int64) error
 
 	// The high-level message that is displayed in a spinner
 	// that contains the current top level step, for bib this
@@ -107,7 +107,7 @@ func NewTerminalProgressBar() (ProgressBar, error) {
 	return b, nil
 }
 
-func (b *terminalProgressBar) SetProgress(subLevel int, msg string, done int, total int) error {
+func (b *terminalProgressBar) SetProgress(subLevel int, msg string, done int64, total int64) error {
 	// auto-add as needed, requires sublevels to get added in order
 	// i.e. adding 0 and then 2 will fail
 	switch {
@@ -282,7 +282,7 @@ func (b *plainProgressBar) Start() {
 func (b *plainProgressBar) Stop() {
 }
 
-func (b *plainProgressBar) SetProgress(subLevel int, msg string, done int, total int) error {
+func (b *plainProgressBar) SetProgress(subLevel int, msg string, done int64, total int64) error {
 	return nil
 }
 
@@ -322,7 +322,7 @@ func (b *debugProgressBar) Stop() {
 	fmt.Fprintf(b.w, "Stop progressbar\n")
 }
 
-func (b *debugProgressBar) SetProgress(subLevel int, msg string, done int, total int) error {
+func (b *debugProgressBar) SetProgress(subLevel int, msg string, done int64, total int64) error {
 	fmt.Fprintf(b.w, "%s[%v / %v] %s", strings.Repeat("  ", subLevel), done, total, msg)
 	fmt.Fprintf(b.w, "\n")
 	return nil
@@ -398,7 +398,7 @@ func runOSBuildWithProgress(pb ProgressBar, manifest []byte, store, outputDirect
 		}
 		i := 0
 		for p := st.Progress; p != nil; p = p.SubProgress {
-			if err := pb.SetProgress(i, p.Message, p.Done, p.Total); err != nil {
+			if err := pb.SetProgress(i, p.Message, int64(p.Done), int64(p.Total)); err != nil {
 				logrus.Warnf("cannot set progress: %v", err)
 			}
 			i++

bib/internal/uploader/aws.go

@@ -9,10 +9,11 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/s3/s3manager"
-	"github.com/cheggaaa/pb/v3"
 	"github.com/google/uuid"
 
 	"github.com/osbuild/images/pkg/arch"
+
+	"github.com/osbuild/bootc-image-builder/bib/internal/progress"
 )
 
 var osStdout io.Writer = os.Stdout
@@ -22,7 +23,21 @@ type AwsUploader interface {
 	Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode, importRole *string) (*string, *string, error)
 }
 
-func doUpload(a AwsUploader, file *os.File, bucketName, keyName string, pbar *pb.ProgressBar) (*s3manager.UploadOutput, error) {
+type proxyReader struct {
+	subLevel    int
+	r           io.Reader
+	pbar        progress.ProgressBar
+	done, total int64
+}
+
+func (r *proxyReader) Read(p []byte) (int, error) {
+	n, err := r.r.Read(p)
+	r.done += int64(n)
+	r.pbar.SetProgress(r.subLevel, "Uploading", r.done, r.total)
+	return n, err
+}
+
+func doUpload(a AwsUploader, file *os.File, bucketName, keyName string, pbar progress.ProgressBar) (*s3manager.UploadOutput, error) {
 	var r io.Reader = file
 
 	// TODO: extract this as a helper once we add "uploadAzure" or
@@ -32,26 +47,22 @@ func doUpload(a AwsUploader, file *os.File, bucketName, keyName string, pbar *pb
 		if err != nil {
 			return nil, fmt.Errorf("cannot stat upload: %v", err)
 		}
-		pbar.SetTotal(st.Size())
-		pbar.Set(pb.Bytes, true)
-		pbar.SetWriter(osStdout)
-		r = pbar.NewProxyReader(file)
-		pbar.Start()
-		defer pbar.Finish()
+		pbar.SetMessagef("Uploading %s to %s", file.Name(), bucketName)
+		r = &proxyReader{0, file, pbar, 0, st.Size()}
 	}
 
 	return a.UploadFromReader(r, bucketName, keyName)
 }
 
-func UploadAndRegister(a AwsUploader, filename, bucketName, imageName, targetArch string, pbar *pb.ProgressBar) error {
+func UploadAndRegister(a AwsUploader, filename, bucketName, imageName, targetArch string, pbar progress.ProgressBar) error {
 	file, err := os.Open(filename)
 	if err != nil {
 		return fmt.Errorf("cannot upload: %v", err)
 	}
 	defer file.Close()
 
 	keyName := fmt.Sprintf("%s-%s", uuid.New().String(), filepath.Base(filename))
-	fmt.Fprintf(osStdout, "Uploading %s to %s:%s\n", filename, bucketName, keyName)
+	pbar.SetMessagef("Uploading %s to %s:%s\n", filename, bucketName, keyName)
 	uploadOutput, err := doUpload(a, file, bucketName, keyName, pbar)
 	if err != nil {
 		return err