feat: add backoff retry to all attestation commands (#1537)

Signed-off-by: Miguel Martinez <[email protected]>

app/cli/cmd/attestation_add.go

@@ -16,18 +16,14 @@
 package cmd
 
 import (
-	"errors"
 	"fmt"
 	"os"
-	"time"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 	"google.golang.org/grpc"
 
-	"github.com/cenkalti/backoff/v4"
 	"github.com/chainloop-dev/chainloop/app/cli/internal/action"
-	v1 "github.com/chainloop-dev/chainloop/app/controlplane/api/controlplane/v1"
 	schemaapi "github.com/chainloop-dev/chainloop/app/controlplane/api/workflowcontract/v1"
 )
 
@@ -90,30 +86,17 @@ func newAttestationAddCmd() *cobra.Command {
 
 			// In some cases, the attestation state is stored remotely. To control concurrency we use
 			// optimistic locking. We retry the operation if the state has changed since we last read it.
-			return backoff.RetryNotify(
+			return runWithBackoffRetry(
 				func() error {
 					if err := a.Run(cmd.Context(), attestationID, name, value, kind, annotations); err != nil {
-						if errors.Is(err, action.ErrAttestationNotInitialized) {
-							return err
-						}
-
-						// We want to retry if the error is a conflict
-						if v1.IsAttestationStateErrorConflict(err) {
-							return err
-						}
-
-						// if it's another kind of error we want to stop retrying
-						return backoff.Permanent(newGracefulError(err))
+						return err
 					}
 
 					logger.Info().Msg("material added to attestation")
 
 					return nil
 				},
-				backoff.NewExponentialBackOff(backoff.WithMaxElapsedTime(3*time.Minute)),
-				func(err error, delay time.Duration) {
-					logger.Err(err).Msgf("retrying in %s", delay)
-				})
+			)
 		},
 
 		PostRunE: func(cmd *cobra.Command, args []string) error {

app/cli/cmd/attestation_init.go

@@ -80,16 +80,22 @@ func newAttestationInitCmd() *cobra.Command {
 				return fmt.Errorf("failed to initialize attestation: %w", err)
 			}
 
-			// Initialize it
-			attestationID, err := a.Run(cmd.Context(), &action.AttestationInitRunOpts{
-				ContractRevision:             contractRevision,
-				ProjectName:                  projectName,
-				ProjectVersion:               projectVersion,
-				WorkflowName:                 workflowName,
-				NewWorkflowContractName:      newWorkflowcontractName,
-				ProjectVersionMarkAsReleased: projectVersionRelease,
-			})
+			var attestationID string
+			err = runWithBackoffRetry(
+				func() error {
+					// Initialize it
+					attestationID, err = a.Run(cmd.Context(), &action.AttestationInitRunOpts{
+						ContractRevision:             contractRevision,
+						ProjectName:                  projectName,
+						ProjectVersion:               projectVersion,
+						WorkflowName:                 workflowName,
+						NewWorkflowContractName:      newWorkflowcontractName,
+						ProjectVersionMarkAsReleased: projectVersionRelease,
+					})
 
+					return err
+				},
+			)
 			if err != nil {
 				if errors.Is(err, action.ErrAttestationAlreadyExist) {
 					return err

app/cli/cmd/attestation_push.go

@@ -76,7 +76,14 @@ func newAttestationPushCmd() *cobra.Command {
 				return err
 			}
 
-			res, err := a.Run(cmd.Context(), attestationID, annotations)
+			var res *action.AttestationResult
+			err = runWithBackoffRetry(
+				func() error {
+					res, err = a.Run(cmd.Context(), attestationID, annotations)
+					return err
+				},
+			)
+
 			if err != nil {
 				if errors.Is(err, action.ErrAttestationNotInitialized) {
 					return err

app/cli/cmd/errors.go

@@ -17,6 +17,12 @@ package cmd
 
 import (
 	"errors"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	v1 "github.com/chainloop-dev/chainloop/app/controlplane/api/controlplane/v1"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // GracefulError represents an error that has been marked as gracefully handled
@@ -42,3 +48,45 @@ var ErrAttestationNotInitialized = errors.New("attestation not yet initialized,
 var ErrAttestationAlreadyExist = errors.New("attestation already initialized, to override it use the --replace flag`")
 var ErrAttestationTokenRequired = errors.New("chainloop Token required, please provide it via --token flag or CHAINLOOP_TOKEN environment variable")
 var ErrKeylessNotSupported = errors.New("keyless signing not supported, please provide a private key reference with --key instead")
+
+func isRetriableAPIError(err error) bool {
+	// we retry state conflicts and other transient errors
+	if v1.IsAttestationStateErrorConflict(err) {
+		return true
+	}
+
+	st, ok := status.FromError(err)
+	if !ok {
+		return false
+	}
+
+	retriableCodes := []codes.Code{
+		codes.Unavailable,
+		codes.Internal,
+		codes.ResourceExhausted,
+		codes.DeadlineExceeded,
+	}
+
+	for _, code := range retriableCodes {
+		if st.Code() == code {
+			return true
+		}
+	}
+
+	return false
+}
+
+func runWithBackoffRetry(fn func() error) error {
+	return backoff.RetryNotify(
+		func() error {
+			err := fn()
+			if !isRetriableAPIError(err) {
+				return backoff.Permanent(err)
+			}
+			return err
+		},
+		backoff.NewExponentialBackOff(backoff.WithMaxElapsedTime(3*time.Minute)),
+		func(err error, delay time.Duration) {
+			logger.Err(err).Msgf("retrying in %s", delay)
+		})
+}