Merge pull request #1176 from josephschorr/skip-rev-write

Skip loading of head revision on write calls

internal/middleware/consistency/consistency.go

@@ -6,15 +6,14 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/authzed/spicedb/internal/services/shared"
-
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
 	log "github.com/authzed/spicedb/internal/logging"
 	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
+	"github.com/authzed/spicedb/internal/services/shared"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/zedtoken"
 )
@@ -39,25 +38,18 @@ func ContextWithHandle(ctx context.Context) context.Context {
 	return context.WithValue(ctx, revisionKey, &revisionHandle{})
 }
 
-// RevisionFromContext reads the selected revision out of a context.Context and returns nil if it
-// does not exist.
-func RevisionFromContext(ctx context.Context) datastore.Revision {
+// RevisionFromContext reads the selected revision out of a context.Context, computes a zedtoken
+// from it, and returns an error if it has not been set on the context.
+func RevisionFromContext(ctx context.Context) (datastore.Revision, *v1.ZedToken, error) {
 	if c := ctx.Value(revisionKey); c != nil {
 		handle := c.(*revisionHandle)
-		return handle.revision
-	}
-	return nil
-}
-
-// MustRevisionFromContext reads the selected revision out of a context.Context, computes a zedtoken
-// from it, and panics if it has not been set on the context.
-func MustRevisionFromContext(ctx context.Context) (datastore.Revision, *v1.ZedToken) {
-	rev := RevisionFromContext(ctx)
-	if rev == nil {
-		panic("consistency middleware did not inject revision")
+		rev := handle.revision
+		if rev != nil {
+			return rev, zedtoken.MustNewFromRevision(rev), nil
+		}
 	}
 
-	return rev, zedtoken.MustNewFromRevision(rev)
+	return nil, nil, fmt.Errorf("consistency middleware did not inject revision")
 }
 
 // AddRevisionToContext adds a revision to the given context, based on the consistency block found
@@ -67,23 +59,8 @@ func AddRevisionToContext(ctx context.Context, req interface{}, ds datastore.Dat
 	case hasConsistency:
 		return addRevisionToContextFromConsistency(ctx, req, ds)
 	default:
-		return addHeadRevision(ctx, ds)
-	}
-}
-
-// addHeadRevision sets the value of the revision in the context to the current head revision in the datastore
-func addHeadRevision(ctx context.Context, ds datastore.Datastore) error {
-	handle := ctx.Value(revisionKey)
-	if handle == nil {
 		return nil
 	}
-
-	revision, err := ds.HeadRevision(ctx)
-	if err != nil {
-		return rewriteDatastoreError(ctx, err)
-	}
-	handle.(*revisionHandle).revision = revision
-	return nil
 }
 
 // addRevisionToContextFromConsistency adds a revision to the given context, based on the consistency block found

internal/middleware/consistency/consistency_test.go

@@ -3,19 +3,13 @@ package consistency
 import (
 	"context"
 	"errors"
-	"io"
 	"testing"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
-	"github.com/grpc-ecosystem/go-grpc-middleware/v2/testing/testpb"
 	"github.com/shopspring/decimal"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"github.com/stretchr/testify/suite"
-	"google.golang.org/grpc"
 
 	"github.com/authzed/spicedb/internal/datastore/proxy/proxy_test"
-	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
 	"github.com/authzed/spicedb/pkg/datastore/revision"
 	"github.com/authzed/spicedb/pkg/zedtoken"
 )
@@ -36,7 +30,11 @@ func TestAddRevisionToContextNoneSupplied(t *testing.T) {
 	updated := ContextWithHandle(context.Background())
 	err := AddRevisionToContext(updated, &v1.ReadRelationshipsRequest{}, ds)
 	require.NoError(err)
-	require.True(optimized.Equal(RevisionFromContext(updated)))
+
+	rev, _, err := RevisionFromContext(updated)
+	require.NoError(err)
+
+	require.True(optimized.Equal(rev))
 	ds.AssertExpectations(t)
 }
 
@@ -55,7 +53,11 @@ func TestAddRevisionToContextMinimizeLatency(t *testing.T) {
 		},
 	}, ds)
 	require.NoError(err)
-	require.True(optimized.Equal(RevisionFromContext(updated)))
+
+	rev, _, err := RevisionFromContext(updated)
+	require.NoError(err)
+
+	require.True(optimized.Equal(rev))
 	ds.AssertExpectations(t)
 }
 
@@ -74,7 +76,11 @@ func TestAddRevisionToContextFullyConsistent(t *testing.T) {
 		},
 	}, ds)
 	require.NoError(err)
-	require.True(head.Equal(RevisionFromContext(updated)))
+
+	rev, _, err := RevisionFromContext(updated)
+	require.NoError(err)
+
+	require.True(head.Equal(rev))
 	ds.AssertExpectations(t)
 }
 
@@ -94,7 +100,11 @@ func TestAddRevisionToContextAtLeastAsFresh(t *testing.T) {
 		},
 	}, ds)
 	require.NoError(err)
