refactor: change IAM updateRoleCAS to use kvutil (#879)

This commit is contained in:
Josh Siviour
2026-08-25 21:24:18 +10:00
committed by GitHub
parent 602247af80
commit a123cbe24c
4 changed files with 64 additions and 81 deletions
+10 -41
View File
@@ -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.
+12 -40
View File
@@ -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
+15
View File
@@ -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
}
+27
View File
@@ -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)
}