bib: integrate upload progress into main progressbar

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 @@
 	"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 @@
 	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 @@
 		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