Skip to content

Commit

Permalink
bib: integrate upload progress into overall progress
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
mvo5 committed Dec 19, 2024
1 parent 86526a8 commit 060ddbc
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 37 deletions.
22 changes: 6 additions & 16 deletions bib/cmd/bootc-image-builder/cloud.go
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
}
9 changes: 3 additions & 6 deletions bib/cmd/bootc-image-builder/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
10 changes: 5 additions & 5 deletions bib/internal/progress/progress.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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++
Expand Down
31 changes: 21 additions & 10 deletions bib/internal/uploader/aws.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Check failure on line 36 in bib/internal/uploader/aws.go

View workflow job for this annotation

GitHub Actions / ⌨ Lint & unittests

Error return value of `r.pbar.SetProgress` is not checked (errcheck)
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
Expand All @@ -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
Expand Down

0 comments on commit 060ddbc

Please sign in to comment.