Merge pull request #3706 from buildkite/bm/secret_retry

Add retry logic to secret getting

Author: moskyb

clicommand/secret_get.go

@@ -94,7 +94,7 @@ Examples:
 		}
 
 		agentClient := api.NewClient(l, loadAPIClientConfig(cfg, "AgentAccessToken"))
-		secrets, errs := secrets.FetchSecrets(ctx, agentClient, cfg.Job, cfg.Keys, 10)
+		secrets, errs := secrets.FetchSecrets(ctx, l, agentClient, cfg.Job, cfg.Keys, 10)
 		if len(errs) > 0 {
 			sb := &strings.Builder{}
 			sb.WriteString("Failed to fetch some secrets:\n")

internal/job/executor.go

@@ -909,13 +909,14 @@ func (e *Executor) fetchAndSetSecrets(ctx context.Context) error {
 	}
 
 	// Create API client for fetching secrets.
-	apiClient := api.NewClient(logger.NewBuffer(), api.Config{
+	secretLogger := logger.NewBuffer()
+	apiClient := api.NewClient(secretLogger, api.Config{
 		Endpoint: e.shell.Env.GetString("BUILDKITE_AGENT_ENDPOINT", ""),
 		Token:    e.shell.Env.GetString("BUILDKITE_AGENT_ACCESS_TOKEN", ""),
 	})
 
 	// Fetch all secrets
-	fetchedSecrets, errs := secrets.FetchSecrets(ctx, apiClient, e.JobID, keys, 10)
+	fetchedSecrets, errs := secrets.FetchSecrets(ctx, secretLogger, apiClient, e.JobID, keys, 10)
 	if len(errs) > 0 {
 		var errorMsg strings.Builder
 		for _, err := range errs {

internal/secrets/secret.go

@@ -4,8 +4,11 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/buildkite/agent/v3/api"
+	"github.com/buildkite/agent/v3/logger"
+	"github.com/buildkite/roko"
 	"golang.org/x/sync/semaphore"
 )
 
@@ -34,9 +37,31 @@ func (e *SecretError) Unwrap() error {
 	return e.Err
 }
 
-// FetchSecrets retrieves all secret values from the API sequentially.
-// If any secret fails, returns error with details of all failed secrets.
-func FetchSecrets(ctx context.Context, client APIClient, jobID string, keys []string, concurrency int) ([]Secret, []error) {
+// FetchSecretsOpt is a functional option for FetchSecrets.
+type FetchSecretsOpt func(*fetchSecretsConfig)
+
+type fetchSecretsConfig struct {
+	retrySleepFunc func(time.Duration)
+}
+
+// WithRetrySleepFunc overrides the sleep function used between retries.
+// This is primarily useful for unit tests.
+func WithRetrySleepFunc(f func(time.Duration)) FetchSecretsOpt {
+	return func(c *fetchSecretsConfig) {
+		c.retrySleepFunc = f
+	}
+}
+
+// FetchSecrets retrieves all secret values from the API concurrently.
+// Each individual secret fetch is retried up to 3 times with exponential
+// backoff on retryable errors (TLS handshake failures, timeouts, 5xx, 429).
+// If any secret fails after retries, returns error with details of all failed secrets.
+func FetchSecrets(ctx context.Context, l logger.Logger, client APIClient, jobID string, keys []string, concurrency int, opts ...FetchSecretsOpt) ([]Secret, []error) {
+	var cfg fetchSecretsConfig
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+
 	secrets := make([]Secret, 0, len(keys))
 	secretsMu := sync.Mutex{}
 
@@ -55,7 +80,33 @@ func FetchSecrets(ctx context.Context, client APIClient, jobID string, keys []st
 
 		go func() {
 			defer sem.Release(1)
-			apiSecret, _, err := client.GetSecret(ctx, &api.GetSecretRequest{Key: key, JobID: jobID})
+
+			r := roko.NewRetrier(
+				roko.WithMaxAttempts(3),
+				roko.WithStrategy(roko.Exponential(2*time.Second, 0)),
+				roko.WithJitterRange(-1*time.Second, 5*time.Second),
+				roko.WithSleepFunc(cfg.retrySleepFunc),
+			)
+
+			apiSecret, err := roko.DoFunc(ctx, r, func(r *roko.Retrier) (*api.Secret, error) {
+				secret, resp, err := client.GetSecret(ctx, &api.GetSecretRequest{Key: key, JobID: jobID})
+				if err != nil {
+					if resp != nil && api.IsRetryableStatus(resp) {
+						l.Warn("Retrying secret %q fetch after retryable HTTP status %d (%s)", key, resp.StatusCode, r)
+						return nil, err
+					}
+
+					if api.IsRetryableError(err) {
+						l.Warn("Retrying secret %q fetch after retryable error: %v (%s)", key, err, r)
+						return nil, err
+					}
+
+					// Non-retryable error, stop retrying
+					r.Break()
+					return nil, err
+				}
+				return secret, nil
+			})
 			if err != nil {
 				errsMu.Lock()
 				errs = append(errs, &SecretError{

internal/secrets/secret_test.go

@@ -8,8 +8,10 @@ import (
 	"net/http/httptest"
 	"runtime"
 	"strings"
+	"sync/atomic"
 	"syscall"
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
@@ -18,6 +20,8 @@ import (
 	"github.com/buildkite/agent/v3/logger"
 )
 
+var noSleep = WithRetrySleepFunc(func(time.Duration) {})
+
 func TestFetchSecrets_Success(t *testing.T) {
 	t.Parallel()
 
@@ -44,7 +48,7 @@ func TestFetchSecrets_Success(t *testing.T) {
 		Token:    "llamas",
 	})
 
-	secrets, errs := FetchSecrets(t.Context(), apiClient, "test-job-id", []string{"DATABASE_URL", "API_TOKEN"}, 10)
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", []string{"DATABASE_URL", "API_TOKEN"}, 10)
 	if len(errs) > 0 {
 		t.Fatalf("expected no errors, got: %v", errs)
 	}
@@ -73,7 +77,7 @@ func TestFetchSecrets_EmptyKeys(t *testing.T) {
 	t.Cleanup(server.Close)
 
 	apiClient := api.NewClient(logger.Discard, api.Config{Endpoint: server.URL, Token: "llamas"})
-	secrets, errs := FetchSecrets(t.Context(), apiClient, "test-job-id", []string{}, 10)
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", []string{}, 10)
 
 	if len(errs) > 0 {
 		t.Fatalf("expected no errors, got: %v", errs)
@@ -95,7 +99,7 @@ func TestFetchSecrets_NilKeys(t *testing.T) {
 
 	apiClient := api.NewClient(logger.Discard, api.Config{Endpoint: server.URL, Token: "llamas"})
 
-	secrets, errs := FetchSecrets(t.Context(), apiClient, "test-job-id", nil, 10)
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", nil, 10)
 
 	if len(errs) > 0 {
 		t.Fatalf("expected no errors, got: %v", errs)
@@ -133,7 +137,7 @@ func TestFetchSecrets_SomeSecretsFail(t *testing.T) {
 	})
 
 	keys := []string{"DATABASE_URL", "MISSING"}
-	secrets, errs := FetchSecrets(t.Context(), apiClient, "test-job-id", keys, 10)
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", keys, 10)
 
 	if len(errs) != 1 {
 		t.Fatalf("expected 1 errors, got %d: %v", len(errs), errs)
@@ -177,7 +181,7 @@ func TestFetchSecrets_AllSecretsFail(t *testing.T) {
 	})
 
 	keys := []string{"API_TOKEN", "DATABASE_URL"}
-	secrets, errs := FetchSecrets(t.Context(), apiClient, "test-job-id", keys, 10)
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", keys, 10)
 
 	if len(errs) != 2 {
 		t.Fatalf("expected 2 errors, got %d: %v", len(errs), errs)
@@ -218,7 +222,7 @@ func TestFetchSecrets_APIClientError(t *testing.T) {
 	})
 
 	keys := []string{"TEST_SECRET"}
-	secrets, errs := FetchSecrets(t.Context(), apiClient, "test-job-id", keys, 10)
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", keys, 10, noSleep)
 
 	if len(errs) != 1 {
 		t.Fatalf("expected 1 error, got %d: %v", len(errs), errs)
@@ -266,3 +270,78 @@ func TestFetchSecrets_APIClientError(t *testing.T) {
 		t.Errorf("expected connection refused error, got: %v (type: %T)", netErr.Err, netErr.Err)
 	}
 }
+
+func TestFetchSecrets_RetriesOnServerError(t *testing.T) {
+	t.Parallel()
+
+	var attempts atomic.Int32
+
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		n := attempts.Add(1)
+		if n <= 2 {
+			// First two attempts return 502 Bad Gateway (retryable)
+			rw.WriteHeader(http.StatusBadGateway)
+			_, _ = fmt.Fprintf(rw, `{"message": "bad gateway"}`)
+			return
+		}
+		// Third attempt succeeds
+		rw.WriteHeader(http.StatusOK)
+		_, _ = fmt.Fprintf(rw, `{"key": "MY_SECRET", "value": "secret-value"}`)
+	}))
+	t.Cleanup(server.Close)
+
+	apiClient := api.NewClient(logger.Discard, api.Config{
+		Endpoint: server.URL,
+		Token:    "llamas",
+	})
+
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", []string{"MY_SECRET"}, 10, noSleep)
+	if len(errs) > 0 {
+		t.Fatalf("expected no errors after retries, got: %v", errs)
+	}
+
+	if len(secrets) != 1 {
+		t.Fatalf("expected 1 secret, got %d", len(secrets))
+	}
+
+	if secrets[0].Value != "secret-value" {
+		t.Errorf("expected secret value %q, got %q", "secret-value", secrets[0].Value)
+	}
+
+	if got := attempts.Load(); got != 3 {
+		t.Errorf("expected 3 attempts (2 failures + 1 success), got %d", got)
+	}
+}
+
+func TestFetchSecrets_NoRetryOnNonRetryableStatus(t *testing.T) {
+	t.Parallel()
+
+	var attempts atomic.Int32
+
+	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		attempts.Add(1)
+		// 404 is not retryable
+		rw.WriteHeader(http.StatusNotFound)
+		_, _ = fmt.Fprintf(rw, `{"message": "secret not found"}`)
+	}))
+	t.Cleanup(server.Close)
+
+	apiClient := api.NewClient(logger.Discard, api.Config{
+		Endpoint: server.URL,
+		Token:    "llamas",
+	})
+
+	secrets, errs := FetchSecrets(t.Context(), logger.Discard, apiClient, "test-job-id", []string{"MISSING"}, 10, noSleep)
+	if len(errs) != 1 {
+		t.Fatalf("expected 1 error, got %d: %v", len(errs), errs)
+	}
+
+	if secrets != nil {
+		t.Errorf("expected nil secrets, got: %v", secrets)
+	}
+
+	// Should have only attempted once since 404 is not retryable
+	if got := attempts.Load(); got != 1 {
+		t.Errorf("expected 1 attempt (no retries for 404), got %d", got)
+	}
+}