-	require.True(exact.Equal(RevisionFromContext(updated)))
+
+	rev, _, err := RevisionFromContext(updated)
+	require.NoError(err)
+
+	require.True(exact.Equal(rev))
 	ds.AssertExpectations(t)
 }
 
@@ -114,7 +124,11 @@ func TestAddRevisionToContextAtValidExactSnapshot(t *testing.T) {
 		},
 	}, ds)
 	require.NoError(err)
-	require.True(exact.Equal(RevisionFromContext(updated)))
+
+	rev, _, err := RevisionFromContext(updated)
+	require.NoError(err)
+
+	require.True(exact.Equal(rev))
 	ds.AssertExpectations(t)
 }
 
@@ -137,65 +151,11 @@ func TestAddRevisionToContextAtInvalidExactSnapshot(t *testing.T) {
 	ds.AssertExpectations(t)
 }
 
-func TestAddRevisionToContextAPIAlwaysFullyConsistent(t *testing.T) {
+func TestAddRevisionToContextNoConsistencyAPI(t *testing.T) {
 	require := require.New(t)
 
-	ds := &proxy_test.MockDatastore{}
-	ds.On("HeadRevision").Return(head, nil).Once()
-
 	updated := ContextWithHandle(context.Background())
-	err := AddRevisionToContext(updated, &v1.WriteSchemaRequest{}, ds)
-	require.NoError(err)
-	require.True(head.Equal(RevisionFromContext(updated)))
-	ds.AssertExpectations(t)
-}
-
-func TestMiddlewareConsistencyTestSuite(t *testing.T) {
-	ds := &proxy_test.MockDatastore{}
-	ds.On("HeadRevision").Return(head, nil)
-
-	s := &ConsistencyTestSuite{
-		InterceptorTestSuite: &testpb.InterceptorTestSuite{
-			ServerOpts: []grpc.ServerOption{
-				grpc.ChainStreamInterceptor(
-					datastoremw.StreamServerInterceptor(ds),
-					StreamServerInterceptor(),
-				),
-				grpc.ChainUnaryInterceptor(
-					datastoremw.UnaryServerInterceptor(ds),
-					UnaryServerInterceptor(),
-				),
-			},
-		},
-	}
-	suite.Run(t, s)
-	ds.AssertExpectations(t)
-}
-
-var (
-	goodPing = &testpb.PingRequest{Value: "something"}
-	goodList = &testpb.PingListRequest{Value: "something"}
-)
 
-type ConsistencyTestSuite struct {
-	*testpb.InterceptorTestSuite
-}
-
-func (s *ConsistencyTestSuite) TestValidPasses_Unary() {
-	require := require.New(s.T())
-	_, err := s.Client.Ping(s.SimpleCtx(), goodPing)
-	require.NoError(err)
-}
-
-func (s *ConsistencyTestSuite) TestValidPasses_ServerStream() {
-	require := require.New(s.T())
-	stream, err := s.Client.PingList(s.SimpleCtx(), goodList)
-	require.NoError(err)
-	for {
-		_, err := stream.Recv()
-		if errors.Is(err, io.EOF) {
-			break
-		}
-		assert.NoError(s.T(), err, "no error on messages sent occurred")
-	}
+	_, _, err := RevisionFromContext(updated)
+	require.Error(err)
 }

internal/services/v1/permissions.go

@@ -34,7 +34,11 @@ import (
 const maxCaveatContextBytes = 4096
 
 func (ps *permissionServer) CheckPermission(ctx context.Context, req *v1.CheckPermissionRequest) (*v1.CheckPermissionResponse, error) {
-	atRevision, checkedAt := consistency.MustRevisionFromContext(ctx)
+	atRevision, checkedAt, err := consistency.RevisionFromContext(ctx)
+	if err != nil {
+		return nil, rewriteError(ctx, err)
+	}
+
 	ds := datastoremw.MustFromContext(ctx).SnapshotReader(atRevision)
 
 	caveatContext, err := getCaveatContext(ctx, req.Context)
@@ -138,10 +142,14 @@ func (ps *permissionServer) CheckPermission(ctx context.Context, req *v1.CheckPe
 }
 
 func (ps *permissionServer) ExpandPermissionTree(ctx context.Context, req *v1.ExpandPermissionTreeRequest) (*v1.ExpandPermissionTreeResponse, error) {
-	atRevision, expandedAt := consistency.MustRevisionFromContext(ctx)
+	atRevision, expandedAt, err := consistency.RevisionFromContext(ctx)
+	if err != nil {
+		return nil, rewriteError(ctx, err)
+	}
+
 	ds := datastoremw.MustFromContext(ctx).SnapshotReader(atRevision)
 
-	err := namespace.CheckNamespaceAndRelation(ctx, req.Resource.ObjectType, req.Permission, false, ds)
+	err = namespace.CheckNamespaceAndRelation(ctx, req.Resource.ObjectType, req.Permission, false, ds)
 	if err != nil {
 		return nil, rewriteError(ctx, err)
 	}
@@ -318,7 +326,11 @@ func TranslateExpansionTree(node *core.RelationTupleTreeNode) *v1.PermissionRela
 
 func (ps *permissionServer) LookupResources(req *v1.LookupResourcesRequest, resp v1.PermissionsService_LookupResourcesServer) error {
 	ctx := resp.Context()
-	atRevision, revisionReadAt := consistency.MustRevisionFromContext(ctx)
+	atRevision, revisionReadAt, err := consistency.RevisionFromContext(ctx)
+	if err != nil {
+		return rewriteError(ctx, err)
+	}
+
 	ds := datastoremw.MustFromContext(ctx).SnapshotReader(atRevision)
 
 	// Perform our preflight checks in parallel
@@ -393,7 +405,10 @@ func (ps *permissionServer) LookupResources(req *v1.LookupResourcesRequest, resp
 
 func (ps *permissionServer) LookupSubjects(req *v1.LookupSubjectsRequest, resp v1.PermissionsService_LookupSubjectsServer) error {
 	ctx := resp.Context()
-	atRevision, revisionReadAt := consistency.MustRevisionFromContext(ctx)
+	atRevision, revisionReadAt, err := consistency.RevisionFromContext(ctx)
+	if err != nil {
+		return rewriteError(ctx, err)
+	}
 
 	ds := datastoremw.MustFromContext(ctx).SnapshotReader(atRevision)
 

internal/services/v1/relationships.go

@@ -111,7 +111,11 @@ func (ps *permissionServer) checkFilterNamespaces(ctx context.Context, filter *v
 
 func (ps *permissionServer) ReadRelationships(req *v1.ReadRelationshipsRequest, resp v1.PermissionsService_ReadRelationshipsServer) error {
 	ctx := resp.Context()
-	atRevision, revisionReadAt := consistency.MustRevisionFromContext(ctx)
+	atRevision, revisionReadAt, err := consistency.RevisionFromContext(ctx)
+	if err != nil {
+		return rewriteError(ctx, err)
+	}
+
 	ds := datastoremw.MustFromContext(ctx).SnapshotReader(atRevision)
 
 	if err := ps.checkFilterNamespaces(ctx, req.RelationshipFilter, ds); err != nil {

internal/services/v1/schema.go

@@ -10,7 +10,6 @@ import (
 
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/middleware"
-	"github.com/authzed/spicedb/internal/middleware/consistency"
 	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
 	"github.com/authzed/spicedb/internal/middleware/usagemetrics"
 	"github.com/authzed/spicedb/internal/services/shared"
@@ -46,15 +45,21 @@ type schemaServer struct {
 }
 
 func (ss *schemaServer) ReadSchema(ctx context.Context, in *v1.ReadSchemaRequest) (*v1.ReadSchemaResponse, error) {
-	readRevision, _ := consistency.MustRevisionFromContext(ctx)
-	ds := datastoremw.MustFromContext(ctx).SnapshotReader(readRevision)
+	// Schema is always read from the head revision.
+	ds := datastoremw.MustFromContext(ctx)
+	headRevision, err := ds.HeadRevision(ctx)
+	if err != nil {
+		return nil, rewriteError(ctx, err)
+	}
+
+	reader := ds.SnapshotReader(headRevision)
 
-	nsDefs, err := ds.ListAllNamespaces(ctx)
+	nsDefs, err := reader.ListAllNamespaces(ctx)
 	if err != nil {
 		return nil, rewriteError(ctx, err)
 	}
 
-	caveatDefs, err := ds.ListAllCaveats(ctx)
+	caveatDefs, err := reader.ListAllCaveats(ctx)
 	if err != nil {
 		return nil, rewriteError(ctx, err)
 	}

pkg/middleware/consistency/consistency.go

@@ -7,22 +7,10 @@ import (
 
 	"github.com/authzed/spicedb/internal/middleware/consistency"
 	"github.com/authzed/spicedb/pkg/datastore"
-	"github.com/authzed/spicedb/pkg/zedtoken"
 )
 
-// RevisionFromContext reads the selected revision out of a context.Context and returns nil if it
-// does not exist.
-func RevisionFromContext(ctx context.Context) datastore.Revision {
+// RevisionFromContext reads the selected revision out of a context.Context, computes a zedtoken
+// from it, and returns an internal error if it has not been set on the context.
+func RevisionFromContext(ctx context.Context) (datastore.Revision, *v1.ZedToken, error) {
 	return consistency.RevisionFromContext(ctx)
 }
-
-// MustRevisionFromContext reads the selected revision out of a context.Context, computes a zedtoken
-// from it, and panics if it has not been set on the context.
-func MustRevisionFromContext(ctx context.Context) (datastore.Revision, *v1.ZedToken) {
-	rev := consistency.RevisionFromContext(ctx)
-	if rev == nil {
-		panic("consistency middleware did not inject revision")
-	}
-
-	return rev, zedtoken.MustNewFromRevision(rev)
-}