Skip to content

Commit

Permalink
bib: show progress on ami upload
Browse files Browse the repository at this point in the history
Add a new `--progress` option that defaults to `text` and show
upload progress when uploading an AMI.
  • Loading branch information
mvo5 committed Mar 7, 2024
1 parent e1cf3df commit 9ac04fb
Show file tree
Hide file tree
Showing 9 changed files with 203 additions and 15 deletions.
16 changes: 15 additions & 1 deletion bib/cmd/bootc-image-builder/cloud.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package main

import (
"github.com/cheggaaa/pb"
"github.com/osbuild/bootc-image-builder/bib/internal/uploader"
"github.com/osbuild/images/pkg/cloud/awscloud"
"github.com/spf13/pflag"
Expand All @@ -19,10 +20,23 @@ 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
}
return uploader.UploadAndRegister(client, path, bucketName, imageName, targetArch)

// 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 "text":
pbar = pb.New(0)
}

return uploader.UploadAndRegister(client, path, bucketName, imageName, targetArch, pbar)
}
1 change: 1 addition & 0 deletions bib/cmd/bootc-image-builder/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -359,6 +359,7 @@ func run() error {
buildCmd.Flags().String("aws-region", "", "target region for AWS uploads (only for type=ami)")
buildCmd.Flags().String("aws-bucket", "", "target S3 bucket name for intermediate storage when creating AMI (only for type=ami)")
buildCmd.Flags().String("aws-ami-name", "", "name for the AMI in AWS (only for type=ami)")
buildCmd.Flags().String("progress", "text", "type of progress bar to use")

// flag rules
for _, dname := range []string{"output", "store", "rpmmd"} {
Expand Down
2 changes: 1 addition & 1 deletion bib/cmd/upload/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ func uploadAMI(cmd *cobra.Command, args []string) {
targetArch, err := flags.GetString("target-arch")
check(err)

check(uploader.UploadAndRegister(client, filename, bucketName, imageName, targetArch))
check(uploader.UploadAndRegister(client, filename, bucketName, imageName, targetArch, nil))
}

func setupCLI() *cobra.Command {
Expand Down
1 change: 1 addition & 0 deletions bib/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ go 1.19

require (
github.com/aws/aws-sdk-go v1.50.23
github.com/cheggaaa/pb v1.0.29
github.com/google/uuid v1.6.0
github.com/osbuild/images v0.38.0
github.com/sirupsen/logrus v1.9.3
Expand Down
12 changes: 12 additions & 0 deletions bib/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ github.com/aws/aws-sdk-go v1.50.23/go.mod h1:LF8svs817+Nz+DmiMQKTO3ubZ/6IaTpq3Tj
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
github.com/cheggaaa/pb v1.0.29 h1:FckUN5ngEk2LpvuG0fw1GEFx6LtyY2pWI/Z2QgCnEYo=
github.com/cheggaaa/pb v1.0.29/go.mod h1:W40334L7FMC5JKWldsTWbdGjLo0RxUKK73K+TuPxX30=
github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc=
github.com/containerd/cgroups/v3 v3.0.2 h1:f5WFqIVSgo5IZmtTT3qVBo6TzI1ON6sycSBKkymb9L0=
Expand Down Expand Up @@ -72,6 +74,8 @@ github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7
github.com/facebookgo/clock v0.0.0-20150410010913-600d898af40a h1:yDWHCSQ40h88yih2JAcL6Ls/kVkSE8GFACTGVnMPruw=
github.com/facebookgo/limitgroup v0.0.0-20150612190941-6abd8d71ec01 h1:IeaD1VDVBPlx3viJT9Md8if8IxxJnO+x0JCGb054heg=
github.com/facebookgo/muster v0.0.0-20150708232844-fd3d7953fd52 h1:a4DFiKFJiDRGFD1qIcqGLX/WlUMD9dyLSLDt+9QZgt8=
github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU=
github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs=
github.com/go-jose/go-jose/v3 v3.0.1 h1:pWmKFVtt+Jl0vBZTIpz/eAKwsm6LkIxDVVbFHKkchhA=
github.com/go-jose/go-jose/v3 v3.0.1/go.mod h1:RNkWWRld676jZEYoV3+XK8L2ZnNSvIsxFMht0mSX+u8=
github.com/go-openapi/analysis v0.21.2/go.mod h1:HZwRk4RRisyG8vx2Oe6aqeSQcoxRp47Xkp3+K6q+LdY=
Expand Down Expand Up @@ -230,6 +234,12 @@ github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
github.com/markbates/oncer v0.0.0-20181203154359-bf2de49a0be2/go.mod h1:Ld9puTsIW75CHf65OeIOkyKbteujpZVXDpWK6YGZbxE=
github.com/markbates/safe v1.0.1/go.mod h1:nAqgmRi7cY2nqMc92/bSEeQA+R4OheNU2T1kNSCBdG0=
github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
github.com/mattn/go-isatty v0.0.11/go.mod h1:PhnuNfih5lzO57/f3n+odYbM4JtupLOxQOAqxQCu2WE=
github.com/mattn/go-isatty v0.0.17 h1:BTarxUcIeDqL27Mc+vyvdWYSL28zpIhv3RoTdsLMPng=
github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/mattn/go-shellwords v1.0.12 h1:M2zGm7EW6UQJvDeQxo4T51eKPurbeFbe8WtebGE2xrk=
Expand Down Expand Up @@ -416,11 +426,13 @@ golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190403152447-81d4e9dc473e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190419153524-e8e3143a4f4a/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190531175056-4c3a928424d2/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210420072515-93ed5bcd2bfe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
Expand Down
55 changes: 46 additions & 9 deletions bib/internal/uploader/aws.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,33 +2,70 @@ package uploader

import (
"fmt"
"io"
"os"
"path/filepath"

"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"
"github.com/google/uuid"

"github.com/osbuild/images/pkg/arch"
"github.com/osbuild/images/pkg/cloud/awscloud"
)

func UploadAndRegister(a *awscloud.AWS, filename, bucketName, imageName, targetArch string) error {
keyName := fmt.Sprintf("%s-%s", uuid.New().String(), filepath.Base(filename))
var osStdout io.Writer = os.Stdout

type AwsUploader interface {
UploadFromReader(r io.Reader, bucketName, keyName string) (*s3manager.UploadOutput, error)
Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode *string) (*string, *string, error)
}

func doUpload(a AwsUploader, file *os.File, bucketName, keyName string, pbar *pb.ProgressBar) (*s3manager.UploadOutput, error) {
var r io.Reader = file

// TODO: extract this as a helper once we add "uploadAzure" or
// similar.
if pbar != nil {
st, err := file.Stat()
if err != nil {
return nil, fmt.Errorf("cannot stat upload: %v", err)
}
pbar.Total = st.Size()
pbar.Units = pb.U_BYTES
pbar.Output = osStdout
r = pbar.NewProxyReader(file)
pbar.Start()
defer pbar.Finish()
}

fmt.Printf("Uploading %s to %s:%s\n", filename, bucketName, keyName)
uploadOutput, err := a.Upload(filename, bucketName, keyName)
return a.UploadFromReader(r, bucketName, keyName)
}

func UploadAndRegister(a AwsUploader, filename, bucketName, imageName, targetArch string, pbar *pb.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)
uploadOutput, err := doUpload(a, file, bucketName, keyName, pbar)
if err != nil {
return err
}
fmt.Printf("File uploaded to %s\n", aws.StringValue(&uploadOutput.Location))
fmt.Fprintf(osStdout, "File uploaded to %s\n", aws.StringValue(&uploadOutput.Location))

if targetArch == "" {
targetArch = arch.Current().String()
}
bootMode := ec2.BootModeValuesUefiPreferred
fmt.Printf("Registering AMI %s\n", imageName)
fmt.Fprintf(osStdout, "Registering AMI %s\n", imageName)
ami, snapshot, err := a.Register(imageName, bucketName, keyName, nil, targetArch, &bootMode)
fmt.Printf("Deleted S3 object %s:%s\n", bucketName, keyName)
fmt.Printf("AMI registered: %s\nSnapshot ID: %s\n", aws.StringValue(ami), aws.StringValue(snapshot))
fmt.Fprintf(osStdout, "Deleted S3 object %s:%s\n", bucketName, keyName)
fmt.Fprintf(osStdout, "AMI registered: %s\nSnapshot ID: %s\n", aws.StringValue(ami), aws.StringValue(snapshot))
if err != nil {
return err
}
Expand Down
86 changes: 86 additions & 0 deletions bib/internal/uploader/aws_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
package uploader_test

import (
"bytes"
"io"
"os"
"path/filepath"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"github.com/aws/aws-sdk-go/service/s3/s3manager"
"github.com/cheggaaa/pb"

"github.com/osbuild/bootc-image-builder/bib/internal/uploader"
)

type FakeAwsUploader struct {
uploadCalled int
registerCalled int
}

func (f *FakeAwsUploader) UploadFromReader(r io.Reader, bucketName, keyName string) (*s3manager.UploadOutput, error) {
f.uploadCalled++

if _, err := io.ReadAll(r); err != nil {
panic(err)
}

return &s3manager.UploadOutput{Location: "some-location"}, nil
}

func (f *FakeAwsUploader) Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode *string) (*string, *string, error) {
f.registerCalled++

s1 := "ret1"
s2 := "ret2"
return &s1, &s2, nil
}

func TestUploadAndRegisterNoProgressBar(t *testing.T) {
fakeStdout := bytes.NewBuffer(nil)
restore := uploader.MockOsStdout(fakeStdout)
defer restore()

fakeDiskFile := filepath.Join(t.TempDir(), "fake-disk.img")
err := os.WriteFile(fakeDiskFile, nil, 0644)
require.Nil(t, err)
fakeUploader := &FakeAwsUploader{}

err = uploader.UploadAndRegister(fakeUploader, fakeDiskFile, "bucketName", "imageName", "", nil)
require.Nil(t, err)

assert.Equal(t, fakeUploader.uploadCalled, 1)
assert.Equal(t, fakeUploader.registerCalled, 1)

assert.Contains(t, fakeStdout.String(), "Uploading ")
assert.Contains(t, fakeStdout.String(), "Registering AMI ")
}

func TestUploadAndRegisterProgressBar(t *testing.T) {
fakeStdout := bytes.NewBuffer(nil)
restore := uploader.MockOsStdout(fakeStdout)
defer restore()

fakeDiskFile := filepath.Join(t.TempDir(), "fake-disk.img")
err := os.WriteFile(fakeDiskFile, nil, 0644)
require.Nil(t, err)
err = os.Truncate(fakeDiskFile, 10*1024*1024)
require.Nil(t, err)

fakeUploader := &FakeAwsUploader{}

pbar := pb.New(0)

err = uploader.UploadAndRegister(fakeUploader, fakeDiskFile, "bucketName", "imageName", "", pbar)
require.Nil(t, err)

assert.Equal(t, fakeUploader.uploadCalled, 1)
assert.Equal(t, fakeUploader.registerCalled, 1)

assert.Contains(t, fakeStdout.String(), "Uploading ")
assert.Contains(t, fakeStdout.String(), "10.00 MiB / 10.00 MiB [============================================] 100.00%")
assert.Contains(t, fakeStdout.String(), "Registering AMI ")
}
13 changes: 13 additions & 0 deletions bib/internal/uploader/export_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
package uploader

import (
"io"
)

func MockOsStdout(new io.Writer) (restore func()) {
saved := osStdout
osStdout = new
return func() {
osStdout = saved
}
}
32 changes: 28 additions & 4 deletions test/test_build.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ class ImageBuildResult(NamedTuple):
img_arch: str
username: str
password: str
bib_output: str
journal_output: str
metadata: dict = {}

Expand Down Expand Up @@ -71,6 +72,7 @@ def image_type_fixture(shared_tmpdir, build_container, request, force_aws_upload
output_path.mkdir(exist_ok=True)

journal_log_path = output_path / "journal.log"
bib_output_path = output_path / "bib-output.log"
artifact = {
"qcow2": pathlib.Path(output_path) / "qcow2/disk.qcow2",
"ami": pathlib.Path(output_path) / "image/disk.raw",
Expand All @@ -84,7 +86,8 @@ def image_type_fixture(shared_tmpdir, build_container, request, force_aws_upload
if generated_img.exists():
print(f"NOTE: reusing cached image {generated_img}")
journal_output = journal_log_path.read_text(encoding="utf8")
yield ImageBuildResult(image_type, generated_img, target_arch, username, password, journal_output)
bib_output = bib_output_path.read_text(encoding="utf8")
yield ImageBuildResult(image_type, generated_img, target_arch, username, password, bib_output, journal_output)
return

# no image yet, build it
Expand Down Expand Up @@ -130,7 +133,7 @@ def image_type_fixture(shared_tmpdir, build_container, request, force_aws_upload
raise RuntimeError("AWS credentials not available (upload forced)")

# run container to deploy an image into a bootable disk and upload to a cloud service if applicable
subprocess.check_call([
p = subprocess.Popen([
"podman", "run", "--rm",
"--privileged",
"--security-opt", "label=type:unconfined_t",
Expand All @@ -143,7 +146,18 @@ def image_type_fixture(shared_tmpdir, build_container, request, force_aws_upload
"--type", image_type,
*upload_args,
*target_arch_args,
])
], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
# not using subprocss.check_output() to ensure we get live output
# during the text
bib_output = ""
while True:
line = p.stdout.readline()
if not line:
break
print(line, end="")
bib_output += line
p.wait(timeout=10)

journal_output = testutil.journal_after_cursor(cursor)
metadata = {}
if image_type == "ami" and upload_args:
Expand All @@ -154,8 +168,11 @@ def del_ami():
request.addfinalizer(del_ami)

journal_log_path.write_text(journal_output, encoding="utf8")
bib_output_path.write_text(bib_output, encoding="utf8")

yield ImageBuildResult(image_type, generated_img, target_arch, username, password, journal_output, metadata)
yield ImageBuildResult(
image_type, generated_img, target_arch, username, password,
bib_output, journal_output, metadata)
# Try to cache as much as possible
disk_usage = shutil.disk_usage(generated_img)
print(f"NOTE: disk usage after {generated_img}: {disk_usage.free / 1_000_000} / {disk_usage.total / 1_000_000}")
Expand Down Expand Up @@ -197,6 +214,13 @@ def test_ami_boots_in_aws(image_type, force_aws_upload):
raise RuntimeError("AWS credentials not available")
pytest.skip("AWS credentials not available (upload not forced)")

# check that upload progress is in the output log. Uploads looks like:
#
# Uploading /output/image/disk.raw to bootc-image-builder-ci:aac64b64-6e57-47df-9730-54763061d84b-disk.raw
# 0 B / 10.00 GiB 0.00%
# In the tests with no pty no progress bar is shown in the output just
# xx / yy zz%
assert " 100.00%\n" in image_type.bib_output
with AWS(image_type.metadata["ami_id"]) as test_vm:
exit_status, _ = test_vm.run("true", user=image_type.username, password=image_type.password)
assert exit_status == 0
Expand Down

0 comments on commit 9ac04fb

Please sign in to comment.