Merge pull request coreos#1162 from jlebon/pr/ore-copy-image-json

ore/aws/copy-image: stream results back to stdout
jcajka · Jan 14, 2020 · 383bdf5 · 383bdf5
2 parents 81b5eec + 5d6b39e
commit 383bdf5
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 23 deletions.
diff --git a/cmd/ore/aws/copy-image.go b/cmd/ore/aws/copy-image.go
@@ -19,6 +19,8 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/coreos/mantle/platform/api/aws"
+
 	"github.com/spf13/cobra"
 )
 
@@ -47,16 +49,19 @@ func runCopyImage(cmd *cobra.Command, args []string) error {
 		os.Exit(2)
 	}
 
-	amis, err := API.CopyImage(sourceImageID, args)
+	enc := json.NewEncoder(os.Stdout)
+	err := API.CopyImage(sourceImageID, args, func(region string, ami aws.ImageData) {
+		enc_err := enc.Encode(map[string]aws.ImageData{region: ami})
+		if enc_err != nil {
+			fmt.Fprintf(os.Stderr, "Couldn't encode result: %v\n", enc_err)
+			os.Exit(1)
+		}
+	})
+
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Couldn't copy images: %v\n", err)
 		os.Exit(1)
 	}
 
-	err = json.NewEncoder(os.Stdout).Encode(amis)
-	if err != nil {
-		fmt.Fprintf(os.Stderr, "Couldn't encode result: %v\n", err)
-		os.Exit(1)
-	}
 	return nil
 }
diff --git a/cmd/plume/prerelease.go b/cmd/plume/prerelease.go
@@ -542,14 +542,12 @@ func awsUploadToPartition(spec *channelSpec, part *awsPartitionSpec, imageName,
 		amis := map[string]string{}
 		if len(destRegions) > 0 {
 			plog.Printf("Replicating AMI %v...", imageID)
-			copiedAmis, err := api.CopyImage(imageID, destRegions)
+			err := api.CopyImage(imageID, destRegions, func(region string, ami aws.ImageData) {
+				amis[region] = ami.AMI
+			})
 			if err != nil {
 				return nil, fmt.Errorf("couldn't copy image: %v", err)
 			}
-
-			for region, data := range copiedAmis {
-				amis[region] = data.AMI
-			}
 		}
 		amis[part.BucketRegion] = imageID
 

diff --git a/platform/api/aws/images.go b/platform/api/aws/images.go
@@ -524,7 +524,7 @@ func (a *API) GrantLaunchPermission(imageID string, userIDs []string) error {
 	return nil
 }
 
-func (a *API) CopyImage(sourceImageID string, regions []string) (map[string]ImageData, error) {
+func (a *API) CopyImage(sourceImageID string, regions []string, cb func(string, ImageData)) error {
 	type result struct {
 		region string
 		data   ImageData
@@ -533,26 +533,26 @@ func (a *API) CopyImage(sourceImageID string, regions []string) (map[string]Imag
 
 	image, err := a.describeImage(sourceImageID)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
 	if *image.VirtualizationType == ec2.VirtualizationTypeParavirtual {
 		for _, region := range regions {
 			if !RegionSupportsPV(region) {
-				return nil, NoRegionPVSupport
+				return NoRegionPVSupport
 			}
 		}
 	}
 
 	snapshotID, err := getImageSnapshotID(image)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	describeSnapshotRes, err := a.ec2.DescribeSnapshots(&ec2.DescribeSnapshotsInput{
 		SnapshotIds: []*string{&snapshotID},
 	})
 	if err != nil {
-		return nil, fmt.Errorf("couldn't describe snapshot: %v", err)
+		return fmt.Errorf("couldn't describe snapshot: %v", err)
 	}
 	snapshot := describeSnapshotRes.Snapshots[0]
 
@@ -561,7 +561,7 @@ func (a *API) CopyImage(sourceImageID string, regions []string) (map[string]Imag
 		SnapshotId: aws.String(snapshotID),
 	})
 	if err != nil {
-		return nil, fmt.Errorf("couldn't describe createVolumePermission: %v", err)
+		return fmt.Errorf("couldn't describe createVolumePermission: %v", err)
 	}
 	createVolumePermissions := describeSnapshotAttributeRes.CreateVolumePermissions
 
@@ -570,7 +570,7 @@ func (a *API) CopyImage(sourceImageID string, regions []string) (map[string]Imag
 		ImageId:   aws.String(sourceImageID),
 	})
 	if err != nil {
-		return nil, fmt.Errorf("couldn't describe launch permissions: %v", err)
+		return fmt.Errorf("couldn't describe launch permissions: %v", err)
 	}
 	launchPermissions := describeAttributeRes.LaunchPermissions
 
@@ -594,19 +594,21 @@ func (a *API) CopyImage(sourceImageID string, regions []string) (map[string]Imag
 			ch <- res
 		}()
 	}
-	wg.Wait()
-	close(ch)
+	go func() {
+		wg.Wait()
+		close(ch)
+	}()
 
-	amis := make(map[string]ImageData)
 	for res := range ch {
 		if res.data.AMI != "" {
-			amis[res.region] = res.data
+			cb(res.region, res.data)
 		}
 		if err == nil {
 			err = res.err
 		}
 	}
-	return amis, err
+
+	return err
 }
 
 func (a *API) copyImageIn(sourceRegion, sourceImageID, name, description string, imageTags, snapshotTags []*ec2.Tag, launchPermissions []*ec2.LaunchPermission, createVolumePermissions []*ec2.CreateVolumePermission) (ImageData, error) {