fixes lack of idempotency in WriteToSpiceDB

I started all of this because of flakes that were trivially
to reproduce, and the underlying issue was WriteToSpiceDB was not idempotent.

It was used by both the optimistic and pesimistic locking workflows, and
the lack of idempotency meant the workflows couldn't recover
by themselves from transient errors. When the tests failed
for the wrong reasons, they flaked.

We introduced `failpoints` as a mechanism to simulate panics
similar to hardware or network failures. However, some recent
changes meant the failpoints were no longer unique and
were causing failures _in the wrong places_.

- when the pessimistic wrote an exclusive lock, retrying it due to
  failure wouldn't work, because the precondition prevented it
  (_make sure no one has a lock on this resource_!)
- when the optimistic wrote with CREATE semantics, a
  subsequent retry after a successful write due to failure
  response failure would fail because the tuples already exist.

This commit proposes making WriteToSpiceDB truly idempotent by
introducing idempotency keys in the SpiceDB schema. All writes
will include a relationships that identify the workflow and the hash
of the payload as the idempotency key.

The flow is as follows:
- perform write
- if failure happens, check if idempotency key was written in previous request
- if exists, assume operation was successful
- if it does not, bubble up the error

The cost of the extra ReadRelationships is only paid in the even of a retry
due to a failure.

The tests were written with the assumption that the system would bubble up
errors after recovery. This goes against the expectations of a durable
workflow engine, which embraces idempotency and is expected to retry
on errors, rather than have the client retry, unless those are unrecoverable.

This was all by design: the workflow wouldn't be responsible to retry things, but rather
execute compensatory operations after an operation failed. This meant
the client had to retry those errors, and it turns out troubleshooting
what happened on a transient error is not that trivial for folks
building on top of the spicedb-kubeapi-proxy.

Author: vroldanbet

e2e/proxy_test.go

@@ -793,56 +793,45 @@ var _ = Describe("Proxy", func() {
 				Expect(owners[0].Relationship.Subject.Object.ObjectId).To(Equal("chani"))
 			})
 
-			It("recovers writes when spicedb write succeeds but crashes", func(ctx context.Context) {
+			It("recovers a `create` operation after a SpiceDB write followed by crash, retries and succeeds", func(ctx context.Context) {
 				// paul creates his namespace
 				Expect(CreateNamespace(ctx, paulClient, paulNamespace)).To(Succeed())
 
-				// make spicedb write crash on chani's namespace write
-				failpoints.EnableFailPoint("panicSpiceDBReadResp", 1)
+				// make spicedb write crash on chani's namespace write, it should be retried idempotently and succeed
+				failpoints.EnableFailPoint("panicSpiceDBWriteResp", 1)
 				err := CreateNamespace(ctx, chaniClient, chaniNamespace)
-				Expect(err).ToNot(BeNil())
-				// pessimistic locking reports a conflict, optimistic locking reports already exists
-				Expect(k8serrors.IsConflict(err) || k8serrors.IsAlreadyExists(err)).To(BeTrue())
+				Expect(err).To(Succeed())
 
-				// paul creates chani's namespace so that the namespace exists
-				Expect(CreateNamespace(ctx, paulClient, chaniNamespace)).To(Succeed())
+				// check that chani can get her namespace
+				Expect(GetNamespace(ctx, chaniClient, chaniNamespace)).To(Succeed())
 
-				// check that chani can't get her namespace, indirectly showing
-				// that the spicedb write was rolled back
-				Expect(k8serrors.IsUnauthorized(GetNamespace(ctx, chaniClient, chaniNamespace))).To(BeTrue())
+				// check that paul can't read chani's namespace
+				Expect(k8serrors.IsUnauthorized(GetNamespace(ctx, paulClient, chaniNamespace))).To(BeTrue())
 
 				// confirm the relationship doesn't exist
 				Expect(len(GetAllTuples(ctx, &v1.RelationshipFilter{
 					ResourceType:          "namespace",
 					OptionalResourceId:    chaniNamespace,
 					OptionalRelation:      "creator",
 					OptionalSubjectFilter: &v1.SubjectFilter{SubjectType: "user", OptionalSubjectId: "chani"},
-				}))).To(BeZero())
-
-				// confirm paul can get the namespace
-				Expect(GetNamespace(ctx, paulClient, chaniNamespace)).To(Succeed())
+				}))).To(Equal(1))
 			})
 
