Fix Workflow state store contention (dapr#8767)

Running Workflows in high throughput scenarios would see contention on processing or blocking Workflow state store operations.

Degraded performance of Workflow operations, leading to increased latency and potential timeouts.

Circular Placement locking of Actor Workflow state operations.

Remove Placement locking in Workflow state operations as it is already locked at the Actor level.

Signed-off-by: joshvanl <me@joshvanl.dev>
Co-authored-by: Cassie Coyle <cassie@diagrid.io>
Co-authored-by: Dapr Bot <56698301+dapr-bot@users.noreply.github.com>

Author: 3 people

pkg/actors/state/fake/fake.go

@@ -20,48 +20,48 @@ import (
 )
 
 type Fake struct {
-	getFn                         func(ctx context.Context, req *api.GetStateRequest) (*api.StateResponse, error)
-	getBulkFn                     func(ctx context.Context, req *api.GetBulkStateRequest) (api.BulkStateResponse, error)
-	transactionalStateOperationFn func(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest) error
+	getFn                         func(ctx context.Context, req *api.GetStateRequest, lock bool) (*api.StateResponse, error)
+	getBulkFn                     func(ctx context.Context, req *api.GetBulkStateRequest, lock bool) (api.BulkStateResponse, error)
+	transactionalStateOperationFn func(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest, lock bool) error
 }
 
 func New() *Fake {
 	return &Fake{
-		getFn: func(ctx context.Context, req *api.GetStateRequest) (*api.StateResponse, error) {
+		getFn: func(ctx context.Context, req *api.GetStateRequest, lock bool) (*api.StateResponse, error) {
 			return nil, nil
 		},
-		getBulkFn: func(ctx context.Context, req *api.GetBulkStateRequest) (api.BulkStateResponse, error) {
+		getBulkFn: func(ctx context.Context, req *api.GetBulkStateRequest, lock bool) (api.BulkStateResponse, error) {
 			return api.BulkStateResponse{}, nil
 		},
-		transactionalStateOperationFn: func(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest) error {
+		transactionalStateOperationFn: func(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest, lock bool) error {
 			return nil
 		},
 	}
 }
 
-func (f *Fake) WithGetFn(fn func(ctx context.Context, req *api.GetStateRequest) (*api.StateResponse, error)) *Fake {
+func (f *Fake) WithGetFn(fn func(ctx context.Context, req *api.GetStateRequest, lock bool) (*api.StateResponse, error)) *Fake {
 	f.getFn = fn
 	return f
 }
 
-func (f *Fake) WithGetBulkFn(fn func(ctx context.Context, req *api.GetBulkStateRequest) (api.BulkStateResponse, error)) *Fake {
+func (f *Fake) WithGetBulkFn(fn func(ctx context.Context, req *api.GetBulkStateRequest, lock bool) (api.BulkStateResponse, error)) *Fake {
 	f.getBulkFn = fn
 	return f
 }
 
-func (f *Fake) WithTransactionalStateOperationFn(fn func(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest) error) *Fake {
+func (f *Fake) WithTransactionalStateOperationFn(fn func(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest, lock bool) error) *Fake {
 	f.transactionalStateOperationFn = fn
 	return f
 }
 
-func (f *Fake) Get(ctx context.Context, req *api.GetStateRequest) (*api.StateResponse, error) {
-	return f.getFn(ctx, req)
+func (f *Fake) Get(ctx context.Context, req *api.GetStateRequest, lock bool) (*api.StateResponse, error) {
+	return f.getFn(ctx, req, lock)
 }
 
-func (f *Fake) GetBulk(ctx context.Context, req *api.GetBulkStateRequest) (api.BulkStateResponse, error) {
-	return f.getBulkFn(ctx, req)
+func (f *Fake) GetBulk(ctx context.Context, req *api.GetBulkStateRequest, lock bool) (api.BulkStateResponse, error) {
+	return f.getBulkFn(ctx, req, lock)
 }
 
-func (f *Fake) TransactionalStateOperation(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest) error {
-	return f.transactionalStateOperationFn(ctx, ignoreHosted, req)
+func (f *Fake) TransactionalStateOperation(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest, lock bool) error {
+	return f.transactionalStateOperationFn(ctx, ignoreHosted, req, lock)
 }

pkg/actors/state/state.go

@@ -40,13 +40,13 @@ var ErrTransactionsTooManyOperations = errors.New("the transaction contains more
 
 type Interface interface {
 	// Get retrieves actor state.
-	Get(ctx context.Context, req *api.GetStateRequest) (*api.StateResponse, error)
+	Get(ctx context.Context, req *api.GetStateRequest, lock bool) (*api.StateResponse, error)
 
 	// GetBulk retrieves actor state in bulk.
-	GetBulk(ctx context.Context, req *api.GetBulkStateRequest) (api.BulkStateResponse, error)
+	GetBulk(ctx context.Context, req *api.GetBulkStateRequest, lock bool) (api.BulkStateResponse, error)
 
 	// TransactionalStateOperation performs a transactional state operation with the actor state store.
-	TransactionalStateOperation(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest) error
+	TransactionalStateOperation(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest, lock bool) error
 }
 
 type Backend interface {
@@ -90,12 +90,16 @@ func New(opts Options) Interface {
 	}
 }
 
-func (s *state) Get(ctx context.Context, req *api.GetStateRequest) (*api.StateResponse, error) {
-	ctx, cancel, err := s.placement.Lock(ctx)
-	if err != nil {
-		return nil, err
+func (s *state) Get(ctx context.Context, req *api.GetStateRequest, lock bool) (*api.StateResponse, error) {
+	if lock {
+		var cancel context.CancelFunc
+		var err error
+		ctx, cancel, err = s.placement.Lock(ctx)
+		if err != nil {
+			return nil, err
+		}
+		defer cancel()
 	}
-	defer cancel()
 
 	storeName, store, err := s.stateStore()
 	if err != nil {
@@ -133,12 +137,16 @@ func (s *state) Get(ctx context.Context, req *api.GetStateRequest) (*api.StateRe
 	}, nil
 }
 
-func (s *state) GetBulk(ctx context.Context, req *api.GetBulkStateRequest) (api.BulkStateResponse, error) {
-	ctx, cancel, err := s.placement.Lock(ctx)
-	if err != nil {
-		return nil, err
+func (s *state) GetBulk(ctx context.Context, req *api.GetBulkStateRequest, lock bool) (api.BulkStateResponse, error) {
+	if lock {
+		var cancel context.CancelFunc
+		var err error
+		ctx, cancel, err = s.placement.Lock(ctx)
+		if err != nil {
+			return nil, err
+		}
+		defer cancel()
 	}
-	defer cancel()
 
 	storeName, store, err := s.stateStore()
 	if err != nil {
@@ -183,12 +191,16 @@ func (s *state) GetBulk(ctx context.Context, req *api.GetBulkStateRequest) (api.
 	return bulkRes, nil
 }
 
-func (s *state) TransactionalStateOperation(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest) error {
-	ctx, cancel, err := s.placement.Lock(ctx)
-	if err != nil {
-		return err
+func (s *state) TransactionalStateOperation(ctx context.Context, ignoreHosted bool, req *api.TransactionalRequest, lock bool) error {
+	if lock {
+		var cancel context.CancelFunc
+		var err error
+		ctx, cancel, err = s.placement.Lock(ctx)
+		if err != nil {
+			return err
+		}
+		defer cancel()
 	}
-	defer cancel()
 
 	if !ignoreHosted {
 		if _, ok := s.table.HostedTarget(req.ActorType, req.ActorID); !ok {
@@ -200,6 +212,7 @@ func (s *state) TransactionalStateOperation(ctx context.Context, ignoreHosted bo
 	baseKey := key.ConstructComposite(s.appID, req.ActorKey())
 	metadata := map[string]string{metadataPartitionKey: baseKey}
 	baseKey += api.DaprSeparator
+	var err error
 	for i, o := range req.Operations {
 		operations[i], err = o.StateOperation(baseKey, api.StateOperationOpts{
 			Metadata: metadata,

pkg/actors/targets/workflow/state.go

@@ -65,7 +65,7 @@ func (w *workflow) saveInternalState(ctx context.Context, state *wfenginestate.S
 
 	log.Debugf("Workflow actor '%s': saving %d keys to actor state store", w.actorID, len(req.Operations))
 
-	if err = w.actorState.TransactionalStateOperation(ctx, true, req); err != nil {
+	if err = w.actorState.TransactionalStateOperation(ctx, true, req, false); err != nil {
 		return err
 	}
 
@@ -97,7 +97,7 @@ func (w *workflow) cleanupWorkflowStateInternal(ctx context.Context, state *wfen
 	}
 
 	// This will do the purging
-	err = w.actorState.TransactionalStateOperation(ctx, true, req)
+	err = w.actorState.TransactionalStateOperation(ctx, true, req, false)
 	if err != nil {
 		return err
 	}

pkg/api/grpc/actor_test.go

@@ -35,7 +35,7 @@ func TestGetActorState(t *testing.T) {
 
 		actors := actorsfake.New()
 		actors.WithState(func(context.Context) (state.Interface, error) {
-			return statefake.New().WithGetFn(func(ctx context.Context, req *actorsapi.GetStateRequest) (*actorsapi.StateResponse, error) {
+			return statefake.New().WithGetFn(func(ctx context.Context, req *actorsapi.GetStateRequest, _ bool) (*actorsapi.StateResponse, error) {
 				return &actorsapi.StateResponse{
 					Data: data,
 					Metadata: map[string]string{
@@ -80,7 +80,7 @@ func TestExecuteActorStateTransaction(t *testing.T) {
 
 		actors := actorsfake.New()
 		actors.WithState(func(context.Context) (state.Interface, error) {
-			return statefake.New().WithTransactionalStateOperationFn(func(ctx context.Context, _ bool, req *actorsapi.TransactionalRequest) error {
+			return statefake.New().WithTransactionalStateOperationFn(func(ctx context.Context, _ bool, req *actorsapi.TransactionalRequest, _ bool) error {
 				return nil
 			}), nil
 		})

pkg/api/grpc/grpc.go

@@ -1071,7 +1071,7 @@ func (a *api) GetActorState(ctx context.Context, in *runtimev1pb.GetActorStateRe
 		Key:       key,
 	}
 
-	resp, err := astate.Get(ctx, &req)
+	resp, err := astate.Get(ctx, &req, true)
 	if err != nil {
 		if _, ok := status.FromError(err); ok {
 			apiServerLogger.Debug(err)
@@ -1143,7 +1143,7 @@ func (a *api) ExecuteActorStateTransaction(ctx context.Context, in *runtimev1pb.
 		Operations: actorOps,
 	}
 
-	err = astate.TransactionalStateOperation(ctx, false, &req)
+	err = astate.TransactionalStateOperation(ctx, false, &req, true)
 	if err != nil {
 		if _, ok := status.FromError(err); ok {
 			apiServerLogger.Debug(err)

pkg/api/http/actors.go

@@ -284,7 +284,7 @@ func (a *api) onActorStateTransaction(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	err = state.TransactionalStateOperation(ctx, false, req)
+	err = state.TransactionalStateOperation(ctx, false, req, true)
 	if err != nil {
 		if errors.As(err, new(messages.APIError)) {
 			respondWithError(w, err)
@@ -449,7 +449,7 @@ func (a *api) onGetActorState(w http.ResponseWriter, r *http.Request) {
 		ActorType: actorType,
 		ActorID:   actorID,
 		Key:       key,
-	})
+	}, true)
 	if err != nil {
 		if errors.As(err, new(messages.APIError)) {
 			respondWithError(w, err)

pkg/api/http/http_test.go

@@ -1081,7 +1081,7 @@ func TestV1ActorEndpoints(t *testing.T) {
 
 	t.Run("Get actor state - 200 OK", func(t *testing.T) {
 		actors.WithState(func(context.Context) (actorsstate.Interface, error) {
-			return statefake.New().WithGetFn(func(context.Context, *actorsapi.GetStateRequest) (*actorsapi.StateResponse, error) {
+			return statefake.New().WithGetFn(func(context.Context, *actorsapi.GetStateRequest, bool) (*actorsapi.StateResponse, error) {
 				return &actorsapi.StateResponse{
 					Data: fakeData,
 					Metadata: map[string]string{
@@ -1104,7 +1104,7 @@ func TestV1ActorEndpoints(t *testing.T) {
 
 	t.Run("Get actor state - 204 No Content", func(t *testing.T) {
 		actors.WithState(func(context.Context) (actorsstate.Interface, error) {
-			return statefake.New().WithGetFn(func(context.Context, *actorsapi.GetStateRequest) (*actorsapi.StateResponse, error) {
+			return statefake.New().WithGetFn(func(context.Context, *actorsapi.GetStateRequest, bool) (*actorsapi.StateResponse, error) {
 				return nil, nil
 			}), nil
 		})
@@ -1122,7 +1122,7 @@ func TestV1ActorEndpoints(t *testing.T) {
 
 	t.Run("Get actor state - 500 on GetState failure", func(t *testing.T) {
 		actors.WithState(func(context.Context) (actorsstate.Interface, error) {
-			return statefake.New().WithGetFn(func(context.Context, *actorsapi.GetStateRequest) (*actorsapi.StateResponse, error) {
+			return statefake.New().WithGetFn(func(context.Context, *actorsapi.GetStateRequest, bool) (*actorsapi.StateResponse, error) {
 				return nil, errors.New("UPSTREAM_ERROR")
 			}), nil
 		})
@@ -1243,7 +1243,7 @@ func TestV1ActorEndpoints(t *testing.T) {
 		}
 
 		actors.WithState(func(context.Context) (actorsstate.Interface, error) {
-			return statefake.New().WithTransactionalStateOperationFn(func(context.Context, bool, *actorsapi.TransactionalRequest) error {
+			return statefake.New().WithTransactionalStateOperationFn(func(context.Context, bool, *actorsapi.TransactionalRequest, bool) error {
 				return errors.New("UPSTREAM_ERROR")
 			}), nil
 		})
@@ -1735,7 +1735,7 @@ func TestV1ActorEndpointsWithTracer(t *testing.T) {
 		buffer = ""
 		apiPath := "v1.0/actors/fakeActorType/fakeActorID/state/key1"
 		actors.WithState(func(context.Context) (actorsstate.Interface, error) {
-			return astate.WithGetFn(func(ctx context.Context, req *actorsapi.GetStateRequest) (*actorsapi.StateResponse, error) {
+			return astate.WithGetFn(func(ctx context.Context, req *actorsapi.GetStateRequest, _ bool) (*actorsapi.StateResponse, error) {
 				called.Add(1)
 				return &actorsapi.StateResponse{
 					Data: fakeData,

pkg/runtime/wfengine/state/state.go

@@ -264,7 +264,7 @@ func LoadWorkflowState(ctx context.Context, state state.Interface, actorID strin
 		ActorID:   actorID,
 		Key:       metadataKey,
 	}
-	res, err := state.Get(ctx, &req)
+	res, err := state.Get(ctx, &req, false)
 	if err != nil {
 		return nil, fmt.Errorf("failed to load workflow metadata: %w", err)
 	}
@@ -312,7 +312,7 @@ func LoadWorkflowState(ctx context.Context, state state.Interface, actorID strin
 	}
 
 	// Perform the request
-	bulkRes, err := state.GetBulk(ctx, bulkReq)
+	bulkRes, err := state.GetBulk(ctx, bulkReq, false)
 	if err != nil {
 		return nil, fmt.Errorf("failed to load workflow state: %w", err)
 	}