feat: honour Retry-After for artifact retries

internal/artifact/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "downloader.go",
         "gs_downloader.go",
         "gs_uploader.go",
+        "retry_after.go",
         "s3.go",
         "s3_downloader.go",
         "s3_uploader.go",

internal/artifact/batch_creator.go

@@ -99,6 +99,9 @@ func (a *BatchCreator) Create(ctx context.Context) ([]*api.Artifact, error) {
 			if resp != nil && resp.StatusCode == 429 {
 				saw429 = true
 			}
+			if applyRetryAfterHeader(resp, r) {
+				a.logger.Debug("CreateArtifacts retry interval updated from Retry-After header")
+			}
 			// the server returns a 403 code if the artifact has exceeded the service quota
 			// Break the retry on any 4xx code except for 429 Too Many Requests.
 			if resp != nil && (resp.StatusCode != 429 && resp.StatusCode >= 400 && resp.StatusCode <= 499) {

internal/artifact/batching_test.go

@@ -6,12 +6,21 @@ import (
 	"fmt"
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/buildkite/agent/v3/api"
 	"github.com/buildkite/agent/v3/logger"
 	"github.com/stretchr/testify/require"
 )
 
+type nextIntervalRecorder struct {
+	next time.Duration
+}
+
+func (n *nextIntervalRecorder) SetNextInterval(d time.Duration) {
+	n.next = d
+}
+
 type stubArtifactAPIClient struct {
 	createFn func(context.Context, string, *api.ArtifactBatch) (*api.ArtifactBatchCreateResponse, *api.Response, error)
 	updateFn func(context.Context, string, []api.ArtifactState) (*api.Response, error)
@@ -117,3 +126,25 @@ func TestUpdateStatesRespectsConfiguredBatchMax(t *testing.T) {
 		require.Equal(t, "sent", tracker.State)
 	}
 }
+
+func TestApplyRetryAfterHeaderSetsRetryInterval(t *testing.T) {
+	t.Parallel()
+
+	r := &nextIntervalRecorder{}
+	resp := &api.Response{Response: &http.Response{Header: http.Header{"Retry-After": []string{"3"}}}}
+
+	updated := applyRetryAfterHeader(resp, r)
+	require.True(t, updated)
+	require.Equal(t, 3*time.Second, r.next)
+}
+
+func TestApplyRetryAfterHeaderIgnoresInvalidValues(t *testing.T) {
+	t.Parallel()
+
+	r := &nextIntervalRecorder{}
+	resp := &api.Response{Response: &http.Response{Header: http.Header{"Retry-After": []string{"not-a-number"}}}}
+
+	updated := applyRetryAfterHeader(resp, r)
+	require.False(t, updated)
+	require.Equal(t, time.Duration(0), r.next)
+}

internal/artifact/retry_after.go

@@ -0,0 +1,30 @@
+package artifact
+
+import (
+	"time"
+
+	"github.com/buildkite/agent/v3/api"
+)
+
+type retryIntervalSetter interface {
+	SetNextInterval(time.Duration)
+}
+
+func applyRetryAfterHeader(resp *api.Response, r retryIntervalSetter) bool {
+	if resp == nil {
+		return false
+	}
+
+	retryAfter := resp.Header.Get("Retry-After")
+	if retryAfter == "" {
+		return false
+	}
+
+	duration, err := time.ParseDuration(retryAfter + "s")
+	if err != nil {
+		return false
+	}
+
+	r.SetNextInterval(duration)
+	return true
+}

internal/artifact/uploader.go

@@ -773,7 +773,10 @@ func (a *artifactUploadWorker) updateStates(ctx context.Context) error {
 				defer cancel()
 			}
 
-			_, err := a.apiClient.UpdateArtifacts(ctxTimeout, a.conf.JobID, chunk)
+			resp, err := a.apiClient.UpdateArtifacts(ctxTimeout, a.conf.JobID, chunk)
+			if applyRetryAfterHeader(resp, r) {
+				a.logger.Debug("UpdateArtifacts retry interval updated from Retry-After header")
+			}
 			if err != nil {
 				a.logger.Warn("%s (%s)", err, r)
 			}