From a123cbe24c89b783d8fe4f7b1da4f6d152cf47da Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Tue, 25 Aug 2026 21:24:18 +1000 Subject: [PATCH] refactor: change IAM updateRoleCAS to use kvutil (#879) --- spinifex/handlers/iam/instance_profiles.go | 51 +++++---------------- spinifex/handlers/iam/roles.go | 52 +++++----------------- spinifex/kvutil/cas.go | 15 +++++++ spinifex/kvutil/cas_conflict_test.go | 27 +++++++++++ 4 files changed, 64 insertions(+), 81 deletions(-) diff --git a/spinifex/handlers/iam/instance_profiles.go b/spinifex/handlers/iam/instance_profiles.go index 38c0ae3e3..57feb7d28 100644 --- a/spinifex/handlers/iam/instance_profiles.go +++ b/spinifex/handlers/iam/instance_profiles.go @@ -373,47 +373,16 @@ func (s *IAMServiceImpl) getInstanceProfile(ctx context.Context, accountID, prof // re-reading and re-running mutate when a concurrent writer wins the race. // mutate reports whether it changed the record; a false return commits nothing. func (s *IAMServiceImpl) updateInstanceProfileCAS(ctx context.Context, accountID, profileName string, mutate func(*InstanceProfile) (bool, error)) error { - key := accountID + "." + profileName - for range instanceProfileCASMaxRetries { - entry, err := s.instanceProfilesBucket.Get(ctx, key) - if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { - return errors.New(awserrors.ErrorIAMNoSuchEntity) - } - return fmt.Errorf("get instance profile: %w", err) - } - - var profile InstanceProfile - if err := json.Unmarshal(entry.Value(), &profile); err != nil { - return fmt.Errorf("unmarshal instance profile: %w", err) - } - - changed, err := mutate(&profile) - if err != nil { - return err - } - if !changed { - return nil - } - - data, err := json.Marshal(&profile) - if err != nil { - return fmt.Errorf("marshal instance profile: %w", err) - } - // ErrKeyRevisionMismatch is the only conflict signal that holds on both - // replicated and single-replica buckets; ErrKeyExists misses R>1. - if _, err := s.instanceProfilesBucket.Update(ctx, key, data, entry.Revision()); err != nil { - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { - continue // CAS conflict — another writer won, re-read and retry. - } - return fmt.Errorf("update instance profile: %w", err) - } - return nil - } - - slog.Warn("IAM instance profile CAS exhausted retries", - "accountID", accountID, "instanceProfileName", profileName, "attempts", instanceProfileCASMaxRetries) - return errors.New(awserrors.ErrorServerInternal) + _, err := kvutil.Update(ctx, s.instanceProfilesBucket, accountID+"."+profileName, kvutil.CASConfig{ + Attempts: instanceProfileCASMaxRetries, + NotFound: errors.New(awserrors.ErrorIAMNoSuchEntity), + Exhausted: func(string, int) error { + slog.Warn("IAM instance profile CAS exhausted retries", + "accountID", accountID, "instanceProfileName", profileName, "attempts", instanceProfileCASMaxRetries) + return errors.New(awserrors.ErrorServerInternal) + }, + }, mutate) + return err } // profileToSDK converts the internal InstanceProfile to the AWS SDK shape. diff --git a/spinifex/handlers/iam/roles.go b/spinifex/handlers/iam/roles.go index bf9802ddc..77f11cc86 100644 --- a/spinifex/handlers/iam/roles.go +++ b/spinifex/handlers/iam/roles.go @@ -572,46 +572,18 @@ func (s *IAMServiceImpl) getRole(ctx context.Context, accountID, roleName string // to one role at once) write the same record concurrently. mutate reports // whether it changed the record; a false return commits nothing. func (s *IAMServiceImpl) updateRoleCAS(ctx context.Context, accountID, roleName string, mutate func(*Role) (bool, error)) error { - key := accountID + "." + roleName - for range roleCASMaxRetries { - entry, err := s.rolesBucket.Get(ctx, key) - if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { - return errors.New(awserrors.ErrorIAMNoSuchEntity) - } - return fmt.Errorf("get role: %w", err) - } - - var role Role - if err := json.Unmarshal(entry.Value(), &role); err != nil { - return fmt.Errorf("unmarshal role: %w", err) - } - - changed, err := mutate(&role) - if err != nil { - return err - } - if !changed { - return nil - } - - data, err := json.Marshal(&role) - if err != nil { - return fmt.Errorf("marshal role: %w", err) - } - if _, err := s.rolesBucket.Update(ctx, key, data, entry.Revision()); err != nil { - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { - continue // CAS conflict — another writer won, re-read and retry. - } - return fmt.Errorf("update role: %w", err) - } - return nil - } - // Only a CAS conflict reaches here, so the role is contended rather than - // broken; say so, or the caller's InternalError has no cause anywhere. - slog.Error("IAM role CAS retries exhausted under contention", - "accountID", accountID, "roleName", roleName, "attempts", roleCASMaxRetries) - return errors.New(awserrors.ErrorServerInternal) + _, err := kvutil.Update(ctx, s.rolesBucket, accountID+"."+roleName, kvutil.CASConfig{ + Attempts: roleCASMaxRetries, + NotFound: errors.New(awserrors.ErrorIAMNoSuchEntity), + Exhausted: func(string, int) error { + // Only a CAS conflict reaches here, so the role is contended rather than + // broken; say so, or the caller's InternalError has no cause anywhere. + slog.Error("IAM role CAS retries exhausted under contention", + "accountID", accountID, "roleName", roleName, "attempts", roleCASMaxRetries) + return errors.New(awserrors.ErrorServerInternal) + }, + }, mutate) + return err } // findInstanceProfilesForRole scans the instance-profiles bucket for any diff --git a/spinifex/kvutil/cas.go b/spinifex/kvutil/cas.go index 8464f0b41..3b94c0371 100644 --- a/spinifex/kvutil/cas.go +++ b/spinifex/kvutil/cas.go @@ -21,6 +21,14 @@ const casBackoffBase = 2 * time.Millisecond type CASConfig struct { Attempts int // 0 selects casAttempts CreateIfAbsent bool // start from a zero value when the key is missing + + // NotFounc replaces the raw jetstream error when the key is absent and + // CreateIfAbsent is unset, so a caller can map it to its own API error. + NotFound error + + // Exhausted builds the error returned when the attempt budget is spent, for callers that log the contention + // or return a specific API error. + Exhausted func(key string, attempts int) error } // isConflict reports a lost CAS race. Create reports one as ErrKeyExists and Update as ErrKeyRevisionMismatch. @@ -54,6 +62,9 @@ func retryCAS(ctx context.Context, cfg CASConfig, op, key string, attempt func() return ctx.Err() } } + if cfg.Exhausted != nil { + return cfg.Exhausted(key, attempts) + } return fmt.Errorf("CAS %s exhausted %d attempts for key %s: %w", op, attempts, key, lastErr) } @@ -74,6 +85,8 @@ func Update[T any](ctx context.Context, kv jetstream.KeyValue, key string, cfg C revision = entry.Revision() case errors.Is(err, jetstream.ErrKeyNotFound) && cfg.CreateIfAbsent: revision = 0 + case errors.Is(err, jetstream.ErrKeyNotFound) && cfg.NotFound != nil: + return cfg.NotFound default: return err } @@ -118,6 +131,8 @@ func Put(ctx context.Context, kv jetstream.KeyValue, key string, cfg CASConfig, revision = entry.Revision() case errors.Is(err, jetstream.ErrKeyNotFound) && cfg.CreateIfAbsent: revision = 0 + case errors.Is(err, jetstream.ErrKeyNotFound) && cfg.NotFound != nil: + return cfg.NotFound default: return err } diff --git a/spinifex/kvutil/cas_conflict_test.go b/spinifex/kvutil/cas_conflict_test.go index 837f80a5e..a8dc3cf57 100644 --- a/spinifex/kvutil/cas_conflict_test.go +++ b/spinifex/kvutil/cas_conflict_test.go @@ -2,6 +2,7 @@ package kvutil_test import ( "context" + "errors" "fmt" "testing" @@ -152,3 +153,29 @@ func TestCASClaim_RetriesRevisionConflict(t *testing.T) { func jetstreamCodeSuffix(code jetstream.ErrorCode) string { return fmt.Sprintf("%d", code) } + +func TestUpdate_AbsentKeyReturnsCallerError(t *testing.T) { + kv := casTestBucket(t, "cas-notfound") + + sentinel := errors.New("no such entity") + _, err := kvutil.Update(t.Context(), kv, "missing", kvutil.CASConfig{NotFound: sentinel}, + func(*casRecord) (bool, error) { t.Fatal("mutate ran against an absent key"); return false, nil }) + + assert.ErrorIs(t, err, sentinel) +} + +func TestUpdate_ExhaustionUsesCallerError(t *testing.T) { + kv := casTestBucket(t, "cas-exhausted-hook") + _, err := kv.Put(t.Context(), "key", []byte(`{"counter":1}`)) + require.NoError(t, err) + + stub := &casConflictKV{KeyValue: kv, conflictCode: jetstream.JSErrCodeStreamWrongLastSequenceConstant, + updateFailures: casTestAttempts + 1} + sentinel := errors.New("contended") + _, err = kvutil.Update(t.Context(), stub, "key", kvutil.CASConfig{ + Attempts: casTestAttempts, + Exhausted: func(string, int) error { return sentinel }, + }, bumpCounter) + + assert.ErrorIs(t, err, sentinel) +}