Harden staged push polling

Author: chruffins

pkg/cmd/push.go

@@ -95,10 +95,31 @@ func handleLocalPush(ctx context.Context, cmd *cli.Command) error {
 		targetName = args[1]
 	}
 
-	return pushLocalImage(ctx, cmd, sourceImage, targetName, nil)
+	return pushLocalImage(ctx, cmd, sourceImage, targetName)
 }
 
-func pushLocalImage(ctx context.Context, cmd *cli.Command, sourceImage, targetName string, img v1.Image) error {
+func pushLocalImage(ctx context.Context, cmd *cli.Command, sourceImage, targetName string) error {
+	fmt.Fprintf(os.Stderr, "Loading image %s from Docker...\n", sourceImage)
+	img, err := loadDockerImage(sourceImage)
+	if err != nil {
+		return err
+	}
+	return uploadLocalImage(ctx, cmd, targetName, img)
+}
+
+func loadDockerImage(image string) (v1.Image, error) {
+	srcRef, err := name.ParseReference(image)
+	if err != nil {
+		return nil, fmt.Errorf("invalid source image: %w", err)
+	}
+	img, err := daemon.Image(srcRef)
+	if err != nil {
+		return nil, fmt.Errorf("load image: %w", err)
+	}
+	return img, nil
+}
+
+func uploadLocalImage(ctx context.Context, cmd *cli.Command, targetName string, img v1.Image) error {
 	baseURL := resolveBaseURL(cmd)
 
 	parsedURL, err := url.Parse(baseURL)
@@ -114,11 +135,6 @@ func pushLocalImage(ctx context.Context, cmd *cli.Command, sourceImage, targetNa
 
 	registryHost := parsedURL.Host
 
-	srcRef, err := name.ParseReference(sourceImage)
-	if err != nil {
-		return fmt.Errorf("invalid source image: %w", err)
-	}
-
 	// Build and validate the target before opening the Docker daemon. The
 	// server computes the image digest from the manifest, while the tag keeps
 	// the image addressable with Docker-like image names after the push.
@@ -132,14 +148,6 @@ func pushLocalImage(ctx context.Context, cmd *cli.Command, sourceImage, targetNa
 		return fmt.Errorf("invalid target: %w", err)
 	}
 
-	if img == nil {
-		fmt.Fprintf(os.Stderr, "Loading image %s from Docker...\n", sourceImage)
-		img, err = daemon.Image(srcRef)
-		if err != nil {
-			return fmt.Errorf("load image: %w", err)
-		}
-	}
-
 	fmt.Fprintf(os.Stderr, "The push refers to repository [%s]\n", dstRef.Context().Name())
 
 	token := resolveAPIKey()

pkg/cmd/push_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/kernel/hypeman-go"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -32,6 +33,30 @@ func TestRenderPushProgressNonInteractive(t *testing.T) {
 	assert.Empty(t, output.String())
 }
 
+func TestPushStatusText(t *testing.T) {
+	lastBytes := int64(0)
+
+	assert.Equal(t, "queued · registry.example.com/app:v1", pushStatusText(&hypeman.Push{
+		Status: hypeman.PushStatusQueued,
+		Target: "registry.example.com/app:v1",
+	}, &lastBytes))
+	assert.Equal(t, "pushing 2.0 KB · 2 layers · registry.example.com/app:v1", pushStatusText(&hypeman.Push{
+		Status: hypeman.PushStatusPushing,
+		Bytes:  2048,
+		Layers: 2,
+		Target: "registry.example.com/app:v1",
+	}, &lastBytes))
+	assert.Equal(t, int64(2048), lastBytes)
+	assert.Equal(t, "pushed · digest: sha256:abc", pushStatusText(&hypeman.Push{
+		Status: hypeman.PushStatusPushed,
+		Digest: "sha256:abc",
+	}, &lastBytes))
+	assert.Equal(t, "failed · registry unavailable", pushStatusText(&hypeman.Push{
+		Status: hypeman.PushStatusFailed,
+		Error:  "registry unavailable",
+	}, &lastBytes))
+}
+
 func TestPushStatusRenderer(t *testing.T) {
 	var output bytes.Buffer
 	renderer := &pushStatusRenderer{output: &output, interactive: true}

pkg/cmd/pushcmd.go

@@ -10,7 +10,6 @@ import (
 	"time"
 
 	"github.com/google/go-containerregistry/pkg/name"
-	"github.com/google/go-containerregistry/pkg/v1/daemon"
 	"github.com/kernel/hypeman-go"
 	"github.com/kernel/hypeman-go/option"
 	"github.com/tidwall/gjson"
@@ -95,24 +94,20 @@ func handleRemotePushTarget(ctx context.Context, cmd *cli.Command, target string
 	// The one-argument form follows Docker's local-tag flow: TARGET must be
 	// present in the local Docker daemon before it can be staged and pushed.
 	// Cached Hypeman images use the explicit IMAGE TARGET form instead.
-	srcRef, err := name.ParseReference(target)
-	if err != nil {
-		return err
-	}
-	img, err := daemon.Image(srcRef)
+	img, err := loadDockerImage(target)
 	if err != nil {
 		return fmt.Errorf("load local Docker image %q: %w; tag it first or use hypeman push <image> <target> for a cached Hypeman image", target, err)
 	}
 
 	fmt.Fprintf(os.Stderr, "Staging local image %s in Hypeman...\n", target)
-	if err := pushLocalImage(ctx, cmd, target, target, img); err != nil {
+	if err := uploadLocalImage(ctx, cmd, target, img); err != nil {
 		return err
 	}
 
 	client := hypeman.NewClient(getDefaultRequestOptions(cmd)...)
-	imported, err := client.Images.Get(ctx, url.PathEscape(target))
+	imported, err := waitForImageRecord(ctx, &client, target)
 	if err != nil {
-		return fmt.Errorf("get staged image %s: %w", target, err)
+		return err
 	}
 	if err := waitForImageReady(ctx, &client, imported); err != nil {
 		return err
@@ -134,6 +129,27 @@ func validateRemotePushTarget(target string) error {
 	return nil
 }
 
+func waitForImageRecord(ctx context.Context, client *hypeman.Client, imageName string) (*hypeman.Image, error) {
+	ticker := time.NewTicker(300 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		img, err := client.Images.Get(ctx, url.PathEscape(imageName))
+		if err == nil {
+			return img, nil
+		}
+		if !isNotFoundError(err) {
+			return nil, fmt.Errorf("get staged image %s: %w", imageName, err)
+		}
+
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-ticker.C:
+		}
+	}
+}
+
 func handlePushCreate(ctx context.Context, cmd *cli.Command) error {
 	args := cmd.Args().Slice()
 	if len(args) != 2 {
@@ -226,41 +242,29 @@ func waitForPush(ctx context.Context, client *hypeman.Client, push *hypeman.Push
 			output:      os.Stderr,
 			interactive: term.IsTerminal(int(os.Stderr.Fd())),
 		}
+		defer renderer.finish()
 	}
 
-	current := push
 	var lastBytes int64
-	for {
-		if renderer != nil {
-			switch current.Status {
-			case hypeman.PushStatusQueued:
-				renderer.update(fmt.Sprintf("queued · %s", current.Target))
-			case hypeman.PushStatusPushing:
-				if current.Bytes > lastBytes {
-					lastBytes = current.Bytes
-				}
-				renderer.update(fmt.Sprintf("pushing %s · %d layers · %s", formatBytes(lastBytes), current.Layers, current.Target))
-			case hypeman.PushStatusPushed:
-				renderer.update(fmt.Sprintf("pushed · digest: %s", current.Digest))
-			case hypeman.PushStatusFailed:
-				message := current.Error
-				if message == "" {
-					message = "unknown error"
-				}
-				renderer.update("failed · " + message)
-			}
+	return pollPush(ctx, client, push, opts, func(current *hypeman.Push) {
+		if renderer == nil {
+			return
 		}
+		if message := pushStatusText(current, &lastBytes); message != "" {
+			renderer.update(message)
+		}
+	})
+}
+
+func pollPush(ctx context.Context, client *hypeman.Client, push *hypeman.Push, opts []option.RequestOption, update func(*hypeman.Push)) (*hypeman.Push, error) {
+	current := push
+	for {
+		update(current)
 
 		switch current.Status {
 		case hypeman.PushStatusPushed:
-			if renderer != nil {
-				renderer.finish()
-			}
 			return current, nil
 		case hypeman.PushStatusFailed:
-			if renderer != nil {
-				renderer.finish()
-			}
 			if current.Error != "" {
 				return nil, fmt.Errorf("push %s failed: %s", push.ID, current.Error)
 			}
@@ -283,6 +287,28 @@ func waitForPush(ctx context.Context, client *hypeman.Client, push *hypeman.Push
 	}
 }
 
+func pushStatusText(push *hypeman.Push, lastBytes *int64) string {
+	switch push.Status {
+	case hypeman.PushStatusQueued:
+		return fmt.Sprintf("queued · %s", push.Target)
+	case hypeman.PushStatusPushing:
+		if push.Bytes > *lastBytes {
+			*lastBytes = push.Bytes
+		}
+		return fmt.Sprintf("pushing %s · %d layers · %s", formatBytes(*lastBytes), push.Layers, push.Target)
+	case hypeman.PushStatusPushed:
+		return fmt.Sprintf("pushed · digest: %s", push.Digest)
+	case hypeman.PushStatusFailed:
+		message := push.Error
+		if message == "" {
+			message = "unknown error"
+		}
+		return "failed · " + message
+	default:
+		return ""
+	}
+}
+
 type pushStatusRenderer struct {
 	output      io.Writer
 	interactive bool