From 86526a8034a6e610ca9d7cafd1f2730fa4a02839 Mon Sep 17 00:00:00 2001 From: Michael Vogt Date: Wed, 18 Dec 2024 16:06:55 +0100 Subject: [PATCH 1/2] progress: add new Clear() method to the progressBar This commit adds a new `Clear()` method on the progressBar so that we can clear subprogress easily. This is useful when switching from image building to image uploading, here we need a single progress again. --- bib/internal/progress/progress.go | 40 ++++++++++++++++++++++---- bib/internal/progress/progress_test.go | 9 ++++++ 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/bib/internal/progress/progress.go b/bib/internal/progress/progress.go index 3e7bf4ce..239a6058 100644 --- a/bib/internal/progress/progress.go +++ b/bib/internal/progress/progress.go @@ -8,6 +8,7 @@ import ( "os" "os/exec" "strings" + "sync" "time" "github.com/cheggaaa/pb/v3" @@ -39,10 +40,6 @@ 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 // The high-level message that is displayed in a spinner @@ -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 @@ -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++ @@ -262,6 +286,9 @@ func (b *plainProgressBar) SetProgress(subLevel int, msg string, done int, total return nil } +func (b *plainProgressBar) Clear() { +} + type debugProgressBar struct { w io.Writer } @@ -300,6 +327,9 @@ func (b *debugProgressBar) SetProgress(subLevel int, msg string, done int, total 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 { diff --git a/bib/internal/progress/progress_test.go b/bib/internal/progress/progress_test.go index 6b11d6f6..2f09a2fb 100644 --- a/bib/internal/progress/progress_test.go +++ b/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) } From 060ddbc38dcaba829a9289d476d1d32085cfe395 Mon Sep 17 00:00:00 2001 From: Michael Vogt Date: Wed, 18 Dec 2024 15:47:32 +0100 Subject: [PATCH 2/2] 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. --- bib/cmd/bootc-image-builder/cloud.go | 22 ++++++-------------- bib/cmd/bootc-image-builder/main.go | 9 +++----- bib/internal/progress/progress.go | 10 ++++----- bib/internal/uploader/aws.go | 31 +++++++++++++++++++--------- 4 files changed, 35 insertions(+), 37 deletions(-) diff --git a/bib/cmd/bootc-image-builder/cloud.go b/bib/cmd/bootc-image-builder/cloud.go index 98de143b..c1e02ede 100644 --- a/bib/cmd/bootc-image-builder/cloud.go +++ b/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) } diff --git a/bib/cmd/bootc-image-builder/main.go b/bib/cmd/bootc-image-builder/main.go index 9a5be5ea..eb7376b0 100644 --- a/bib/cmd/bootc-image-builder/main.go +++ b/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: diff --git a/bib/internal/progress/progress.go b/bib/internal/progress/progress.go index 239a6058..3dfb3ab3 100644 --- a/bib/internal/progress/progress.go +++ b/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++ diff --git a/bib/internal/uploader/aws.go b/bib/internal/uploader/aws.go index 359f5a4a..b839f874 100644 --- a/bib/internal/uploader/aws.go +++ b/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,18 +47,14 @@ 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) @@ -51,7 +62,7 @@ func UploadAndRegister(a AwsUploader, filename, bucketName, imageName, targetArc 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