From f1c626cf235c2389a2af39435ed38816a6b78bfe Mon Sep 17 00:00:00 2001
From: Jacob Gillespie <jacobwgillespie@gmail.com>
Date: Thu, 21 Nov 2024 18:29:21 +0000
Subject: [PATCH] Retry `depot push` on 5xx errors

---
 pkg/cmd/push/blobs.go | 62 +++++++++++++++++++++++++++----------------
 pkg/cmd/push/push.go  |  4 +--
 2 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/pkg/cmd/push/blobs.go b/pkg/cmd/push/blobs.go
index 6ae88732..ab56e7fb 100644
--- a/pkg/cmd/push/blobs.go
+++ b/pkg/cmd/push/blobs.go
@@ -7,11 +7,12 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/opencontainers/go-digest"
 )
 
-type BlobRequest struct {
+type BlobToPush struct {
 	ParsedTag     *ParsedTag
 	RegistryToken *Token
 	BuildID       string
@@ -19,37 +20,52 @@ type BlobRequest struct {
 }
 
 // PushBlob requests a blob to be pushed from Depot to a destination registry.
-func PushBlob(ctx context.Context, depotToken string, req *BlobRequest) error {
+func PushBlob(ctx context.Context, depotToken string, blob *BlobToPush) error {
+	var err error
+	var req *http.Request
+
 	pushRequest := struct {
 		RegistryHost        string `json:"registryHost"`
 		RepositoryNamespace string `json:"repositoryNamespace"`
 		RegistryToken       string `json:"registryToken"`
 		TokenScheme         string `json:"tokenScheme"`
 	}{
-		RegistryHost:        req.ParsedTag.Host,
-		RepositoryNamespace: req.ParsedTag.Path,
-		RegistryToken:       req.RegistryToken.Token,
-		TokenScheme:         req.RegistryToken.Scheme,
+		RegistryHost:        blob.ParsedTag.Host,
+		RepositoryNamespace: blob.ParsedTag.Path,
+		RegistryToken:       blob.RegistryToken.Token,
+		TokenScheme:         blob.RegistryToken.Scheme,
 	}
 	buf, _ := json.MarshalIndent(pushRequest, "", "  ")
+	url := fmt.Sprintf("https://blob.depot.dev/blobs/%s/%s", blob.BuildID, blob.Digest.String())
 
-	url := fmt.Sprintf("https://blob.depot.dev/blobs/%s/%s", req.BuildID, req.Digest.String())
-	pushReq, err := http.NewRequestWithContext(ctx, "POST", url, strings.NewReader(string(buf)))
-	if err != nil {
-		return err
-	}
-	pushReq.Header.Add("Authorization", "Bearer "+depotToken)
-	pushReq.Header.Set("Content-Type", "application/json")
-	resp, err := http.DefaultClient.Do(pushReq)
-	if err != nil {
-		return err
-	}
-	_ = resp.Body.Close()
-	if resp.StatusCode/100 != 2 {
-		body, _ := io.ReadAll(resp.Body)
-		err := fmt.Errorf("unexpected status code: %d %s", resp.StatusCode, string(body))
-		return err
+	attempts := 0
+	for {
+		attempts += 1
+
+		req, err = http.NewRequestWithContext(ctx, "POST", url, strings.NewReader(string(buf)))
+		if err != nil {
+			return err
+		}
+		req.Header.Add("Authorization", "Bearer "+depotToken)
+		req.Header.Set("Content-Type", "application/json")
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			return err
+		}
+		_ = resp.Body.Close()
+
+		if resp.StatusCode/100 != 2 {
+			if resp.StatusCode >= 500 && attempts < 3 {
+				time.Sleep(5 * time.Second)
+				continue
+			}
+
+			body, _ := io.ReadAll(resp.Body)
+			err := fmt.Errorf("unexpected status code: %d %s", resp.StatusCode, string(body))
+			return err
+		}
+
+		return nil
 	}
 
-	return nil
 }
diff --git a/pkg/cmd/push/push.go b/pkg/cmd/push/push.go
index d7c9c6d9..7354d630 100644
--- a/pkg/cmd/push/push.go
+++ b/pkg/cmd/push/push.go
@@ -167,13 +167,13 @@ func Push(ctx context.Context, progressFmt, buildID, target, tag, token string,
 			blob := blobs[i]
 			fin := logger(fmt.Sprintf("Pushing blob %s", blob.Digest.String()))
 
-			req := &BlobRequest{
+			blobToPush := &BlobToPush{
 				ParsedTag:     parsedTag,
 				RegistryToken: registryToken,
 				BuildID:       buildID,
 				Digest:        blob.Digest,
 			}
-			err := PushBlob(blobCtx, token, req)
+			err := PushBlob(blobCtx, token, blobToPush)
 			fin()
 			return err
 		})