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

@@ -8,6 +8,7 @@ import (
 	"os"
 	"os/exec"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/cheggaaa/pb/v3"
@@ -39,11 +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.
-	//
-	// Note that reducing depth is currently not supported, once
-	// a sub-progress is added it cannot be removed/hidden
-	// (but if required it can be added, its a SMOP)
-	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
@@ -58,6 +55,10 @@ type ProgressBar interface {
 	// like "Starting module org.osbuild.selinux"
 	SetMessagef(fmt string, args ...interface{})
 
+	// Clear clears all the subprogress/pulse/messages. This is
+	// useful when there is a need to reduce the sublevels.
+	Clear()
+
 	// Start will start rendering the progress information
 	Start()
 
@@ -89,7 +90,8 @@ type terminalProgressBar struct {
 
 	shutdownCh chan bool
 
-	out io.Writer
+	outMu sync.Mutex
+	out   io.Writer
 }
 
 // NewTerminalProgressBar creates a new default pb3 based progressbar suitable for
@@ -105,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 {
@@ -149,7 +151,29 @@ func (b *terminalProgressBar) SetMessagef(msg string, args ...interface{}) {
 	b.msgPb.Set("msg", shorten(fmt.Sprintf(msg, args...)))
 }
 
+func (b *terminalProgressBar) Clear() {
+	// ensure anything pending is output
+	b.render()
+
+	b.outMu.Lock()
+	defer b.outMu.Unlock()
+
+	// pulseMsg line + subLevel lines + message line
+	linesToClear := len(b.subLevelPbs) + 2
+	for i := 0; i < linesToClear; i++ {
+		fmt.Fprintf(b.out, "%s\n", ERASE_LINE)
+	}
+	fmt.Fprint(b.out, cursorUp(linesToClear))
+
+	b.subLevelPbs = nil
+	b.SetPulseMsgf("")
+	b.SetMessagef("")
+}
+
 func (b *terminalProgressBar) render() {
+	b.outMu.Lock()
+	defer b.outMu.Unlock()
+
 	var renderedLines int
 	fmt.Fprintf(b.out, "%s%s\n", ERASE_LINE, b.spinnerPb.String())
 	renderedLines++
@@ -258,10 +282,13 @@ 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
 }
 
+func (b *plainProgressBar) Clear() {
+}
+
 type debugProgressBar struct {
 	w io.Writer
 }
@@ -295,11 +322,14 @@ 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
 }
+func (b *debugProgressBar) Clear() {
+	fmt.Fprintf(b.w, "Clear progressbar\n")
+}
 
 // XXX: merge variant back into images/pkg/osbuild/osbuild-exec.go
 func RunOSBuild(pb ProgressBar, manifest []byte, store, outputDirectory string, exports, extraEnv []string) error {
@@ -368,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/progress/progress_test.go

@@ -58,6 +58,8 @@ func TestPlainProgress(t *testing.T) {
 	assert.Equal(t, "", buf.String())
 	pbar.Stop()
 	assert.Equal(t, "", buf.String())
+	pbar.Clear()
+	assert.Equal(t, "", buf.String())
 }
 
 func TestDebugProgress(t *testing.T) {
@@ -87,6 +89,10 @@ func TestDebugProgress(t *testing.T) {
 	pbar.Stop()
 	assert.Equal(t, "Stop progressbar\n", buf.String())
 	buf.Reset()
+
+	pbar.Clear()
+	assert.Equal(t, "Clear progressbar\n", buf.String())
+	buf.Reset()
 }
 
 func TestTermProgress(t *testing.T) {
@@ -102,12 +108,15 @@ func TestTermProgress(t *testing.T) {
 	pbar.SetMessagef("some-message")
 	err = pbar.SetProgress(0, "set-progress-msg", 0, 5)
 	assert.NoError(t, err)
+	pbar.Clear()
 	pbar.Stop()
 	assert.NoError(t, pbar.(*progress.TerminalProgressBar).Err())
 
 	assert.Contains(t, buf.String(), "[1 / 6] set-progress-msg")
 	assert.Contains(t, buf.String(), "[|] pulse-msg\n")
 	assert.Contains(t, buf.String(), "Message: some-message\n")
+	// check clear (clear 3 lines)
+	assert.Contains(t, buf.String(), "\x1b[2K\n\x1b[2K\n\x1b[2K\n")
 	// check shutdown
 	assert.Contains(t, buf.String(), progress.CURSOR_SHOW)
 }

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