-			It("recovers deletes when spicedb write succeeds but crashes", func(ctx context.Context) {
+			It("recovers a `delete` operation after a SpiceDB write followed by crash, retries and succeeds", func(ctx context.Context) {
 				// paul creates his namespace
 				Expect(CreateNamespace(ctx, paulClient, paulNamespace)).To(Succeed())
 				Expect(CreatePod(ctx, paulClient, paulNamespace, paulPod)).To(Succeed())
 
 				// chani can't create the same pod
 				Expect(k8serrors.IsAlreadyExists(CreatePod(ctx, chaniClient, paulNamespace, paulPod))).To(BeTrue())
-				fmt.Println(GetPod(ctx, chaniClient, paulNamespace, paulPod))
-				Expect(k8serrors.IsUnauthorized(GetPod(ctx, chaniClient, paulNamespace, paulPod))).To(BeTrue())
+				// chani isn't authorized to get the pod
+				err := GetPod(ctx, chaniClient, paulNamespace, paulPod)
+				Expect(k8serrors.IsUnauthorized(err)).To(BeTrue())
 
 				// make spicedb write crash on pod delete
-				failpoints.EnableFailPoint("panicSpiceDBReadResp", 1)
-				err := DeletePod(ctx, paulClient, paulNamespace, paulPod)
-				if lockMode == proxyrule.OptimisticLockMode {
-					Expect(err).To(Succeed())
-				} else {
-					Expect(k8serrors.IsConflict(err) || k8serrors.IsUnauthorized(err)).To(BeTrue())
-					// paul sees the request fail, so he tries again:
-					Expect(DeletePod(ctx, paulClient, paulNamespace, paulPod)).To(Succeed())
-				}
+				failpoints.EnableFailPoint("panicSpiceDBWriteResp", 1)
+				err = DeletePod(ctx, paulClient, paulNamespace, paulPod)
+				Expect(err).To(Succeed())
 
 				// chani can now re-create paul's pod and take ownership
 				Expect(CreatePod(ctx, chaniClient, paulNamespace, paulPod)).To(Succeed())

pkg/authz/distributedtx/activity.go

@@ -7,10 +7,12 @@ import (
 	"io"
 	"net/http"
 
+	"github.com/cespare/xxhash/v2"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apiserver/pkg/endpoints/request"
 	"k8s.io/client-go/rest"
+	"k8s.io/klog/v2"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
@@ -38,11 +40,97 @@ type ActivityHandler struct {
 }
 
 // WriteToSpiceDB writes relationships to spicedb and returns any errors.
-func (h *ActivityHandler) WriteToSpiceDB(ctx context.Context, input *v1.WriteRelationshipsRequest) (*v1.WriteRelationshipsResponse, error) {
+func (h *ActivityHandler) WriteToSpiceDB(ctx context.Context, input *v1.WriteRelationshipsRequest, workflowID string) (struct{}, error) {
 	failpoints.FailPoint("panicWriteSpiceDB")
-	out, err := h.PermissionClient.WriteRelationships(ctx, input)
-	failpoints.FailPoint("panicSpiceDBReadResp")
-	return out, err
+	idempotencyKey, err := idempotencyKeyForPayload(input, workflowID)
+	if err != nil {
+		return struct{}{}, fmt.Errorf("failed to create idempotency key for payload: %w", err)
+	}
+	cloned := input
+	cloned.Updates = append(cloned.Updates, &v1.RelationshipUpdate{
+		Operation:    v1.RelationshipUpdate_OPERATION_CREATE,
+		Relationship: idempotencyKey,
+	})
+
+	_, err = h.PermissionClient.WriteRelationships(ctx, input)
+	failpoints.FailPoint("panicSpiceDBWriteResp")
+	if err != nil {
+		exists, err := isRelExists(ctx, h.PermissionClient, idempotencyKey)
+		if err != nil {
+			return struct{}{}, err
+		}
+
+		if exists {
+			klog.Infof("idempotency key already exists, relationships were already written: %v", idempotencyKey)
+			return struct{}{}, nil // idempotent write, key already exists
+		}
+
+		return struct{}{}, err
+	}
+
+	return struct{}{}, nil
+}
+
+// idempotencyKeyForPayload computes an idempotency key off a proto payload for a request, and the workflow ID that executed it.
+func idempotencyKeyForPayload(input *v1.WriteRelationshipsRequest, workflowID string) (*v1.Relationship, error) {
+	bytes, err := input.MarshalVT()
+	if err != nil {
+		return nil, err
+	}
+
+	return &v1.Relationship{
+		Resource: &v1.ObjectReference{
+			ObjectType: "workflow",
+			ObjectId:   workflowID,
+		},
+		Relation: "idempotency_key",
+		Subject: &v1.SubjectReference{
+			Object: &v1.ObjectReference{
+				ObjectType: "activity",
+				ObjectId:   fmt.Sprintf("%x", xxhash.Sum64(bytes)),
+			},
+		},
+	}, nil
+}
+
+func isRelExists(ctx context.Context, client v1.PermissionsServiceClient, toCheck *v1.Relationship) (bool, error) {
+	req := readRequestForRel(toCheck)
+
+	resp, err := client.ReadRelationships(ctx, req)
+	if err != nil {
+		klog.ErrorS(err, "failed to check existence of relationship", "rel", toCheck)
+		return false, fmt.Errorf("failed to determine if relationship exists: %w", err)
+	}
+
+	var exists bool
+	_, err = resp.Recv()
+	if err != nil {
+		if !errors.Is(err, io.EOF) {
+			return false, fmt.Errorf("failed to determine existence of relationship: %w", err)
+		}
+	} else {
+		exists = true
+	}
+
+	return exists, nil
+}
+
+func readRequestForRel(lock *v1.Relationship) *v1.ReadRelationshipsRequest {
+	req := &v1.ReadRelationshipsRequest{
+		RelationshipFilter: &v1.RelationshipFilter{
+			ResourceType:       lock.Resource.ObjectType,
+			OptionalResourceId: lock.Resource.ObjectId,
+			OptionalRelation:   lock.Relation,
+			OptionalSubjectFilter: &v1.SubjectFilter{
+				SubjectType:       lock.Subject.Object.ObjectType,
+				OptionalSubjectId: lock.Subject.Object.ObjectId,
+				OptionalRelation: &v1.SubjectFilter_RelationFilter{
+					Relation: lock.Subject.OptionalRelation,
+				},
+			},
+		},
+	}
+	return req
 }
 
 // ReadRelationships reads relationships from spicedb and returns any errors.

pkg/authz/distributedtx/workflow.go

@@ -83,7 +83,7 @@ func (r *RollbackRelationships) WithRels(relationships ...*v1.RelationshipUpdate
 	return r
 }
 
-func (r *RollbackRelationships) Cleanup(ctx workflow.Context, reason string) {
+func (r *RollbackRelationships) Cleanup(ctx workflow.Context, workflowID, reason string) {
 	invert := func(op v1.RelationshipUpdate_Operation) v1.RelationshipUpdate_Operation {
 		switch op {
 		case v1.RelationshipUpdate_OPERATION_CREATE:
@@ -106,10 +106,10 @@ func (r *RollbackRelationships) Cleanup(ctx workflow.Context, reason string) {
 	}
 
 	for {
-		f := workflow.ExecuteActivity[*v1.WriteRelationshipsResponse](ctx,
+		f := workflow.ExecuteActivity[struct{}](ctx,
 			workflow.DefaultActivityOptions,
 			activityHandler.WriteToSpiceDB,
-			&v1.WriteRelationshipsRequest{Updates: updates})
+			&v1.WriteRelationshipsRequest{Updates: updates}, workflowID)
 
 		if _, err := f.Get(ctx); err != nil {
 			if s, ok := status.FromError(err); ok {
@@ -181,18 +181,18 @@ func PessimisticWriteToSpiceDBAndKube(ctx workflow.Context, input *WriteObjInput
 		Updates:               append(updates, resourceLockRel),
 	}
 
-	_, err := workflow.ExecuteActivity[*v1.WriteRelationshipsResponse](ctx,
+	_, err := workflow.ExecuteActivity[struct{}](ctx,
 		workflow.DefaultActivityOptions,
 		activityHandler.WriteToSpiceDB,
-		arg).Get(ctx)
+		arg, instance.InstanceID).Get(ctx)
 	if err != nil {
 		// request failed for some reason
 		klog.ErrorS(err, "spicedb write failed")
 		for _, u := range updates {
 			klog.V(3).InfoS("update details", "update", u.String(), "relationship", u)
 		}
 
-		rollback.WithRels(updates...).Cleanup(ctx, "rollback due to failed SpiceDB write")
+		rollback.WithRels(updates...).Cleanup(ctx, instance.InstanceID, "rollback due to failed SpiceDB write")
 
 		// if the spicedb write fails, report it as a kube conflict error
 		// we return this for any error, not just lock conflicts, so that the
@@ -231,21 +231,21 @@ func PessimisticWriteToSpiceDBAndKube(ctx workflow.Context, input *WriteObjInput
 		isSuccessful, err := isSuccessfulKuberentesOperation(input, out)
 		if err != nil {
 			klog.V(1).ErrorS(err, "error checking kube response", "response", out, "verb", input.RequestInfo.Verb)
-			rollback.WithRels(updates...).Cleanup(ctx, "rollback due to failed kube operation after max attempts")
+			rollback.WithRels(updates...).Cleanup(ctx, instance.InstanceID, "rollback due to failed kube operation after max attempts")
 			return nil, fmt.Errorf("failed to communicate with kubernetes after %d attempts: %w", MaxKubeAttempts, err)
 		}
 
 		if isSuccessful {
-			rollback.Cleanup(ctx, fmt.Sprintf("cleanup after successful kube operation: %s", input.RequestInfo.Verb))
+			rollback.Cleanup(ctx, instance.InstanceID, fmt.Sprintf("cleanup after successful kube operation: %s", input.RequestInfo.Verb))
 			return out, nil
 		}
 
 		klog.V(3).ErrorS(err, "unsuccessful Kube API operation on PessimisticWriteToSpiceDBAndKube", "response", out, "verb", input.RequestInfo.Verb)
-		rollback.WithRels(updates...).Cleanup(ctx, "rollback due to unsuccessful kube operation")
+		rollback.WithRels(updates...).Cleanup(ctx, instance.InstanceID, "rollback due to unsuccessful kube operation")
 		return out, nil
 	}
 
-	rollback.WithRels(updates...).Cleanup(ctx, "rollback due to failed kube operation after max attempts")
+	rollback.WithRels(updates...).Cleanup(ctx, instance.InstanceID, "rollback due to failed kube operation after max attempts")
 	return nil, fmt.Errorf("failed to communicate with kubernetes after %d attempts", MaxKubeAttempts)
 }
 
@@ -310,15 +310,16 @@ func OptimisticWriteToSpiceDBAndKube(ctx workflow.Context, input *WriteObjInput)
 		return nil, fmt.Errorf("failed to append deletes from filters: %w", err)
 	}
 
+	instance := workflow.WorkflowInstance(ctx)
 	rollback := NewRollbackRelationships(updates...)
-	_, err := workflow.ExecuteActivity[*v1.WriteRelationshipsResponse](ctx,
+	_, err := workflow.ExecuteActivity[struct{}](ctx,
 		workflow.DefaultActivityOptions,
 		activityHandler.WriteToSpiceDB,
 		&v1.WriteRelationshipsRequest{
 			Updates: updates,
-		}).Get(ctx)
+		}, instance.InstanceID).Get(ctx)
 	if err != nil {
-		rollback.Cleanup(ctx, "rollback due to failed SpiceDB write")
+		rollback.Cleanup(ctx, instance.InstanceID, "rollback due to failed SpiceDB write")
 		klog.ErrorS(err, "SpiceDB write failed")
 		// report spicedb write errors as conflicts
 		return KubeConflict(err, input), nil
@@ -342,7 +343,7 @@ func OptimisticWriteToSpiceDBAndKube(ctx workflow.Context, input *WriteObjInput)
 
 		// if the object doesn't exist, clean up the spicedb write
 		if !exists {
-			rollback.Cleanup(ctx, "rollback due to failed Kube write")
+			rollback.Cleanup(ctx, instance.InstanceID, "rollback due to failed Kube write")
 			return nil, err
 		}
 	}

pkg/spicedb/bootstrap.yaml

@@ -28,6 +28,11 @@ schema: |-
   definition lock {
     relation workflow: workflow
   }
-  definition workflow {}
+
+  definition workflow {
+    relation idempotency_key: activity
+  }
+
+  definition activity{}
 relationships: |
   namespace:spicedb-kubeapi-proxy#viewer@user:rakis