Use CompareAndSwap for OktaAssignments instead of lock. (#24727)

* Use CompareAndSwap for OktaAssignments instead of lock.

CompareAndSwap should be used for the OktaAssignment updates instead of a
lock, as it's lighter wait and fits the intended problem better.

* Use currentAssignment instead of backendAssignment.

* Update the last transition time.

* Set Metadata during update.

api/types/okta.go

@@ -177,6 +177,8 @@ func (o *OktaImportRuleMatchV1) CheckAndSetDefaults() error {
 type OktaAssignment interface {
 	ResourceWithLabels
 
+	// SetMetadata will set the metadata for the Okta assignment.
+	SetMetadata(metadata Metadata)
 	// GetUser will return the user that the Okta assignment actions applies to.
 	GetUser() string
 	// GetActions will return the list of actions that will be performed as part of this assignment.
@@ -201,6 +203,11 @@ func NewOktaAssignment(metadata Metadata, spec OktaAssignmentSpecV1) (OktaAssign
 	return o, nil
 }
 
+// SetMetadata will set the metadata for the Okta assignment.
+func (o *OktaAssignmentV1) SetMetadata(metadata Metadata) {
+	o.Metadata = metadata
+}
+
 // GetUser returns the user that the actions will be applied to.
 func (o *OktaAssignmentV1) GetUser() string {
 	return o.Spec.User

lib/services/local/okta.go

@@ -18,7 +18,6 @@ package local
 
 import (
 	"context"
-	"time"
 
 	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
@@ -34,9 +33,6 @@ const (
 	oktaImportRuleMaxPageSize = 200
 	oktaAssignmentPrefix      = "okta_assignment"
 	oktaAssignmentMaxPageSize = 200
-
-	oktaAssignmentModifyLockTTL = 10 * time.Second
-	oktaAssignmentModifyLock    = "okta_assignment_modify_lock"
 )
 
 // OktaService manages Okta resources in the Backend.
@@ -116,17 +112,7 @@ func (o *OktaService) ListOktaAssignments(ctx context.Context, pageSize int, nex
 
 // GetOktaAssignment returns the specified Okta assignment resources.
 func (o *OktaService) GetOktaAssignment(ctx context.Context, name string) (types.OktaAssignment, error) {
-	var assignment types.OktaAssignment
-	err := o.assignmentSvc.RunWhileLocked(ctx, oktaAssignmentModifyLock, oktaAssignmentModifyLockTTL, func(ctx context.Context, b backend.Backend) error {
-		var err error
-		assignment, err = o.assignmentSvc.GetResource(ctx, name)
-		if err != nil {
-			return trace.Wrap(err)
-		}
-
-		return nil
-	})
-	return assignment, trace.Wrap(err)
+	return o.assignmentSvc.GetResource(ctx, name)
 }
 
 // CreateOktaAssignment creates a new Okta assignment resource.
@@ -137,42 +123,38 @@ func (o *OktaService) CreateOktaAssignment(ctx context.Context, assignment types
 // UpdateOktaAssignment updates an existing Okta assignment resource.
 func (o *OktaService) UpdateOktaAssignment(ctx context.Context, assignment types.OktaAssignment) (types.OktaAssignment, error) {
 	var previousAssignment types.OktaAssignment
-	err := o.assignmentSvc.RunWhileLocked(ctx, oktaAssignmentModifyLock, oktaAssignmentModifyLockTTL, func(ctx context.Context, b backend.Backend) error {
-		var err error
-		previousAssignment, err = o.assignmentSvc.GetResource(ctx, assignment.GetName())
-		if err != nil {
-			return trace.Wrap(err)
-		}
+	err := o.assignmentSvc.UpdateAndSwapResource(ctx, assignment.GetName(), func(currentAssignment types.OktaAssignment) error {
+		previousAssignment = currentAssignment.Copy()
+		currentActions := currentAssignment.GetActions()
 
-		// Make a copy so that we don't modify the previous assignment.
-		previousCopy := previousAssignment.Copy()
-		previousActions := previousCopy.GetActions()
-
-		if len(previousActions) != len(assignment.GetActions()) {
+		if len(currentActions) != len(assignment.GetActions()) {
 			return trace.BadParameter("Update to Okta assignment %s failed because the previous version has a different number of actions", assignment.GetName())
 		}
 
 		// Make sure that the status transitions of the updated assignment are valid.
 		for i, action := range assignment.GetActions() {
-			previousAction := previousActions[i]
+			currentAction := currentActions[i]
 
 			// Ensure that the previous actions are equal
-			if !actionsMatch(previousAction, action) {
+			if !actionsMatch(currentAction, action) {
 				return trace.BadParameter("action mismatch when updating Okta assignment %s", assignment.GetName())
 			}
 
 			// Don't check the status transition if the statuses are equal and the last transitions are equal.
-			if previousAction.GetStatus() == action.GetStatus() &&
-				previousAction.GetLastTransition().Equal(action.GetLastTransition()) {
+			if currentAction.GetStatus() == action.GetStatus() &&
+				currentAction.GetLastTransition().Equal(action.GetLastTransition()) {
 				continue
 			}
 
-			if err := previousAction.SetStatus(action.GetStatus()); err != nil {
+			if err := currentAction.SetStatus(action.GetStatus()); err != nil {
 				return trace.Wrap(err)
 			}
+			currentAction.SetLastTransition(action.GetLastTransition())
 		}
 
-		return o.assignmentSvc.UpdateResource(ctx, assignment)
+		currentAssignment.SetMetadata(assignment.GetMetadata())
+
+		return nil
 	})
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -185,38 +167,29 @@ func (o *OktaService) UpdateOktaAssignment(ctx context.Context, assignment types
 // status is a valid transition. If a transition is invalid, it will be logged and the rest of the action statuses
 // will be updated if possible.
 func (o *OktaService) UpdateOktaAssignmentActionStatuses(ctx context.Context, name, status string) (types.OktaAssignment, error) {
-	var assignment types.OktaAssignment
-	err := o.assignmentSvc.RunWhileLocked(ctx, oktaAssignmentModifyLock, oktaAssignmentModifyLockTTL, func(ctx context.Context, b backend.Backend) error {
-		var err error
-		assignment, err = o.assignmentSvc.GetResource(ctx, name)
-		if err != nil {
-			return trace.Wrap(err)
-		}
-
+	var previousAssignment types.OktaAssignment
+	err := o.assignmentSvc.UpdateAndSwapResource(ctx, name, func(assignment types.OktaAssignment) error {
+		previousAssignment = assignment.Copy()
 		for _, action := range assignment.GetActions() {
 			if err := action.SetStatus(status); err != nil {
 				o.log.Warnf("Unable to transition status from %s -> %s", action.GetStatus(), status)
 			}
 		}
 
-		return o.assignmentSvc.UpdateResource(ctx, assignment)
+		return nil
 	})
-	return assignment, trace.Wrap(err)
+	return previousAssignment, trace.Wrap(err)
 
 }
 
 // DeleteOktaAssignment removes the specified Okta assignment resource.
 func (o *OktaService) DeleteOktaAssignment(ctx context.Context, name string) error {
-	return trace.Wrap(o.assignmentSvc.RunWhileLocked(ctx, oktaAssignmentModifyLock, oktaAssignmentModifyLockTTL, func(ctx context.Context, b backend.Backend) error {
-		return o.assignmentSvc.DeleteResource(ctx, name)
-	}))
+	return o.assignmentSvc.DeleteResource(ctx, name)
 }
 
 // DeleteAllOktaAssignments removes all Okta assignments.
 func (o *OktaService) DeleteAllOktaAssignments(ctx context.Context) error {
-	return trace.Wrap(o.assignmentSvc.RunWhileLocked(ctx, oktaAssignmentModifyLock, oktaAssignmentModifyLockTTL, func(ctx context.Context, b backend.Backend) error {
-		return o.assignmentSvc.DeleteAllResources(ctx)
-	}))
+	return o.assignmentSvc.DeleteAllResources(ctx)
 }
 
 // actionsMatch returns true if two actions match minus the status and last transition.