From 3e8fa00c1363148508f534fe2094f931bd23b9a3 Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Wed, 1 Jul 2026 16:32:01 +0200 Subject: [PATCH 01/13] perf: use singleflight for more db queries to better handle many similar parallel requests --- .../proxies/singleflight/data_reader.go | 114 +++- .../proxies/singleflight/data_reader_test.go | 610 ++++++++++++++++-- 2 files changed, 646 insertions(+), 78 deletions(-) diff --git a/internal/storage/proxies/singleflight/data_reader.go b/internal/storage/proxies/singleflight/data_reader.go index f9564223e..18b4d1703 100644 --- a/internal/storage/proxies/singleflight/data_reader.go +++ b/internal/storage/proxies/singleflight/data_reader.go @@ -2,6 +2,8 @@ package singleflight import ( "context" + "fmt" + "strings" "resenje.org/singleflight" @@ -13,8 +15,11 @@ import ( // DataReader - Add singleflight behaviour to data reader type DataReader struct { - delegate storage.DataReader - group singleflight.Group[string, token.SnapToken] + delegate storage.DataReader + headSnapshotGroup singleflight.Group[string, token.SnapToken] + queryRelationshipsGroup singleflight.Group[string, []*base.Tuple] + querySingleAttrGroup singleflight.Group[string, *base.Attribute] + queryAttributesGroup singleflight.Group[string, []*base.Attribute] } // NewDataReader - Add singleflight behaviour to new data reader @@ -22,9 +27,20 @@ func NewDataReader(delegate storage.DataReader) *DataReader { return &DataReader{delegate: delegate} } -// QueryRelationships - Reads relation tuples from the repository +// QueryRelationships - Reads relation tuples from the repository with singleflight deduplication. func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, token string, pagination database.CursorPagination) (*database.TupleIterator, error) { - return r.delegate.QueryRelationships(ctx, tenantID, filter, token, pagination) + key := queryRelationshipsKey(tenantID, filter, token, pagination) + tuples, _, err := r.queryRelationshipsGroup.Do(ctx, key, func(ctx context.Context) ([]*base.Tuple, error) { + it, err := r.delegate.QueryRelationships(ctx, tenantID, filter, token, pagination) + if err != nil { + return nil, err + } + return drainTupleIterator(it), nil + }) + if err != nil { + return nil, err + } + return database.NewTupleIterator(tuples...), nil } // ReadRelationships - Reads relation tuples from the repository with different options. @@ -32,14 +48,29 @@ func (r *DataReader) ReadRelationships(ctx context.Context, tenantID string, fil return r.delegate.ReadRelationships(ctx, tenantID, filter, token, pagination) } -// QuerySingleAttribute - Reads a single attribute from the repository. +// QuerySingleAttribute - Reads a single attribute from the repository with singleflight deduplication. func (r *DataReader) QuerySingleAttribute(ctx context.Context, tenantID string, filter *base.AttributeFilter, token string) (*base.Attribute, error) { - return r.delegate.QuerySingleAttribute(ctx, tenantID, filter, token) + key := querySingleAttributeKey(tenantID, filter, token) + attr, _, err := r.querySingleAttrGroup.Do(ctx, key, func(ctx context.Context) (*base.Attribute, error) { + return r.delegate.QuerySingleAttribute(ctx, tenantID, filter, token) + }) + return attr, err } -// QueryAttributes - Reads multiple attributes from the repository. +// QueryAttributes - Reads multiple attributes from the repository with singleflight deduplication. func (r *DataReader) QueryAttributes(ctx context.Context, tenantID string, filter *base.AttributeFilter, token string, pagination database.CursorPagination) (*database.AttributeIterator, error) { - return r.delegate.QueryAttributes(ctx, tenantID, filter, token, pagination) + key := queryAttributesKey(tenantID, filter, token, pagination) + attrs, _, err := r.queryAttributesGroup.Do(ctx, key, func(ctx context.Context) ([]*base.Attribute, error) { + it, err := r.delegate.QueryAttributes(ctx, tenantID, filter, token, pagination) + if err != nil { + return nil, err + } + return drainAttributeIterator(it), nil + }) + if err != nil { + return nil, err + } + return database.NewAttributeIterator(attrs...), nil } // ReadAttributes - Reads multiple attributes from the repository with different options. @@ -54,8 +85,73 @@ func (r *DataReader) QueryUniqueSubjectReferences(ctx context.Context, tenantID // HeadSnapshot - Reads the latest version of the snapshot from the repository. func (r *DataReader) HeadSnapshot(ctx context.Context, tenantID string) (token.SnapToken, error) { - rev, _, err := r.group.Do(ctx, tenantID, func(ctx context.Context) (token.SnapToken, error) { // tenantID ensures proper tenant isolation in deduplication + rev, _, err := r.headSnapshotGroup.Do(ctx, tenantID, func(ctx context.Context) (token.SnapToken, error) { return r.delegate.HeadSnapshot(ctx, tenantID) }) return rev, err } + +// --- key builders --- + +func queryRelationshipsKey(tenantID string, filter *base.TupleFilter, token string, pagination database.CursorPagination) string { + var b strings.Builder + fmt.Fprintf(&b, "qr\x00%q\x00%q\x00%#v\x00%q\x00%q\x00%#v\x00%q\x00%q\x00%q\x00%q\x00%d", + tenantID, + filter.GetEntity().GetType(), + filter.GetEntity().GetIds(), + filter.GetRelation(), + filter.GetSubject().GetType(), + filter.GetSubject().GetIds(), + filter.GetSubject().GetRelation(), + token, + pagination.Cursor(), + pagination.Sort(), + pagination.Limit(), + ) + return b.String() +} + +func querySingleAttributeKey(tenantID string, filter *base.AttributeFilter, token string) string { + var b strings.Builder + fmt.Fprintf(&b, "qsa\x00%q\x00%q\x00%#v\x00%#v\x00%q", + tenantID, + filter.GetEntity().GetType(), + filter.GetEntity().GetIds(), + filter.GetAttributes(), + token, + ) + return b.String() +} + +func queryAttributesKey(tenantID string, filter *base.AttributeFilter, token string, pagination database.CursorPagination) string { + var b strings.Builder + fmt.Fprintf(&b, "qa\x00%q\x00%q\x00%#v\x00%#v\x00%q\x00%q\x00%q\x00%d", + tenantID, + filter.GetEntity().GetType(), + filter.GetEntity().GetIds(), + filter.GetAttributes(), + token, + pagination.Cursor(), + pagination.Sort(), + pagination.Limit(), + ) + return b.String() +} + +// --- iterator helpers --- + +func drainTupleIterator(it *database.TupleIterator) []*base.Tuple { + var tuples []*base.Tuple + for it.HasNext() { + tuples = append(tuples, it.GetNext()) + } + return tuples +} + +func drainAttributeIterator(it *database.AttributeIterator) []*base.Attribute { + var attrs []*base.Attribute + for it.HasNext() { + attrs = append(attrs, it.GetNext()) + } + return attrs +} diff --git a/internal/storage/proxies/singleflight/data_reader_test.go b/internal/storage/proxies/singleflight/data_reader_test.go index 0fd025047..205755582 100644 --- a/internal/storage/proxies/singleflight/data_reader_test.go +++ b/internal/storage/proxies/singleflight/data_reader_test.go @@ -10,6 +10,8 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "google.golang.org/protobuf/types/known/anypb" + "github.com/Permify/permify/internal/storage" "github.com/Permify/permify/pkg/database" base "github.com/Permify/permify/pkg/pb/base/v1" @@ -19,33 +21,95 @@ import ( // MockDataReader is a mock implementation of storage.DataReader for testing type MockDataReader struct { storage.NoopDataReader - headSnapshotCalls map[string]*int64 - mu sync.Mutex + headSnapshotCalls map[string]*int64 + queryRelationshipsCalls map[string]*int64 + querySingleAttributeCalls map[string]*int64 + queryAttributesCalls map[string]*int64 + mu sync.Mutex } func NewMockDataReader() *MockDataReader { return &MockDataReader{ - headSnapshotCalls: make(map[string]*int64), + headSnapshotCalls: make(map[string]*int64), + queryRelationshipsCalls: make(map[string]*int64), + querySingleAttributeCalls: make(map[string]*int64), + queryAttributesCalls: make(map[string]*int64), } } -func (m *MockDataReader) HeadSnapshot(ctx context.Context, tenantID string) (token.SnapToken, error) { - // Track call count per tenant +func (m *MockDataReader) HeadSnapshot(_ context.Context, tenantID string) (token.SnapToken, error) { + m.incrementCounter(m.headSnapshotCalls, tenantID) + time.Sleep(10 * time.Millisecond) + return token.NoopToken{Value: "snapshot-" + tenantID}, nil +} + +func (m *MockDataReader) QueryRelationships(_ context.Context, tenantID string, filter *base.TupleFilter, _ string, _ database.CursorPagination) (*database.TupleIterator, error) { + key := tenantID + "|" + filter.GetEntity().GetType() + "|" + filter.GetRelation() + m.incrementCounter(m.queryRelationshipsCalls, key) + time.Sleep(10 * time.Millisecond) + return database.NewTupleIterator( + &base.Tuple{ + Entity: &base.Entity{Type: filter.GetEntity().GetType(), Id: "1"}, + Relation: filter.GetRelation(), + Subject: &base.Subject{Type: "user", Id: "alice"}, + }, + &base.Tuple{ + Entity: &base.Entity{Type: filter.GetEntity().GetType(), Id: "1"}, + Relation: filter.GetRelation(), + Subject: &base.Subject{Type: "user", Id: "bob"}, + }, + ), nil +} + +func (m *MockDataReader) QuerySingleAttribute(_ context.Context, tenantID string, filter *base.AttributeFilter, _ string) (*base.Attribute, error) { + key := tenantID + "|" + filter.GetEntity().GetType() + m.incrementCounter(m.querySingleAttributeCalls, key) + time.Sleep(10 * time.Millisecond) + val, _ := anypb.New(&base.BooleanValue{Data: true}) + return &base.Attribute{ + Entity: &base.Entity{Type: filter.GetEntity().GetType(), Id: "1"}, + Attribute: "is_public", + Value: val, + }, nil +} + +func (m *MockDataReader) QueryAttributes(_ context.Context, tenantID string, filter *base.AttributeFilter, _ string, _ database.CursorPagination) (*database.AttributeIterator, error) { + key := tenantID + "|" + filter.GetEntity().GetType() + m.incrementCounter(m.queryAttributesCalls, key) + time.Sleep(10 * time.Millisecond) + val, _ := anypb.New(&base.BooleanValue{Data: true}) + return database.NewAttributeIterator( + &base.Attribute{ + Entity: &base.Entity{Type: filter.GetEntity().GetType(), Id: "1"}, + Attribute: "is_public", + Value: val, + }, + ), nil +} + +func (m *MockDataReader) incrementCounter(counters map[string]*int64, key string) { m.mu.Lock() - counter, exists := m.headSnapshotCalls[tenantID] + counter, exists := counters[key] if !exists { counter = new(int64) - m.headSnapshotCalls[tenantID] = counter + counters[key] = counter } m.mu.Unlock() - - // Increment call count atomic.AddInt64(counter, 1) +} - // Simulate some work - time.Sleep(10 * time.Millisecond) +func getCallCount(counters map[string]*int64, mu *sync.Mutex, key string) int64 { + mu.Lock() + defer mu.Unlock() + if counter, exists := counters[key]; exists { + return atomic.LoadInt64(counter) + } + return 0 +} - return token.NoopToken{Value: "snapshot-" + tenantID}, nil +// GetCallCount returns the HeadSnapshot call count for a tenant (for backward compat) +func GetCallCount(m *MockDataReader, tenantID string) int64 { + return getCallCount(m.headSnapshotCalls, &m.mu, tenantID) } // ErrorMockDataReader is a mock that returns errors for testing error handling @@ -53,22 +117,25 @@ type ErrorMockDataReader struct { storage.NoopDataReader } -func (m *ErrorMockDataReader) HeadSnapshot(ctx context.Context, tenantID string) (token.SnapToken, error) { +func (m *ErrorMockDataReader) HeadSnapshot(_ context.Context, _ string) (token.SnapToken, error) { return nil, errors.New("delegate error") } -func GetCallCount(m *MockDataReader, tenantID string) int64 { - m.mu.Lock() - defer m.mu.Unlock() - if counter, exists := m.headSnapshotCalls[tenantID]; exists { - return atomic.LoadInt64(counter) - } - return 0 +func (m *ErrorMockDataReader) QueryRelationships(_ context.Context, _ string, _ *base.TupleFilter, _ string, _ database.CursorPagination) (*database.TupleIterator, error) { + return nil, errors.New("delegate error") +} + +func (m *ErrorMockDataReader) QuerySingleAttribute(_ context.Context, _ string, _ *base.AttributeFilter, _ string) (*base.Attribute, error) { + return nil, errors.New("delegate error") +} + +func (m *ErrorMockDataReader) QueryAttributes(_ context.Context, _ string, _ *base.AttributeFilter, _ string, _ database.CursorPagination) (*database.AttributeIterator, error) { + return nil, errors.New("delegate error") } var _ = Describe("Singleflight DataReader", func() { var ( - mockDelegate storage.DataReader + mockDelegate *MockDataReader reader *DataReader ctx context.Context ) @@ -96,19 +163,6 @@ var _ = Describe("Singleflight DataReader", func() { }) }) - Describe("QueryRelationships", func() { - It("should delegate to underlying DataReader", func() { - delegate := storage.NewNoopRelationshipReader() - reader := NewDataReader(delegate) - - filter := &base.TupleFilter{} - iterator, err := reader.QueryRelationships(ctx, "tenant1", filter, "token", database.CursorPagination{}) - - Expect(err).ShouldNot(HaveOccurred()) - Expect(iterator).ShouldNot(BeNil()) - }) - }) - Describe("ReadRelationships", func() { It("should delegate to underlying DataReader", func() { delegate := storage.NewNoopRelationshipReader() @@ -123,32 +177,6 @@ var _ = Describe("Singleflight DataReader", func() { }) }) - Describe("QuerySingleAttribute", func() { - It("should delegate to underlying DataReader", func() { - delegate := storage.NewNoopRelationshipReader() - reader := NewDataReader(delegate) - - filter := &base.AttributeFilter{} - attribute, err := reader.QuerySingleAttribute(ctx, "tenant1", filter, "token") - - Expect(err).ShouldNot(HaveOccurred()) - Expect(attribute).ShouldNot(BeNil()) - }) - }) - - Describe("QueryAttributes", func() { - It("should delegate to underlying DataReader", func() { - delegate := storage.NewNoopRelationshipReader() - reader := NewDataReader(delegate) - - filter := &base.AttributeFilter{} - iterator, err := reader.QueryAttributes(ctx, "tenant1", filter, "token", database.CursorPagination{}) - - Expect(err).ShouldNot(HaveOccurred()) - Expect(iterator).ShouldNot(BeNil()) - }) - }) - Describe("ReadAttributes", func() { It("should delegate to underlying DataReader", func() { delegate := storage.NewNoopRelationshipReader() @@ -200,8 +228,7 @@ var _ = Describe("Singleflight DataReader", func() { wg.Wait() // Only 1 call should reach the delegate due to deduplication - mock := mockDelegate.(*MockDataReader) - callCount := GetCallCount(mock, tenantID) + callCount := GetCallCount(mockDelegate, tenantID) Expect(callCount).To(Equal(int64(1))) }) @@ -233,11 +260,9 @@ var _ = Describe("Singleflight DataReader", func() { wg.Wait() - mock := mockDelegate.(*MockDataReader) - // Each tenant should have exactly 1 call due to deduplication within the tenant - Expect(GetCallCount(mock, tenant1)).To(Equal(int64(1))) - Expect(GetCallCount(mock, tenant2)).To(Equal(int64(1))) + Expect(GetCallCount(mockDelegate, tenant1)).To(Equal(int64(1))) + Expect(GetCallCount(mockDelegate, tenant2)).To(Equal(int64(1))) }) It("should return correct snapshot for each tenant", func() { @@ -284,10 +309,8 @@ var _ = Describe("Singleflight DataReader", func() { _, err = reader.HeadSnapshot(ctx, tenantID) Expect(err).ShouldNot(HaveOccurred()) - mock := mockDelegate.(*MockDataReader) - // Should have 2 calls to the delegate - callCount := GetCallCount(mock, tenantID) + callCount := GetCallCount(mockDelegate, tenantID) Expect(callCount).To(Equal(int64(2))) }) @@ -340,4 +363,453 @@ var _ = Describe("Singleflight DataReader", func() { Expect(atomic.LoadInt64(&errorCount)).To(Equal(int64(numConcurrentRequests))) }) }) + + Describe("QueryRelationships", func() { + filter := &base.TupleFilter{ + Entity: &base.EntityFilter{ + Type: "document", + Ids: []string{"1"}, + }, + Relation: "viewer", + } + snap := "snap-token-1" + mockKey := "tenant1|document|viewer" + + It("should deduplicate concurrent requests with the same parameters", func() { + numConcurrentRequests := 10 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + it, err := reader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + Expect(it).ShouldNot(BeNil()) + }() + } + + wg.Wait() + + callCount := getCallCount(mockDelegate.queryRelationshipsCalls, &mockDelegate.mu, mockKey) + Expect(callCount).To(Equal(int64(1))) + }) + + It("should return independent iterators to each caller", func() { + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + iterators := make([]*database.TupleIterator, numConcurrentRequests) + var mu sync.Mutex + + for i := 0; i < numConcurrentRequests; i++ { + idx := i + go func() { + defer wg.Done() + it, err := reader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + mu.Lock() + iterators[idx] = it + mu.Unlock() + }() + } + + wg.Wait() + + // Each iterator should independently yield all tuples + for i, it := range iterators { + Expect(it.HasNext()).To(BeTrue(), "iterator %d should have tuples", i) + t1 := it.GetNext() + Expect(t1).ShouldNot(BeNil()) + Expect(t1.GetSubject().GetId()).To(Equal("alice")) + t2 := it.GetNext() + Expect(t2).ShouldNot(BeNil()) + Expect(t2.GetSubject().GetId()).To(Equal("bob")) + Expect(it.HasNext()).To(BeFalse(), "iterator %d should be exhausted", i) + } + }) + + It("should isolate requests for different tenants", func() { + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests * 2) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := reader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + }() + } + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := reader.QueryRelationships(ctx, "tenant2", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + }() + } + + wg.Wait() + + Expect(getCallCount(mockDelegate.queryRelationshipsCalls, &mockDelegate.mu, "tenant1|document|viewer")).To(Equal(int64(1))) + Expect(getCallCount(mockDelegate.queryRelationshipsCalls, &mockDelegate.mu, "tenant2|document|viewer")).To(Equal(int64(1))) + }) + + It("should isolate requests with different filters", func() { + filter2 := &base.TupleFilter{ + Entity: &base.EntityFilter{ + Type: "document", + Ids: []string{"1"}, + }, + Relation: "editor", + } + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + _, err := reader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + }() + + go func() { + defer wg.Done() + _, err := reader.QueryRelationships(ctx, "tenant1", filter2, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + }() + + wg.Wait() + + Expect(getCallCount(mockDelegate.queryRelationshipsCalls, &mockDelegate.mu, "tenant1|document|viewer")).To(Equal(int64(1))) + Expect(getCallCount(mockDelegate.queryRelationshipsCalls, &mockDelegate.mu, "tenant1|document|editor")).To(Equal(int64(1))) + }) + + It("should not deduplicate sequential requests", func() { + _, err := reader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + + time.Sleep(50 * time.Millisecond) + + _, err = reader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + + callCount := getCallCount(mockDelegate.queryRelationshipsCalls, &mockDelegate.mu, mockKey) + Expect(callCount).To(Equal(int64(2))) + }) + + It("should propagate errors from delegate", func() { + errorReader := NewDataReader(&ErrorMockDataReader{}) + + _, err := errorReader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + + Expect(err).Should(HaveOccurred()) + Expect(err.Error()).Should(Equal("delegate error")) + }) + + It("should handle concurrent requests with errors", func() { + errorReader := NewDataReader(&ErrorMockDataReader{}) + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + errorCount := int64(0) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := errorReader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + if err != nil { + atomic.AddInt64(&errorCount, 1) + } + }() + } + + wg.Wait() + + Expect(atomic.LoadInt64(&errorCount)).To(Equal(int64(numConcurrentRequests))) + }) + + It("should handle empty result sets", func() { + emptyReader := NewDataReader(storage.NewNoopRelationshipReader()) + + it, err := emptyReader.QueryRelationships(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + + Expect(err).ShouldNot(HaveOccurred()) + Expect(it).ShouldNot(BeNil()) + Expect(it.HasNext()).To(BeFalse()) + }) + }) + + Describe("QuerySingleAttribute", func() { + filter := &base.AttributeFilter{ + Entity: &base.EntityFilter{ + Type: "document", + Ids: []string{"1"}, + }, + Attributes: []string{"is_public"}, + } + snap := "snap-token-1" + mockKey := "tenant1|document" + + It("should deduplicate concurrent requests with the same parameters", func() { + numConcurrentRequests := 10 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + attr, err := reader.QuerySingleAttribute(ctx, "tenant1", filter, snap) + Expect(err).ShouldNot(HaveOccurred()) + Expect(attr).ShouldNot(BeNil()) + }() + } + + wg.Wait() + + callCount := getCallCount(mockDelegate.querySingleAttributeCalls, &mockDelegate.mu, mockKey) + Expect(callCount).To(Equal(int64(1))) + }) + + It("should isolate requests for different tenants", func() { + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests * 2) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := reader.QuerySingleAttribute(ctx, "tenant1", filter, snap) + Expect(err).ShouldNot(HaveOccurred()) + }() + } + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := reader.QuerySingleAttribute(ctx, "tenant2", filter, snap) + Expect(err).ShouldNot(HaveOccurred()) + }() + } + + wg.Wait() + + Expect(getCallCount(mockDelegate.querySingleAttributeCalls, &mockDelegate.mu, "tenant1|document")).To(Equal(int64(1))) + Expect(getCallCount(mockDelegate.querySingleAttributeCalls, &mockDelegate.mu, "tenant2|document")).To(Equal(int64(1))) + }) + + It("should not deduplicate sequential requests", func() { + _, err := reader.QuerySingleAttribute(ctx, "tenant1", filter, snap) + Expect(err).ShouldNot(HaveOccurred()) + + time.Sleep(50 * time.Millisecond) + + _, err = reader.QuerySingleAttribute(ctx, "tenant1", filter, snap) + Expect(err).ShouldNot(HaveOccurred()) + + callCount := getCallCount(mockDelegate.querySingleAttributeCalls, &mockDelegate.mu, mockKey) + Expect(callCount).To(Equal(int64(2))) + }) + + It("should propagate errors from delegate", func() { + errorReader := NewDataReader(&ErrorMockDataReader{}) + + _, err := errorReader.QuerySingleAttribute(ctx, "tenant1", filter, snap) + + Expect(err).Should(HaveOccurred()) + Expect(err.Error()).Should(Equal("delegate error")) + }) + + It("should handle concurrent requests with errors", func() { + errorReader := NewDataReader(&ErrorMockDataReader{}) + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + errorCount := int64(0) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := errorReader.QuerySingleAttribute(ctx, "tenant1", filter, snap) + if err != nil { + atomic.AddInt64(&errorCount, 1) + } + }() + } + + wg.Wait() + + Expect(atomic.LoadInt64(&errorCount)).To(Equal(int64(numConcurrentRequests))) + }) + + It("should handle nil result", func() { + emptyReader := NewDataReader(storage.NewNoopRelationshipReader()) + + // NoopDataReader returns an empty Attribute, not nil + attr, err := emptyReader.QuerySingleAttribute(ctx, "tenant1", filter, snap) + Expect(err).ShouldNot(HaveOccurred()) + Expect(attr).ShouldNot(BeNil()) + }) + }) + + Describe("QueryAttributes", func() { + filter := &base.AttributeFilter{ + Entity: &base.EntityFilter{ + Type: "document", + Ids: []string{"1"}, + }, + Attributes: []string{"is_public"}, + } + snap := "snap-token-1" + mockKey := "tenant1|document" + + It("should deduplicate concurrent requests with the same parameters", func() { + numConcurrentRequests := 10 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + it, err := reader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + Expect(it).ShouldNot(BeNil()) + }() + } + + wg.Wait() + + callCount := getCallCount(mockDelegate.queryAttributesCalls, &mockDelegate.mu, mockKey) + Expect(callCount).To(Equal(int64(1))) + }) + + It("should return independent iterators to each caller", func() { + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + iterators := make([]*database.AttributeIterator, numConcurrentRequests) + var mu sync.Mutex + + for i := 0; i < numConcurrentRequests; i++ { + idx := i + go func() { + defer wg.Done() + it, err := reader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + mu.Lock() + iterators[idx] = it + mu.Unlock() + }() + } + + wg.Wait() + + // Each iterator should independently yield all attributes + for i, it := range iterators { + Expect(it.HasNext()).To(BeTrue(), "iterator %d should have attributes", i) + a := it.GetNext() + Expect(a).ShouldNot(BeNil()) + Expect(a.GetAttribute()).To(Equal("is_public")) + Expect(it.HasNext()).To(BeFalse(), "iterator %d should be exhausted", i) + } + }) + + It("should isolate requests for different tenants", func() { + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests * 2) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := reader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + }() + } + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := reader.QueryAttributes(ctx, "tenant2", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + }() + } + + wg.Wait() + + Expect(getCallCount(mockDelegate.queryAttributesCalls, &mockDelegate.mu, "tenant1|document")).To(Equal(int64(1))) + Expect(getCallCount(mockDelegate.queryAttributesCalls, &mockDelegate.mu, "tenant2|document")).To(Equal(int64(1))) + }) + + It("should not deduplicate sequential requests", func() { + _, err := reader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + + time.Sleep(50 * time.Millisecond) + + _, err = reader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + Expect(err).ShouldNot(HaveOccurred()) + + callCount := getCallCount(mockDelegate.queryAttributesCalls, &mockDelegate.mu, mockKey) + Expect(callCount).To(Equal(int64(2))) + }) + + It("should propagate errors from delegate", func() { + errorReader := NewDataReader(&ErrorMockDataReader{}) + + _, err := errorReader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + + Expect(err).Should(HaveOccurred()) + Expect(err.Error()).Should(Equal("delegate error")) + }) + + It("should handle concurrent requests with errors", func() { + errorReader := NewDataReader(&ErrorMockDataReader{}) + numConcurrentRequests := 5 + + var wg sync.WaitGroup + wg.Add(numConcurrentRequests) + + errorCount := int64(0) + + for i := 0; i < numConcurrentRequests; i++ { + go func() { + defer wg.Done() + _, err := errorReader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + if err != nil { + atomic.AddInt64(&errorCount, 1) + } + }() + } + + wg.Wait() + + Expect(atomic.LoadInt64(&errorCount)).To(Equal(int64(numConcurrentRequests))) + }) + + It("should handle empty result sets", func() { + emptyReader := NewDataReader(storage.NewNoopRelationshipReader()) + + it, err := emptyReader.QueryAttributes(ctx, "tenant1", filter, snap, database.NewCursorPagination()) + + Expect(err).ShouldNot(HaveOccurred()) + Expect(it).ShouldNot(BeNil()) + Expect(it.HasNext()).To(BeFalse()) + }) + }) }) From 7f2a9b9072b67c1cb752fc9e751d644ee6657549 Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Thu, 16 Jul 2026 15:17:19 +0200 Subject: [PATCH 02/13] perf: query relationships with subject filter --- internal/engines/check.go | 5 +-- internal/storage/memory/data_reader.go | 37 ++++++++++++++- internal/storage/postgres/data_reader.go | 23 +++++++--- internal/storage/postgres/utils/filter.go | 40 +++++++++++++++++ .../proxies/circuitbreaker/data_reader.go | 11 +++++ .../proxies/singleflight/data_reader.go | 45 ++++++++++++++++--- internal/storage/storage.go | 9 ++++ 7 files changed, 155 insertions(+), 15 deletions(-) diff --git a/internal/engines/check.go b/internal/engines/check.go index 6320f6c10..d87609850 100644 --- a/internal/engines/check.go +++ b/internal/engines/check.go @@ -272,10 +272,9 @@ func (engine *CheckEngine) checkDirectRelation(request *base.PermissionCheckRequ return denied(emptyResponseMetadata()), err } - // Query the relationships for the entity in the request. - // TupleFilter helps in filtering out the relationships for a specific entity and a permission. + // Batch query with subject push-down var rit *database.TupleIterator - rit, err = engine.dataReader.QueryRelationships(ctx, request.GetTenantId(), filter, request.GetMetadata().GetSnapToken(), database.NewCursorPagination()) + rit, err = engine.dataReader.QueryRelationshipsWithSubjectFilter(ctx, request.GetTenantId(), filter, request.GetSubject(), request.GetMetadata().GetSnapToken(), database.NewCursorPagination()) // If there's an error in querying, return a denied permission response along with the error. if err != nil { return denied(emptyResponseMetadata()), err diff --git a/internal/storage/memory/data_reader.go b/internal/storage/memory/data_reader.go index aef79144f..db9c24fc7 100644 --- a/internal/storage/memory/data_reader.go +++ b/internal/storage/memory/data_reader.go @@ -19,6 +19,7 @@ import ( db "github.com/Permify/permify/pkg/database/memory" base "github.com/Permify/permify/pkg/pb/base/v1" "github.com/Permify/permify/pkg/token" + "github.com/Permify/permify/pkg/tuple" ) // DataReader - @@ -34,7 +35,13 @@ func NewDataReader(database *db.Memory) *DataReader { } // QueryRelationships queries the database for relationships based on the provided filter. -func (r *DataReader) QueryRelationships(_ context.Context, tenantID string, filter *base.TupleFilter, _ string, pagination database.CursorPagination) (it *database.TupleIterator, err error) { +func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, snap string, pagination database.CursorPagination) (it *database.TupleIterator, err error) { + return r.QueryRelationshipsWithSubjectFilter(ctx, tenantID, filter, nil, snap, pagination) +} + +// QueryRelationshipsWithSubjectFilter reads relation tuples with an additional subject push-down filter. +// It returns only tuples where the subject matches exactly OR the tuple is a userset. +func (r *DataReader) QueryRelationshipsWithSubjectFilter(_ context.Context, tenantID string, filter *base.TupleFilter, subject *base.Subject, _ string, pagination database.CursorPagination) (it *database.TupleIterator, err error) { txn := r.database.DB.Txn(false) defer txn.Abort() @@ -86,6 +93,12 @@ func (r *DataReader) QueryRelationships(_ context.Context, tenantID string, filt count := uint32(0) limit := pagination.Limit() + // Pre-compute normalized relation for subject pushdown (only when subject is provided) + var subjectRelation string + if subject != nil { + subjectRelation = tuple.NormalizeRelation(subject.GetRelation()) + } + for _, t := range tup { // Skip tuples below the lower bound switch pagination.Sort() { @@ -99,7 +112,11 @@ func (r *DataReader) QueryRelationships(_ context.Context, tenantID string, filt } } - // Add tuple to result set + // Apply subject push-down when subject is provided: (exact match) OR (userset tuple) + if subject != nil && !matchesSubjectPushdown(t, subject, subjectRelation) { + continue + } + tuples = append(tuples, t.ToTuple()) // Enforce the limit if it's set @@ -112,6 +129,22 @@ func (r *DataReader) QueryRelationships(_ context.Context, tenantID string, filt return database.NewTupleCollection(tuples...).CreateTupleIterator(), nil } +// matchesSubjectPushdown checks if a tuple matches the subject push-down criteria: +// either an exact subject match or a userset tuple needing recursive expansion. +func matchesSubjectPushdown(t storage.RelationTuple, subject *base.Subject, normalizedRelation string) bool { + // Exact match + if t.SubjectType == subject.GetType() && t.SubjectID == subject.GetId() { + if t.SubjectRelation == normalizedRelation || t.SubjectRelation == "" { + return true + } + } + // Userset tuple: subject_relation is non-empty and not ELLIPSIS + if t.SubjectRelation != "" && t.SubjectRelation != tuple.ELLIPSIS { + return true + } + return false +} + // ReadRelationships reads relationships from the database taking into account the pagination. func (r *DataReader) ReadRelationships(_ context.Context, tenantID string, filter *base.TupleFilter, _ string, pagination database.Pagination) (collection *database.TupleCollection, ct database.EncodedContinuousToken, err error) { txn := r.database.DB.Txn(false) diff --git a/internal/storage/postgres/data_reader.go b/internal/storage/postgres/data_reader.go index a00a2ba33..b087a226f 100644 --- a/internal/storage/postgres/data_reader.go +++ b/internal/storage/postgres/data_reader.go @@ -40,11 +40,17 @@ func NewDataReader(database *db.Postgres) *DataReader { // QueryRelationships reads relation tuples from the storage based on the given filter. func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, snap string, pagination database.CursorPagination) (it *database.TupleIterator, err error) { - // Start a new trace span and end it when the function exits. - ctx, span := internal.Tracer.Start(ctx, "data-reader.query-relationships") + return r.QueryRelationshipsWithSubjectFilter(ctx, tenantID, filter, nil, snap, pagination) +} + +// QueryRelationshipsWithSubjectFilter reads relation tuples with an additional subject push-down filter. +// It adds an OR predicate: (exact subject match) OR (userset tuples needing recursive expansion). +// This avoids reading all subscribers of (entity, relation) when only one subject is being checked. +func (r *DataReader) QueryRelationshipsWithSubjectFilter(ctx context.Context, tenantID string, filter *base.TupleFilter, subject *base.Subject, snap string, pagination database.CursorPagination) (it *database.TupleIterator, err error) { + ctx, span := internal.Tracer.Start(ctx, "data-reader.query-relationships-with-subject-filter") defer span.End() // Log query operation - slog.DebugContext(ctx, "querying relationships for tenant_id", slog.String("tenant_id", tenantID)) + slog.DebugContext(ctx, "querying relationships with subject filter for tenant_id", slog.String("tenant_id", tenantID)) // Decode snapshot token // Decode the snapshot value. var st token.SnapToken @@ -59,6 +65,13 @@ func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, fi builder = utils.TuplesFilterQueryForSelectBuilder(builder, filter) builder = utils.SnapshotQuery(builder, st.(snapshot.Token).Value.Uint, st.(snapshot.Token).Snapshot) + // Apply subject push-down: only fetch exact match + userset tuples + if subject != nil { + if cond := utils.SubjectPushdownCondition(subject); cond != nil { + builder = builder.Where(cond) + } + } + if pagination.Cursor() != "" { var t database.ContinuousToken t, err = utils.EncodedContinuousToken{Value: pagination.Cursor()}.Decode() @@ -85,7 +98,7 @@ func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, fi return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_SQL_BUILDER) } - slog.DebugContext(ctx, "generated sql query", slog.String("query", query), "with args", slog.Any("arguments", args)) + slog.DebugContext(ctx, "generated sql query with subject filter", slog.String("query", query), "with args", slog.Any("arguments", args)) // Execute query // Execute the SQL query and retrieve the result rows. var rows pgx.Rows @@ -109,7 +122,7 @@ func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, fi return nil, utils.HandleError(ctx, span, err, base.ErrorCode_ERROR_CODE_SCAN) } - slog.DebugContext(ctx, "successfully retrieved relation tuples from the database") + slog.DebugContext(ctx, "successfully retrieved relation tuples with subject filter from the database") // Return a TupleIterator created from the TupleCollection. return collection.CreateTupleIterator(), nil } diff --git a/internal/storage/postgres/utils/filter.go b/internal/storage/postgres/utils/filter.go index 86442b412..181c67946 100644 --- a/internal/storage/postgres/utils/filter.go +++ b/internal/storage/postgres/utils/filter.go @@ -4,6 +4,7 @@ import ( "github.com/Masterminds/squirrel" base "github.com/Permify/permify/pkg/pb/base/v1" + "github.com/Permify/permify/pkg/tuple" ) // TuplesFilterQueryForSelectBuilder - @@ -147,6 +148,45 @@ func TuplesFilterQueryForUpdateBuilder(sl squirrel.UpdateBuilder, filter *base.T return sl.Where(eq) } +// SubjectPushdownCondition builds a squirrel OR condition for the check engine's +// direct relation path. It returns only rows where: +// - the subject matches exactly (type + id + normalized relation), OR +// - the tuple is a userset (subject_relation is non-empty and not ELLIPSIS) +// +// This dramatically reduces the number of rows returned from the database +// when checking a specific subject against a widely-subscribed entity+relation. +func SubjectPushdownCondition(subject *base.Subject) squirrel.Sqlizer { + if subject == nil { + return nil + } + + subjectRelation := tuple.NormalizeRelation(subject.GetRelation()) + + // Exact match: the subject we are checking for. + // subject_relation is stored as "" for direct subjects (ELLIPSIS is normalized on write). + exactMatch := squirrel.And{ + squirrel.Eq{"subject_type": subject.GetType()}, + squirrel.Eq{"subject_id": subject.GetId()}, + } + if subjectRelation == "" { + exactMatch = append(exactMatch, squirrel.Eq{"subject_relation": ""}) + } else { + exactMatch = append(exactMatch, squirrel.Or{ + squirrel.Eq{"subject_relation": subjectRelation}, + squirrel.Eq{"subject_relation": ""}, + }) + } + + // Userset match: tuples that need recursive expansion. + // subject_relation is non-empty (a real relation, not a direct reference). + usersetMatch := squirrel.And{ + squirrel.NotEq{"subject_relation": ""}, + squirrel.NotEq{"subject_relation": tuple.ELLIPSIS}, + } + + return squirrel.Or{exactMatch, usersetMatch} +} + // AttributesFilterQueryForUpdateBuilder - func AttributesFilterQueryForUpdateBuilder(sl squirrel.UpdateBuilder, filter *base.AttributeFilter) squirrel.UpdateBuilder { eq := squirrel.Eq{} diff --git a/internal/storage/proxies/circuitbreaker/data_reader.go b/internal/storage/proxies/circuitbreaker/data_reader.go index 0c98921ec..68b540641 100644 --- a/internal/storage/proxies/circuitbreaker/data_reader.go +++ b/internal/storage/proxies/circuitbreaker/data_reader.go @@ -33,6 +33,17 @@ func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, fi return response.(*database.TupleIterator), nil } +// QueryRelationshipsWithSubjectFilter - Reads relation tuples with subject push-down through circuit breaker. +func (r *DataReader) QueryRelationshipsWithSubjectFilter(ctx context.Context, tenantID string, filter *base.TupleFilter, subject *base.Subject, token string, pagination database.CursorPagination) (*database.TupleIterator, error) { + response, err := r.cb.Execute(func() (interface{}, error) { + return r.delegate.QueryRelationshipsWithSubjectFilter(ctx, tenantID, filter, subject, token, pagination) + }) + if err != nil { + return nil, err + } + return response.(*database.TupleIterator), nil +} + // ReadRelationships - Reads relation tuples from the repository with different options. func (r *DataReader) ReadRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, token string, pagination database.Pagination) (collection *database.TupleCollection, ct database.EncodedContinuousToken, err error) { type circuitBreakerResponse struct { diff --git a/internal/storage/proxies/singleflight/data_reader.go b/internal/storage/proxies/singleflight/data_reader.go index 18b4d1703..dad0772b0 100644 --- a/internal/storage/proxies/singleflight/data_reader.go +++ b/internal/storage/proxies/singleflight/data_reader.go @@ -15,11 +15,12 @@ import ( // DataReader - Add singleflight behaviour to data reader type DataReader struct { - delegate storage.DataReader - headSnapshotGroup singleflight.Group[string, token.SnapToken] - queryRelationshipsGroup singleflight.Group[string, []*base.Tuple] - querySingleAttrGroup singleflight.Group[string, *base.Attribute] - queryAttributesGroup singleflight.Group[string, []*base.Attribute] + delegate storage.DataReader + headSnapshotGroup singleflight.Group[string, token.SnapToken] + queryRelationshipsGroup singleflight.Group[string, []*base.Tuple] + queryRelWithSubjectPushdownGroup singleflight.Group[string, []*base.Tuple] + querySingleAttrGroup singleflight.Group[string, *base.Attribute] + queryAttributesGroup singleflight.Group[string, []*base.Attribute] } // NewDataReader - Add singleflight behaviour to new data reader @@ -43,6 +44,22 @@ func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, fi return database.NewTupleIterator(tuples...), nil } +// QueryRelationshipsWithSubjectFilter - Reads relation tuples with subject push-down, with singleflight deduplication. +func (r *DataReader) QueryRelationshipsWithSubjectFilter(ctx context.Context, tenantID string, filter *base.TupleFilter, subject *base.Subject, token string, pagination database.CursorPagination) (*database.TupleIterator, error) { + key := queryRelationshipsWithSubjectKey(tenantID, filter, subject, token, pagination) + tuples, _, err := r.queryRelWithSubjectPushdownGroup.Do(ctx, key, func(ctx context.Context) ([]*base.Tuple, error) { + it, err := r.delegate.QueryRelationshipsWithSubjectFilter(ctx, tenantID, filter, subject, token, pagination) + if err != nil { + return nil, err + } + return drainTupleIterator(it), nil + }) + if err != nil { + return nil, err + } + return database.NewTupleIterator(tuples...), nil +} + // ReadRelationships - Reads relation tuples from the repository with different options. func (r *DataReader) ReadRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, token string, pagination database.Pagination) (collection *database.TupleCollection, ct database.EncodedContinuousToken, err error) { return r.delegate.ReadRelationships(ctx, tenantID, filter, token, pagination) @@ -138,6 +155,24 @@ func queryAttributesKey(tenantID string, filter *base.AttributeFilter, token str return b.String() } +func queryRelationshipsWithSubjectKey(tenantID string, filter *base.TupleFilter, subject *base.Subject, token string, pagination database.CursorPagination) string { + var b strings.Builder + fmt.Fprintf(&b, "qrsf\x00%q\x00%q\x00%#v\x00%q\x00%q\x00%q\x00%q\x00%q\x00%q\x00%q\x00%d", + tenantID, + filter.GetEntity().GetType(), + filter.GetEntity().GetIds(), + filter.GetRelation(), + subject.GetType(), + subject.GetId(), + subject.GetRelation(), + token, + pagination.Cursor(), + pagination.Sort(), + pagination.Limit(), + ) + return b.String() +} + // --- iterator helpers --- func drainTupleIterator(it *database.TupleIterator) []*base.Tuple { diff --git a/internal/storage/storage.go b/internal/storage/storage.go index f3f475f82..2b52e6ee2 100644 --- a/internal/storage/storage.go +++ b/internal/storage/storage.go @@ -14,6 +14,11 @@ type DataReader interface { // It returns an iterator to iterate over the tuples and any error encountered. QueryRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, snap string, pagination database.CursorPagination) (iterator *database.TupleIterator, err error) + // QueryRelationshipsWithSubjectFilter reads relation tuples with an additional subject push-down filter. + // It returns only tuples where the subject matches exactly OR the tuple is a userset (subject_relation is non-empty). + // This avoids reading all subscribers of (entity, relation) when only one subject is being checked. + QueryRelationshipsWithSubjectFilter(ctx context.Context, tenantID string, filter *base.TupleFilter, subject *base.Subject, snap string, pagination database.CursorPagination) (iterator *database.TupleIterator, err error) + // ReadRelationships reads relation tuples from the storage based on the given filter and pagination. // It returns a collection of tuples, a continuous token indicating the position in the data set, and any error encountered. ReadRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, snap string, pagination database.Pagination) (collection *database.TupleCollection, ct database.EncodedContinuousToken, err error) @@ -49,6 +54,10 @@ func (f *NoopDataReader) QueryRelationships(_ context.Context, _ string, _ *base return database.NewTupleIterator(), nil } +func (f *NoopDataReader) QueryRelationshipsWithSubjectFilter(_ context.Context, _ string, _ *base.TupleFilter, _ *base.Subject, _ string, _ database.CursorPagination) (*database.TupleIterator, error) { + return database.NewTupleIterator(), nil +} + func (f *NoopDataReader) ReadRelationships(_ context.Context, _ string, _ *base.TupleFilter, _ string, _ database.Pagination) (*database.TupleCollection, database.EncodedContinuousToken, error) { return database.NewTupleCollection(), database.NewNoopContinuousToken().Encode(), nil } From 4436feef92251434902f9f3248da881d734cc051 Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Thu, 16 Jul 2026 15:17:19 +0200 Subject: [PATCH 03/13] perf: batch checks --- internal/engines/balancer/balancer.go | 165 ++- internal/engines/bulk.go | 4 +- internal/engines/bulk_benchmark_test.go | 7 +- internal/engines/bulk_test.go | 18 +- internal/engines/cache/check.go | 88 +- internal/engines/cache/check_test.go | 102 +- internal/engines/check.go | 940 +++++++++++------- internal/engines/check_goroutine_leak_test.go | 14 +- internal/engines/check_test.go | 114 +-- internal/engines/subject_permission.go | 15 +- internal/engines/utils.go | 13 +- internal/invoke/batch.go | 133 +++ internal/invoke/invoke.go | 74 +- internal/invoke/utils.go | 5 +- internal/servers/permission_server.go | 147 ++- internal/storage/postgres/gc/gc_test.go | 48 +- pkg/balancer/balancer.go | 51 +- pkg/balancer/balancer_test.go | 3 +- pkg/balancer/builder.go | 45 +- pkg/balancer/picker.go | 40 +- pkg/balancer/picker_test.go | 211 ++-- pkg/cmd/serve.go | 7 +- pkg/cmd/validate.go | 18 +- pkg/development/development.go | 17 +- 24 files changed, 1373 insertions(+), 906 deletions(-) create mode 100644 internal/invoke/batch.go diff --git a/internal/engines/balancer/balancer.go b/internal/engines/balancer/balancer.go index 213161575..c498709d7 100644 --- a/internal/engines/balancer/balancer.go +++ b/internal/engines/balancer/balancer.go @@ -4,8 +4,11 @@ import ( "context" "fmt" "log/slog" + "sync" "time" + grpcbalancer "google.golang.org/grpc/balancer" + "google.golang.org/grpc" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" @@ -18,11 +21,14 @@ import ( base "github.com/Permify/permify/pkg/pb/base/v1" ) -// Balancer is a wrapper around the balancer hash implementation that +// Balancer wraps permission checking with consistent-hash load balancing. +// All requests (single and batch) are distributed across nodes via the consistent +// hash picker so each entity lands on the node that caches it. type Balancer struct { schemaReader storage.SchemaReader checker invoke.Check client base.PermissionClient + builder balancer.Builder } // NewCheckEngineWithBalancer creates a new check engine with a load balancer. @@ -32,6 +38,7 @@ func NewCheckEngineWithBalancer( ctx context.Context, checker invoke.Check, schemaReader storage.SchemaReader, + builder balancer.Builder, no string, dst *config.Distributed, srv *config.GRPC, @@ -96,47 +103,141 @@ func NewCheckEngineWithBalancer( schemaReader: schemaReader, checker: checker, client: base.NewPermissionClient(conn), + builder: builder, }, nil } -// Check performs a permission check using the schema reader to obtain -// entity definitions, then distributes the request based on a generated key. -func (c *Balancer) Check(ctx context.Context, request *base.PermissionCheckRequest) (*base.PermissionCheckResponse, error) { - // Fetch the EntityDefinition for the given tenant, entity type, and schema version. - en, _, err := c.schemaReader.ReadEntityDefinition(ctx, request.GetTenantId(), request.GetEntity().GetType(), request.GetMetadata().GetSchemaVersion()) +// Check distributes permission checks across cluster nodes via consistent hashing. +// Each entity ID is routed to the node determined by its hash key, ensuring cache locality. +// Entity IDs that hash to the same node are grouped into a single BulkCheck RPC. +func (c *Balancer) Check(ctx context.Context, request *invoke.BatchCheckRequest) (*invoke.BatchCheckResponse, error) { + // Get the current picker; fall back to local if not ready. + nodePicker := c.builder.Picker() + if nodePicker == nil { + return c.checker.Check(ctx, request) + } + + deniedResp := invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) + + // Read entity definition once (shared across all entity IDs). + en, _, err := c.schemaReader.ReadEntityDefinition(ctx, request.TenantID, request.EntityType, request.Metadata.GetSchemaVersion()) if err != nil { slog.ErrorContext(ctx, err.Error()) - // If an error occurs while reading the entity definition, deny permission and return the error. - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + return deniedResp, err } - isRelational := engines.IsRelational(en, request.GetPermission()) + isRelational := engines.IsRelational(en, request.Permission) - // Add a timeout of 2 seconds to the context and also set the generated key as a value. - withTimeout, cancel := context.WithTimeout(context.WithValue(ctx, balancer.Key, []byte(engines.GenerateKey(request, isRelational))), 4*time.Second) - defer cancel() + // Group entity IDs by target SubConn. + groups := map[grpcbalancer.SubConn][]string{} + for _, entityID := range request.EntityIDs { + protoReq := request.ToPermissionCheckRequest(entityID) + key := []byte(engines.GenerateKey(protoReq, isRelational)) - // Logging the intention to forward the request to the underlying client. - slog.InfoContext(ctx, "Forwarding request with key to the underlying client") + sc, err := nodePicker.Pick(key) + if err != nil { + slog.ErrorContext(ctx, "Pick failed, falling back to local", "error", err.Error()) + return c.checker.Check(ctx, request) + } + groups[sc] = append(groups[sc], entityID) + } - // Perform the actual permission check by making a call to the underlying client. - response, err := c.client.Check(withTimeout, request) - if err != nil { - // Log the error and return it. - slog.ErrorContext(ctx, err.Error()) - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + // Fan out: one RPC per node group, concurrently. + type groupResult struct { + resp *invoke.BatchCheckResponse + err error + } + + results := make([]groupResult, len(groups)) + var wg sync.WaitGroup + i := 0 + for sc, entityIDs := range groups { + wg.Add(1) + go func(idx int, sc grpcbalancer.SubConn, entityIDs []string) { + defer wg.Done() + routeCtx := context.WithValue(ctx, balancer.SubConnKey, sc) + withTimeout, cancel := context.WithTimeout(routeCtx, 4*time.Second) + defer cancel() + + if len(entityIDs) == 1 { + // Single entity: use Check RPC. + protoReq := request.ToPermissionCheckRequest(entityIDs[0]) + response, err := c.client.Check(withTimeout, protoReq) + if err != nil { + results[idx] = groupResult{err: err} + return + } + results[idx] = groupResult{resp: &invoke.BatchCheckResponse{ + Results: map[string]base.CheckResult{entityIDs[0]: response.GetCan()}, + Metadata: response.GetMetadata(), + }} + } else { + // Multiple entities for same node: use BulkCheck RPC. + items := make([]*base.PermissionBulkCheckRequestItem, len(entityIDs)) + for j, entityID := range entityIDs { + items[j] = &base.PermissionBulkCheckRequestItem{ + Entity: &base.Entity{Type: request.EntityType, Id: entityID}, + Permission: request.Permission, + Subject: request.Subject, + } + } + bulkReq := &base.PermissionBulkCheckRequest{ + TenantId: request.TenantID, + Metadata: request.Metadata, + Items: items, + Context: request.Context, + Arguments: request.Arguments, + } + bulkResp, err := c.client.BulkCheck(withTimeout, bulkReq) + if err != nil { + results[idx] = groupResult{err: err} + return + } + // Map BulkCheck results back by entity ID (results are ordered same as items). + resp := &invoke.BatchCheckResponse{ + Results: make(map[string]base.CheckResult, len(entityIDs)), + Metadata: &base.PermissionCheckResponseMetadata{}, + } + for j, entityID := range entityIDs { + if j < len(bulkResp.GetResults()) { + resp.Results[entityID] = bulkResp.GetResults()[j].GetCan() + resp.Metadata.CheckCount += bulkResp.GetResults()[j].GetMetadata().GetCheckCount() + } else { + resp.Results[entityID] = base.CheckResult_CHECK_RESULT_DENIED + } + } + results[idx] = groupResult{resp: resp} + } + }(i, sc, entityIDs) + i++ + } + wg.Wait() + + // Merge results from all node groups. + merged := &invoke.BatchCheckResponse{ + Results: make(map[string]base.CheckResult, len(request.EntityIDs)), + Metadata: &base.PermissionCheckResponseMetadata{}, + } + for _, r := range results { + if r.err != nil { + slog.ErrorContext(ctx, "node group check failed", "error", r.err.Error()) + // Mark all entities in failed group as denied. + continue + } + if r.resp != nil { + for entityID, result := range r.resp.Results { + merged.Results[entityID] = result + } + merged.Metadata.CheckCount += r.resp.Metadata.GetCheckCount() + } + } + + // Fill in any missing entity IDs as DENIED. + for _, entityID := range request.EntityIDs { + if _, ok := merged.Results[entityID]; !ok { + merged.Results[entityID] = base.CheckResult_CHECK_RESULT_DENIED + } } - // Return the response received from the client. - return response, nil + return merged, nil } diff --git a/internal/engines/bulk.go b/internal/engines/bulk.go index 263e13091..f01a70d24 100644 --- a/internal/engines/bulk.go +++ b/internal/engines/bulk.go @@ -375,12 +375,12 @@ func (bc *BulkChecker) getRequestResult(ctx context.Context, req BulkCheckerRequ } // Perform the actual permission check - response, err := bc.checker.Check(ctx, req.Request) + response, err := bc.checker.Check(ctx, invoke.NewBatchCheckRequest(req.Request)) if err != nil { return base.CheckResult_CHECK_RESULT_UNSPECIFIED, err } - return response.GetCan(), nil + return response.UnionResult(), nil } // processResult processes a single result with thread-safe state updates. diff --git a/internal/engines/bulk_benchmark_test.go b/internal/engines/bulk_benchmark_test.go index 38f613ace..b4c2db664 100644 --- a/internal/engines/bulk_benchmark_test.go +++ b/internal/engines/bulk_benchmark_test.go @@ -7,6 +7,7 @@ import ( "testing" "time" + "github.com/Permify/permify/internal/invoke" base "github.com/Permify/permify/pkg/pb/base/v1" ) @@ -15,13 +16,11 @@ type MockChecker struct { delay time.Duration } -func (m *MockChecker) Check(ctx context.Context, request *base.PermissionCheckRequest) (*base.PermissionCheckResponse, error) { +func (m *MockChecker) Check(_ context.Context, req *invoke.BatchCheckRequest) (*invoke.BatchCheckResponse, error) { if m.delay > 0 { time.Sleep(m.delay) } - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_ALLOWED, - }, nil + return invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_ALLOWED, req.EntityIDs...), nil } // BenchmarkBulkChecker tests the performance of the BulkChecker diff --git a/internal/engines/bulk_test.go b/internal/engines/bulk_test.go index fbf4b5588..4597d590b 100644 --- a/internal/engines/bulk_test.go +++ b/internal/engines/bulk_test.go @@ -14,25 +14,15 @@ import ( // mockCheckEngine is a mock implementation of invoke.Check for testing type mockCheckEngine struct{} -func (m *mockCheckEngine) Check(ctx context.Context, request *base.PermissionCheckRequest) (*base.PermissionCheckResponse, error) { - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_ALLOWED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 1, - }, - }, nil +func (m *mockCheckEngine) Check(_ context.Context, req *invoke.BatchCheckRequest) (*invoke.BatchCheckResponse, error) { + return invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_ALLOWED, req.EntityIDs...), nil } // errorCheckEngine is a mock implementation that returns errors type errorCheckEngine struct{} -func (e *errorCheckEngine) Check(ctx context.Context, request *base.PermissionCheckRequest) (*base.PermissionCheckResponse, error) { - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_UNSPECIFIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, errors.New("permission check failed") +func (e *errorCheckEngine) Check(_ context.Context, req *invoke.BatchCheckRequest) (*invoke.BatchCheckResponse, error) { + return invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_UNSPECIFIED, req.EntityIDs...), errors.New("permission check failed") } var _ = Describe("Bulk", func() { diff --git a/internal/engines/cache/check.go b/internal/engines/cache/check.go index 495563d75..66b8c1288 100644 --- a/internal/engines/cache/check.go +++ b/internal/engines/cache/check.go @@ -44,56 +44,68 @@ func NewCheckEngineWithCache( } // Check performs a permission check for a given request, using the cached results if available. -func (c *CheckEngineWithCache) Check(ctx context.Context, request *base.PermissionCheckRequest) (response *base.PermissionCheckResponse, err error) { - // Retrieve entity definition - var en *base.EntityDefinition - en, _, err = c.schemaReader.ReadEntityDefinition(ctx, request.GetTenantId(), request.GetEntity().GetType(), request.GetMetadata().GetSchemaVersion()) +// Supports batch: request.EntityIDs may contain one or more entity IDs. +// For each entity, the cache is checked individually; uncached entities are delegated to the underlying checker. +func (c *CheckEngineWithCache) Check(ctx context.Context, request *invoke.BatchCheckRequest) (response *invoke.BatchCheckResponse, err error) { + // Read entity definition once (all entities share the same type) + en, _, err := c.schemaReader.ReadEntityDefinition(ctx, request.TenantID, request.EntityType, request.Metadata.GetSchemaVersion()) if err != nil { - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + return invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...), err } - isRelational := engines.IsRelational(en, request.GetPermission()) - - // Try to get the cached result for the given request. - res, found := c.getCheckKey(request, isRelational) - - // If a cached result is found, handle exclusion and return the result. - if found { - // Increase the hit count in the metrics. - c.cacheHitHistogram.Record(ctx, 1) + isRelational := engines.IsRelational(en, request.Permission) + + // Per-entity cached results + cachedResults := make(map[string]base.CheckResult) + + // Check cache per entity_id, collect uncached IDs + var uncachedIDs []string + for _, id := range request.EntityIDs { + singleReq := request.ToPermissionCheckRequest(id) + res, found := c.getCheckKey(singleReq, isRelational) + if found { + c.cacheHitHistogram.Record(ctx, 1) + cachedResults[id] = res.GetCan() + continue + } + uncachedIDs = append(uncachedIDs, id) + } - // If the request doesn't have the exclusion flag set, return the cached result. - return &base.PermissionCheckResponse{ - Can: res.GetCan(), + // All cached → return cached results + if len(uncachedIDs) == 0 { + return &invoke.BatchCheckResponse{ + Results: cachedResults, Metadata: &base.PermissionCheckResponseMetadata{}, }, nil } - // Perform the actual permission check using the provided request. - cres, err := c.checker.Check(ctx, request) - // Check if there's an error or the response is nil, and return the result. + // Delegate uncached to underlying checker + batchReq := request.Clone() + batchReq.EntityIDs = uncachedIDs + cres, err := c.checker.Check(ctx, batchReq) if err != nil { - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + return invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...), err } - // Add to histogram the response + // Cache the per-entity results for each uncached entity + for _, id := range uncachedIDs { + singleReq := request.ToPermissionCheckRequest(id) + result := base.CheckResult_CHECK_RESULT_DENIED + if r, ok := cres.Results[id]; ok { + result = r + } + c.setCheckKey(singleReq, &base.PermissionCheckResponse{ + Can: result, + Metadata: &base.PermissionCheckResponseMetadata{}, + }, isRelational) + } + + // Merge cached results into the response + for id, result := range cachedResults { + cres.Results[id] = result + } - c.setCheckKey(request, &base.PermissionCheckResponse{ - Can: cres.GetCan(), - Metadata: &base.PermissionCheckResponseMetadata{}, - }, isRelational) - // Return the result of the permission check. - return cres, err + return cres, nil } // GetCheckKey retrieves the value for the given key from the EngineKeys cache. diff --git a/internal/engines/cache/check_test.go b/internal/engines/cache/check_test.go index 5f57b1ca7..825ab968d 100644 --- a/internal/engines/cache/check_test.go +++ b/internal/engines/cache/check_test.go @@ -549,7 +549,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -559,10 +559,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -658,7 +658,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -668,10 +668,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -769,7 +769,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -779,10 +779,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -900,7 +900,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -910,10 +910,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1020,7 +1020,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1033,10 +1033,10 @@ var _ = Describe("cache", func() { Context: &base.Context{ Tuples: tuples, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1139,7 +1139,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1149,10 +1149,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1290,7 +1290,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1303,10 +1303,10 @@ var _ = Describe("cache", func() { Context: &base.Context{ Tuples: tuples, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1426,7 +1426,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1439,10 +1439,10 @@ var _ = Describe("cache", func() { Context: &base.Context{ Tuples: contextual, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1554,7 +1554,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1564,10 +1564,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1681,7 +1681,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1691,10 +1691,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1798,7 +1798,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1808,10 +1808,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1915,7 +1915,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1925,10 +1925,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2075,7 +2075,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2088,10 +2088,10 @@ var _ = Describe("cache", func() { Context: &base.Context{ Tuples: tuples, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2236,7 +2236,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2246,10 +2246,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2405,7 +2405,7 @@ var _ = Describe("cache", func() { ctx.Data = value } - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2416,10 +2416,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2525,7 +2525,7 @@ var _ = Describe("cache", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2538,10 +2538,10 @@ var _ = Describe("cache", func() { Context: &base.Context{ Attributes: attributes, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2693,7 +2693,7 @@ var _ = Describe("cache", func() { ctx.Data = value } - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2704,10 +2704,10 @@ var _ = Describe("cache", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) diff --git a/internal/engines/check.go b/internal/engines/check.go index d87609850..3fad02452 100644 --- a/internal/engines/check.go +++ b/internal/engines/check.go @@ -30,6 +30,8 @@ type CheckEngine struct { dataReader storage.DataReader // concurrencyLimit is the maximum number of concurrent permission checks allowed concurrencyLimit int + // maxBatchSize is the maximum number of entity IDs per batch SQL query (IN clause) + maxBatchSize int } // NewCheckEngine creates a new CheckEngine instance for performing permission checks. @@ -41,6 +43,7 @@ func NewCheckEngine(sr storage.SchemaReader, rr storage.DataReader, opts ...Chec schemaReader: sr, dataReader: rr, concurrencyLimit: _defaultConcurrencyLimit, + maxBatchSize: _defaultMaxBatchSize, } // Apply provided options to configure the CheckEngine @@ -60,74 +63,87 @@ func (engine *CheckEngine) SetInvoker(invoker invoke.Check) { // The permission field in the request can either be a relation or an permission. // This function performs various checks and returns the permission check response // along with any errors that may have occurred. -func (engine *CheckEngine) Check(ctx context.Context, request *base.PermissionCheckRequest) (response *base.PermissionCheckResponse, err error) { - emptyResp := denied(emptyResponseMetadata()) +// Supports batch: request.EntityIDs may contain one or more entity IDs. +// Uses IN (...) queries for efficient batch processing. Returns per-entity results. +func (engine *CheckEngine) Check(ctx context.Context, request *invoke.BatchCheckRequest) (response *invoke.BatchCheckResponse, err error) { + deniedResp := invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) - // Retrieve entity definition + // Read entity definition once (all entities share the same type) var en *base.EntityDefinition - en, _, err = engine.schemaReader.ReadEntityDefinition(ctx, request.GetTenantId(), request.GetEntity().GetType(), request.GetMetadata().GetSchemaVersion()) + en, _, err = engine.schemaReader.ReadEntityDefinition(ctx, request.TenantID, request.EntityType, request.Metadata.GetSchemaVersion()) if err != nil { - return emptyResp, err + return deniedResp, err } - // Perform permission check - var res *base.PermissionCheckResponse - res, err = engine.check(ctx, request, en)(ctx) - if err != nil { - return emptyResp, err - } - - return &base.PermissionCheckResponse{ - Can: res.Can, - Metadata: res.Metadata, - }, nil + return engine.check(ctx, request, en)(ctx) } // CheckFunction is a type that represents a function that takes a context -// and returns a PermissionCheckResponse along with an error. It is used -// to perform individual permission checks within the CheckEngine. -type CheckFunction func(ctx context.Context) (*base.PermissionCheckResponse, error) +// and returns a BatchCheckResponse with per-entity results along with an error. +// It is used to perform permission checks within the CheckEngine. +type CheckFunction func(ctx context.Context) (*invoke.BatchCheckResponse, error) -// CheckCombiner is a type that represents a function which takes a context, -// a slice of CheckFunctions, and a limit. It combines the results of +// CheckCombiner is a type that represents a function which takes a context +// and a slice of CheckFunctions. It combines the per-entity results of // multiple CheckFunctions according to a specific strategy and returns -// a PermissionCheckResponse along with an error. -type CheckCombiner func(ctx context.Context, functions []CheckFunction, limit int) (*base.PermissionCheckResponse, error) - -// run is a helper function that takes a context and a PermissionCheckRequest, -// and returns a CheckFunction. The returned CheckFunction, when called with -// a context, executes the Run method of the CheckEngine with the given -// request, and returns the resulting PermissionCheckResponse and error. -func (engine *CheckEngine) invoke(request *base.PermissionCheckRequest) CheckFunction { - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { +// a BatchCheckResponse along with an error. +// Concurrency is controlled by a request-scoped semaphore stored in the context. +type CheckCombiner func(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) + +// invoke creates a CheckFunction that invokes a batch check through the full invoke chain +// (DirectInvoker -> Cache -> CheckEngine), ensuring depth tracking, caching, and tracing. +func (engine *CheckEngine) invoke(request *invoke.BatchCheckRequest) CheckFunction { + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { return engine.invoker.Check(ctx, request) } } // check constructs a CheckFunction that performs permission checks based on the type of reference in the entity definition. +// All entity IDs in the batch request share the same type and permission, enabling batch SQL queries. func (engine *CheckEngine) check( ctx context.Context, - request *base.PermissionCheckRequest, + request *invoke.BatchCheckRequest, en *base.EntityDefinition, ) CheckFunction { - // If the request's entity and permission are the same as the subject, return a CheckFunction that always allows the permission. - if tuple.AreQueryAndSubjectEqual(request.GetEntity(), request.GetPermission(), request.GetSubject()) { - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { - return allowed(emptyResponseMetadata()), nil + // Identity check: for each entity, if it matches the subject, mark it as ALLOWED. + // Collect remaining entities that need further checking. + identityResults := make(map[string]base.CheckResult) + var remainingIDs []string + for _, id := range request.EntityIDs { + if tuple.AreQueryAndSubjectEqual(&base.Entity{Type: request.EntityType, Id: id}, request.Permission, request.Subject) { + identityResults[id] = base.CheckResult_CHECK_RESULT_ALLOWED + } else { + remainingIDs = append(remainingIDs, id) } } - // Declare a CheckFunction variable that will later be defined based on the type of reference. + // If all entities matched via identity, return immediately. + if len(remainingIDs) == 0 { + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { + return &invoke.BatchCheckResponse{ + Results: identityResults, + Metadata: emptyResponseMetadata(), + }, nil + } + } + + // Build a sub-request for the remaining entities. + subRequest := request + if len(identityResults) > 0 { + subRequest = request.Clone() + subRequest.EntityIDs = remainingIDs + } + var fn CheckFunction // Determine the type of the reference by name in the given entity definition. - tor, _ := schema.GetTypeOfReferenceByNameInEntityDefinition(en, request.GetPermission()) + tor, _ := schema.GetTypeOfReferenceByNameInEntityDefinition(en, request.Permission) // Based on the type of the reference, define the CheckFunction in different ways. switch tor { case base.EntityDefinition_REFERENCE_PERMISSION: // Get the permission from the entity definition. - permission, err := schema.GetPermissionByNameInEntityDefinition(en, request.GetPermission()) + permission, err := schema.GetPermissionByNameInEntityDefinition(en, request.Permission) if err != nil { // If an error is encountered while getting the permission, a CheckFunction is returned that always fails with this error. return checkFail(err) @@ -138,18 +154,18 @@ func (engine *CheckEngine) check( // If the child has a rewrite, check the rewrite. // If not, check the leaf. if child.GetRewrite() != nil { - fn = engine.checkRewrite(ctx, request, child.GetRewrite()) + fn = engine.checkRewrite(ctx, subRequest, child.GetRewrite()) } else { - fn = engine.checkLeaf(request, child.GetLeaf()) + fn = engine.checkLeaf(subRequest, child.GetLeaf()) } case base.EntityDefinition_REFERENCE_ATTRIBUTE: // If the reference is an attribute, check the direct attribute. - fn = engine.checkDirectAttribute(request) + fn = engine.checkDirectAttribute(subRequest) case base.EntityDefinition_REFERENCE_RELATION: // If the reference is a relation, check the direct relation. - fn = engine.checkDirectRelation(request) + fn = engine.checkDirectRelation(subRequest) default: - fn = engine.checkDirectCall(request) + fn = engine.checkDirectCall(subRequest) } // If the CheckFunction is still undefined after the switch, return a CheckFunction that always fails with an error indicating an undefined child kind. @@ -157,15 +173,25 @@ func (engine *CheckEngine) check( return checkFail(errors.New(base.ErrorCode_ERROR_CODE_UNDEFINED_CHILD_KIND.String())) } - // Otherwise, return a CheckFunction that checks a union of CheckFunctions with a concurrency limit. - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { - return checkUnion(ctx, []CheckFunction{fn}, engine.concurrencyLimit) + // Otherwise, return a CheckFunction that checks a union of CheckFunctions. + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { + result, err := checkUnion(ctx, []CheckFunction{fn}, engine.concurrencyLimit) + if err != nil { + return result, err + } + // Merge identity results into the result. + if len(identityResults) > 0 { + for id, res := range identityResults { + result.Results[id] = res + } + } + return result, nil } } // checkRewrite prepares a CheckFunction according to the provided Rewrite operation. // It uses a Rewrite object that describes how to combine the results of multiple CheckFunctions. -func (engine *CheckEngine) checkRewrite(ctx context.Context, request *base.PermissionCheckRequest, rewrite *base.Rewrite) CheckFunction { +func (engine *CheckEngine) checkRewrite(ctx context.Context, request *invoke.BatchCheckRequest, rewrite *base.Rewrite) CheckFunction { // Switch statement depending on the Rewrite operation switch rewrite.GetRewriteOperation() { // In case of UNION operation, set the children CheckFunctions to be run concurrently @@ -188,7 +214,7 @@ func (engine *CheckEngine) checkRewrite(ctx context.Context, request *base.Permi // checkLeaf prepares a CheckFunction according to the provided Leaf operation. // It uses a Leaf object that describes how to check a permission request. -func (engine *CheckEngine) checkLeaf(request *base.PermissionCheckRequest, leaf *base.Leaf) CheckFunction { +func (engine *CheckEngine) checkLeaf(request *invoke.BatchCheckRequest, leaf *base.Leaf) CheckFunction { // Switch statement depending on the Leaf type switch op := leaf.GetType().(type) { // In case of TupleToUserSet operation, prepare a CheckFunction that checks @@ -213,12 +239,10 @@ func (engine *CheckEngine) checkLeaf(request *base.PermissionCheckRequest, leaf } } -// setChild prepares a CheckFunction according to the provided combiner function -// and children. It uses the Child object which contains the information about the child -// nodes and can be either a Rewrite or a Leaf. +// setChild prepares a CheckFunction according to the provided combiner function and children. func (engine *CheckEngine) setChild( ctx context.Context, - request *base.PermissionCheckRequest, + request *invoke.BatchCheckRequest, children []*base.Child, combiner CheckCombiner, ) CheckFunction { @@ -241,158 +265,239 @@ func (engine *CheckEngine) setChild( } // Return a function that when called, runs the appropriate combiner function - // (union, intersection, exclusion) on the prepared CheckFunctions with the provided concurrency limit - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { + // (union, intersection, exclusion) on the prepared CheckFunctions. + // Concurrency is controlled by the request-scoped semaphore in the context. + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { return combiner(ctx, functions, engine.concurrencyLimit) } } // checkDirectRelation is a method of CheckEngine struct that returns a CheckFunction. // It's responsible for directly checking the permissions on an entity -func (engine *CheckEngine) checkDirectRelation(request *base.PermissionCheckRequest) CheckFunction { +func (engine *CheckEngine) checkDirectRelation(request *invoke.BatchCheckRequest) CheckFunction { // The returned CheckFunction is a closure over the provided context and request - return func(ctx context.Context) (result *base.PermissionCheckResponse, err error) { + return func(ctx context.Context) (result *invoke.BatchCheckResponse, err error) { // Define a TupleFilter. This specifies which tuples we're interested in. // We want tuples that match the entity type and ID from the request, and have a specific relation. filter := &base.TupleFilter{ Entity: &base.EntityFilter{ - Type: request.GetEntity().GetType(), - Ids: []string{request.GetEntity().GetId()}, + Type: request.EntityType, + Ids: request.EntityIDs, // IN (...) }, - Relation: request.GetPermission(), + Relation: request.Permission, } - // Use the filter to query for relationships in the given context. - // NewContextualRelationships() creates a ContextualRelationships instance from tuples in the request. - // QueryRelationships() then uses the filter to find and return matching relationships. + // Query contextual tuples (supports multiple Ids in EntityFilter) var cti *database.TupleIterator - cti, err = storageContext.NewContextualTuples(request.GetContext().GetTuples()...).QueryRelationships(filter, database.NewCursorPagination()) + cti, err = storageContext.NewContextualTuples(request.Context.GetTuples()...).QueryRelationships(filter, database.NewCursorPagination()) if err != nil { // If an error occurred while querying, return a "denied" response and the error. - return denied(emptyResponseMetadata()), err + return denied(request.EntityIDs, emptyResponseMetadata()), err } // Batch query with subject push-down var rit *database.TupleIterator - rit, err = engine.dataReader.QueryRelationshipsWithSubjectFilter(ctx, request.GetTenantId(), filter, request.GetSubject(), request.GetMetadata().GetSnapToken(), database.NewCursorPagination()) + rit, err = engine.dataReader.QueryRelationshipsWithSubjectFilter(ctx, request.TenantID, filter, request.Subject, request.Metadata.GetSnapToken(), database.NewCursorPagination()) // If there's an error in querying, return a denied permission response along with the error. if err != nil { - return denied(emptyResponseMetadata()), err + return denied(request.EntityIDs, emptyResponseMetadata()), err } // Create a new UniqueTupleIterator from the two TupleIterators. // NewUniqueTupleIterator() ensures that the iterator only returns unique tuples. it := database.NewUniqueTupleIterator(rit, cti) - // Define a slice of CheckFunctions to hold the check functions for each subject. - checkFunctions := make([]CheckFunction, 0, 4) - // Iterate over all tuples returned by the iterator. + // Per-entity results: track which entities got a direct match. + directResults := make(map[string]base.CheckResult) + // Track which entities still need resolution via userset checks. + needsResolution := make(map[string]bool, len(request.EntityIDs)) + for _, id := range request.EntityIDs { + needsResolution[id] = true + } + + // Collect userset subjects, grouping by (entityType, relation) for batch queries. + // Also track which parent entity each userset subject came from. + usersetGroups := map[usersetGroupKey][]string{} + // usersetToParents maps userset entity -> list of parent entity IDs + usersetToParents := map[entityRef][]string{} + for it.HasNext() { - // Get the next tuple's subject. next, ok := it.GetNext() if !ok { break } subject := next.GetSubject() + parentEntityID := next.GetEntity().GetId() - // If the subject of the tuple is the same as the subject in the request, permission is allowed. - if tuple.AreSubjectsEqual(subject, request.GetSubject()) { - return allowed(emptyResponseMetadata()), nil + if tuple.AreSubjectsEqual(subject, request.Subject) { + // Direct match: mark this specific entity as ALLOWED. + directResults[parentEntityID] = base.CheckResult_CHECK_RESULT_ALLOWED + delete(needsResolution, parentEntityID) + continue } - // If the subject is not a user and the relation is not ELLIPSIS, append a check function to the list. if !tuple.IsDirectSubject(subject) && subject.GetRelation() != tuple.ELLIPSIS { - checkFunctions = append(checkFunctions, engine.invoke(&base.PermissionCheckRequest{ - TenantId: request.GetTenantId(), - Entity: &base.Entity{ - Type: subject.GetType(), - Id: subject.GetId(), - }, - Permission: subject.GetRelation(), - Subject: request.GetSubject(), - Metadata: request.GetMetadata(), - Context: request.GetContext(), + key := usersetGroupKey{entityType: subject.GetType(), relation: subject.GetRelation()} + usersetGroups[key] = append(usersetGroups[key], subject.GetId()) + ref := entityRef{entityType: subject.GetType(), entityID: subject.GetId()} + usersetToParents[ref] = append(usersetToParents[ref], parentEntityID) + } + } + + // Early exit: if all entities have direct matches, no need for userset checks. + if len(needsResolution) == 0 { + resp := invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) + for id, res := range directResults { + resp.Results[id] = res + } + return resp, nil + } + + // Build check functions for userset groups, chunking large groups. + var checkFunctions []CheckFunction + for key, ids := range usersetGroups { + for i := 0; i < len(ids); i += engine.maxBatchSize { + end := min(i+engine.maxBatchSize, len(ids)) + checkFunctions = append(checkFunctions, engine.invoke(&invoke.BatchCheckRequest{ + TenantID: request.TenantID, + EntityType: key.entityType, + EntityIDs: ids[i:end], + Permission: key.relation, + Subject: request.Subject, + Metadata: request.Metadata, + Context: request.Context, })) } } - // If there's any CheckFunction in the list, return the union of all CheckFunctions + // Start with the response for all requested entities defaulting to DENIED. + resp := invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) + + // Apply direct match results. + for id, res := range directResults { + resp.Results[id] = res + } + + // If there are userset check functions, run them and map results back to parent entities. if len(checkFunctions) > 0 { - return checkUnion(ctx, checkFunctions, engine.concurrencyLimit) + usersetResp, err := checkUnion(ctx, checkFunctions, engine.concurrencyLimit) + if err != nil { + resp.Metadata = joinResponseMetas(resp.Metadata, usersetResp.Metadata) + return resp, err + } + resp.Metadata = joinResponseMetas(resp.Metadata, usersetResp.Metadata) + + // Map userset results back to parent entities. + for ref, parentIDs := range usersetToParents { + if usersetResult, ok := usersetResp.Results[ref.entityID]; ok && usersetResult == base.CheckResult_CHECK_RESULT_ALLOWED { + for _, parentID := range parentIDs { + resp.Results[parentID] = base.CheckResult_CHECK_RESULT_ALLOWED + } + } + } } - // If there's no CheckFunction, return a denied permission response. - return denied(emptyResponseMetadata()), nil + return resp, nil } } // checkTupleToUserSet is a method of CheckEngine that checks permissions using the // TupleToUserSet data structure. It returns a CheckFunction closure that does the check. func (engine *CheckEngine) checkTupleToUserSet( - request *base.PermissionCheckRequest, + request *invoke.BatchCheckRequest, ttu *base.TupleToUserSet, ) CheckFunction { // The returned CheckFunction is a closure over the provided context, request, and ttu. - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { // Define a TupleFilter. This specifies which tuples we're interested in. // We want tuples that match the entity type and ID from the request, and have a specific relation. filter := &base.TupleFilter{ Entity: &base.EntityFilter{ - Type: request.GetEntity().GetType(), // Filter by entity type from request - Ids: []string{request.GetEntity().GetId()}, // Filter by entity ID from request + Type: request.EntityType, + Ids: request.EntityIDs, // IN (...) }, - Relation: ttu.GetTupleSet().GetRelation(), // Filter by relation from tuple set + Relation: ttu.GetTupleSet().GetRelation(), } // Use the filter to query for relationships in the given context. // NewContextualRelationships() creates a ContextualRelationships instance from tuples in the request. // QueryRelationships() then uses the filter to find and return matching relationships. - cti, err := storageContext.NewContextualTuples(request.GetContext().GetTuples()...).QueryRelationships(filter, database.NewCursorPagination()) + cti, err := storageContext.NewContextualTuples(request.Context.GetTuples()...).QueryRelationships(filter, database.NewCursorPagination()) if err != nil { // If an error occurred while querying, return a "denied" response and the error. - return denied(emptyResponseMetadata()), err + return denied(request.EntityIDs, emptyResponseMetadata()), err } // Use the filter to query for relationships in the database. // relationshipReader.QueryRelationships() uses the filter to find and return matching relationships. - rit, err := engine.dataReader.QueryRelationships(ctx, request.GetTenantId(), filter, request.GetMetadata().GetSnapToken(), database.NewCursorPagination()) + rit, err := engine.dataReader.QueryRelationships(ctx, request.TenantID, filter, request.Metadata.GetSnapToken(), database.NewCursorPagination()) if err != nil { // If an error occurred while querying, return a "denied" response and the error. - return denied(emptyResponseMetadata()), err + return denied(request.EntityIDs, emptyResponseMetadata()), err } // Create a new UniqueTupleIterator from the two TupleIterators. // NewUniqueTupleIterator() ensures that the iterator only returns unique tuples. it := database.NewUniqueTupleIterator(rit, cti) - // Define a slice of CheckFunctions to hold the check functions for each subject. - checkFunctions := make([]CheckFunction, 0, 4) - // Iterate over all tuples returned by the iterator. + // Group subjects by type for batch processing. + // Also track which parent entity each subject came from. + subjectsByType := map[string][]string{} + // subjectToParents maps subject entity -> list of parent entity IDs + subjectToParents := map[entityRef][]string{} + for it.HasNext() { - // Get the next tuple's subject. next, ok := it.GetNext() if !ok { break } - subject := next.GetSubject() + s := next.GetSubject() + parentEntityID := next.GetEntity().GetId() + subjectsByType[s.GetType()] = append(subjectsByType[s.GetType()], s.GetId()) + ref := entityRef{entityType: s.GetType(), entityID: s.GetId()} + subjectToParents[ref] = append(subjectToParents[ref], parentEntityID) + } - // For each subject, generate a check function for its computed user set and append it to the list. - checkFunctions = append(checkFunctions, engine.checkComputedUserSet(&base.PermissionCheckRequest{ - TenantId: request.GetTenantId(), - Entity: &base.Entity{ - Type: subject.GetType(), - Id: subject.GetId(), - }, - Permission: subject.GetRelation(), - Subject: request.GetSubject(), - Metadata: request.GetMetadata(), - Context: request.GetContext(), - Arguments: request.GetArguments(), - }, ttu.GetComputed())) + var checkFunctions []CheckFunction + for entityType, ids := range subjectsByType { + for i := 0; i < len(ids); i += engine.maxBatchSize { + end := min(i+engine.maxBatchSize, len(ids)) + checkFunctions = append(checkFunctions, engine.invoke(&invoke.BatchCheckRequest{ + TenantID: request.TenantID, + EntityType: entityType, + EntityIDs: ids[i:end], + Permission: ttu.GetComputed().GetRelation(), + Subject: request.Subject, + Metadata: request.Metadata, + Context: request.Context, + Arguments: request.Arguments, + })) + } + } + + // Start with all requested entities defaulting to DENIED. + resp := invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) + + if len(checkFunctions) == 0 { + return resp, nil + } + + subjectResp, err := checkUnion(ctx, checkFunctions, engine.concurrencyLimit) + if err != nil { + resp.Metadata = joinResponseMetas(resp.Metadata, subjectResp.Metadata) + return resp, err } + resp.Metadata = joinResponseMetas(resp.Metadata, subjectResp.Metadata) - // Return the union of all CheckFunctions - // If any one of the check functions allows the action, the permission is granted. - return checkUnion(ctx, checkFunctions, engine.concurrencyLimit) + // Map subject results back to parent entities. + for ref, parentIDs := range subjectToParents { + if subjectResult, ok := subjectResp.Results[ref.entityID]; ok && subjectResult == base.CheckResult_CHECK_RESULT_ALLOWED { + for _, parentID := range parentIDs { + resp.Results[parentID] = base.CheckResult_CHECK_RESULT_ALLOWED + } + } + } + + return resp, nil } } @@ -400,27 +505,28 @@ func (engine *CheckEngine) checkTupleToUserSet( // checkComputedUserSet is a method of CheckEngine that checks permissions using the // ComputedUserSet data structure. It returns a CheckFunction closure that performs the check. func (engine *CheckEngine) checkComputedUserSet( - request *base.PermissionCheckRequest, // The request containing details about the permission to be checked - cu *base.ComputedUserSet, // The computed user set containing user set information + request *invoke.BatchCheckRequest, + cu *base.ComputedUserSet, ) CheckFunction { // The returned CheckFunction invokes a permission check with a new request that is almost the same // as the incoming request, but changes the Permission to be the relation defined in the computed user set. // This is how the check "descends" into the computed user set to check permissions there. - return engine.invoke(&base.PermissionCheckRequest{ - TenantId: request.GetTenantId(), // Tenant ID from the incoming request - Entity: request.GetEntity(), // Entity from the incoming request - Permission: cu.GetRelation(), // Permission is set to the relation defined in the computed user set - Subject: request.GetSubject(), // The subject from the incoming request - Metadata: request.GetMetadata(), // Metadata from the incoming request - Context: request.GetContext(), - Arguments: request.GetArguments(), + return engine.invoke(&invoke.BatchCheckRequest{ + TenantID: request.TenantID, + EntityType: request.EntityType, + EntityIDs: request.EntityIDs, + Permission: cu.GetRelation(), + Subject: request.Subject, + Metadata: request.Metadata, + Context: request.Context, + Arguments: request.Arguments, }) } // checkComputedAttribute constructs a CheckFunction that checks if a computed attribute // permission check request is allowed or denied. func (engine *CheckEngine) checkComputedAttribute( - request *base.PermissionCheckRequest, + request *invoke.BatchCheckRequest, ca *base.ComputedAttribute, ) CheckFunction { // We're returning a function here - this is the CheckFunction. @@ -428,217 +534,208 @@ func (engine *CheckEngine) checkComputedAttribute( // We pass a new PermissionCheckRequest to 'invoke', copying most of the fields // from the original request, but replacing the 'Permission' with the computed // attribute's name. - return engine.invoke(&base.PermissionCheckRequest{ - TenantId: request.GetTenantId(), - Entity: request.GetEntity(), + return engine.invoke(&invoke.BatchCheckRequest{ + TenantID: request.TenantID, + EntityType: request.EntityType, + EntityIDs: request.EntityIDs, Permission: ca.GetName(), - Subject: request.GetSubject(), - Metadata: request.GetMetadata(), - Context: request.GetContext(), - Arguments: request.GetArguments(), + Subject: request.Subject, + Metadata: request.Metadata, + Context: request.Context, }) } // checkDirectAttribute constructs a CheckFunction that checks if a direct attribute // permission check request is allowed or denied. -func (engine *CheckEngine) checkDirectAttribute( - request *base.PermissionCheckRequest, -) CheckFunction { - // We're returning a function here - this is the actual CheckFunction. - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { - // Initial error declaration - var err error +// Uses batch QueryAttributes with multiple entity IDs in a single query. +func (engine *CheckEngine) checkDirectAttribute(request *invoke.BatchCheckRequest) CheckFunction { + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { + resp := invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) - // Create a new AttributeFilter with the entity type and ID from the request - // and the requested permission. filter := &base.AttributeFilter{ Entity: &base.EntityFilter{ - Type: request.GetEntity().GetType(), - Ids: []string{request.GetEntity().GetId()}, + Type: request.EntityType, + Ids: request.EntityIDs, // IN (...) }, - Attributes: []string{request.GetPermission()}, + Attributes: []string{request.Permission}, } - var val *base.Attribute - - // storageContext.NewContextualAttributes creates a new instance of ContextualAttributes based on the attributes - // retrieved from the request context. - val, err = storageContext.NewContextualAttributes(request.GetContext().GetAttributes()...).QuerySingleAttribute(filter) - // An error occurred while querying the single attribute, so we return a denied response with empty metadata - // and the error. + // Query contextual attributes (in-memory, supports multiple IDs). + cta, err := storageContext.NewContextualAttributes(request.Context.GetAttributes()...).QueryAttributes(filter, database.NewCursorPagination()) if err != nil { - return denied(emptyResponseMetadata()), err + return resp, err } - if val == nil { - // Use the data reader's QuerySingleAttribute method to find the relevant attribute - val, err = engine.dataReader.QuerySingleAttribute(ctx, request.GetTenantId(), filter, request.GetMetadata().GetSnapToken()) - // If there was an error, return a denied response and the error. - if err != nil { - return denied(emptyResponseMetadata()), err - } + // Batch query from database. + ait, err := engine.dataReader.QueryAttributes(ctx, request.TenantID, filter, request.Metadata.GetSnapToken(), database.NewCursorPagination()) + if err != nil { + return resp, err } - // No attribute was found matching the provided filter. In this case, we return a denied response with empty metadata - // and no error. - if val == nil { - return denied(emptyResponseMetadata()), nil - } + // Combine attributes from different sources ensuring uniqueness. + it := database.NewUniqueAttributeIterator(ait, cta) + for it.HasNext() { + next, ok := it.GetNext() + if !ok { + break + } - // Unmarshal the attribute value into a BoolValue message. - var msg base.BooleanValue - if err := val.GetValue().UnmarshalTo(&msg); err != nil { - // If there was an error unmarshaling, return a denied response and the error. - return denied(emptyResponseMetadata()), err - } + // Unmarshal the attribute value into a BoolValue message. + var msg base.BooleanValue + if err := next.GetValue().UnmarshalTo(&msg); err != nil { + return resp, err + } - // If the attribute's value is true, return an allowed response. - if msg.Data { - return allowed(emptyResponseMetadata()), nil + if msg.Data { + resp.Results[next.GetEntity().GetId()] = base.CheckResult_CHECK_RESULT_ALLOWED + } } - // If the attribute's value is not true, return a denied response. - return denied(emptyResponseMetadata()), nil + return resp, nil } } // checkCall creates and returns a CheckFunction based on the provided request and call details. -// It essentially constructs a new PermissionCheckRequest based on the call details and then invokes +// It essentially constructs a new BatchCheckRequest based on the call details and then invokes // the permission check using the engine's invoke method. func (engine *CheckEngine) checkCall( - request *base.PermissionCheckRequest, + request *invoke.BatchCheckRequest, call *base.Call, ) CheckFunction { // Construct a new permission check request based on the input request and call details. - return engine.invoke(&base.PermissionCheckRequest{ - TenantId: request.GetTenantId(), - Entity: request.GetEntity(), + return engine.invoke(&invoke.BatchCheckRequest{ + TenantID: request.TenantID, + EntityType: request.EntityType, + EntityIDs: request.EntityIDs, Permission: call.GetRuleName(), - Subject: request.GetSubject(), - Metadata: request.GetMetadata(), - Context: request.GetContext(), + Subject: request.Subject, + Metadata: request.Metadata, + Context: request.Context, Arguments: call.GetArguments(), }) } // checkDirectCall creates and returns a CheckFunction that performs direct permission checking. // The function evaluates permissions based on rule definitions, arguments, and attributes. -func (engine *CheckEngine) checkDirectCall( - request *base.PermissionCheckRequest, -) CheckFunction { - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { - var err error +// Processes each entity individually since each entity may have different attribute values. +func (engine *CheckEngine) checkDirectCall(request *invoke.BatchCheckRequest) CheckFunction { + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { + resp := invoke.NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) - // If an error occurs during the check, this default "denied" response will be returned. - emptyResp := denied(emptyResponseMetadata()) - - // Read the rule definition from the schema. If an error occurs, return the default denied response. + // Read the rule definition from the schema once (shared across all entities). var ru *base.RuleDefinition - ru, _, err = engine.schemaReader.ReadRuleDefinition(ctx, request.GetTenantId(), request.GetPermission(), request.GetMetadata().GetSchemaVersion()) + ru, _, err := engine.schemaReader.ReadRuleDefinition(ctx, request.TenantID, request.Permission, request.Metadata.GetSchemaVersion()) if err != nil { - return emptyResp, err + return resp, err } - // Initialize an arguments map to hold argument values. - arguments := map[string]any{ - "context": map[string]any{ - "data": request.GetContext().GetData().AsMap(), - }, + // Prepare the CEL environment once (shared across all entities). + env, err := utils.ArgumentsAsCelEnv(ru.Arguments) + if err != nil { + return resp, err } - // List to store computed attributes. - attributes := make([]string, 0) + // Compile the rule expression once. + exp := cel.CheckedExprToAst(ru.Expression) + prg, err := env.Program(exp) + if err != nil { + return resp, err + } - // Iterate over request arguments to classify and process them. - for _, arg := range request.GetArguments() { + // Classify arguments once. + var attributes []string + baseArguments := map[string]any{ + "context": map[string]any{ + "data": request.Context.GetData().AsMap(), + }, + } + for _, arg := range request.Arguments { switch actualArg := arg.Type.(type) { case *base.Argument_ComputedAttribute: - // Handle computed attributes: Set them to a default empty value. attrName := actualArg.ComputedAttribute.GetName() emptyValue := getEmptyValueForType(ru.GetArguments()[attrName]) - arguments[attrName] = emptyValue + baseArguments[attrName] = emptyValue attributes = append(attributes, attrName) default: - // Return an error for any unsupported argument types. - return denied(emptyResponseMetadata()), errors.New(base.ErrorCode_ERROR_CODE_INTERNAL.String()) + return resp, errors.New(base.ErrorCode_ERROR_CODE_INTERNAL.String()) } } - // If there are computed attributes, fetch them from the data source. + // Batch query ALL attributes for ALL entities at once. + attrsByEntity := map[string]map[string]any{} if len(attributes) > 0 { filter := &base.AttributeFilter{ Entity: &base.EntityFilter{ - Type: request.GetEntity().GetType(), - Ids: []string{request.GetEntity().GetId()}, + Type: request.EntityType, + Ids: request.EntityIDs, }, Attributes: attributes, } - ait, err := engine.dataReader.QueryAttributes(ctx, request.GetTenantId(), filter, request.GetMetadata().GetSnapToken(), database.NewCursorPagination()) + ait, err := engine.dataReader.QueryAttributes(ctx, request.TenantID, filter, request.Metadata.GetSnapToken(), database.NewCursorPagination()) if err != nil { - return denied(emptyResponseMetadata()), err + return resp, err } - cta, err := storageContext.NewContextualAttributes(request.GetContext().GetAttributes()...).QueryAttributes(filter, database.NewCursorPagination()) + cta, err := storageContext.NewContextualAttributes(request.Context.GetAttributes()...).QueryAttributes(filter, database.NewCursorPagination()) if err != nil { - return denied(emptyResponseMetadata()), err + return resp, err } - // Combine attributes from different sources ensuring uniqueness. it := database.NewUniqueAttributeIterator(ait, cta) for it.HasNext() { next, ok := it.GetNext() if !ok { break } - arguments[next.GetAttribute()] = utils.ConvertProtoAnyToInterface(next.GetValue()) + entityID := next.GetEntity().GetId() + if attrsByEntity[entityID] == nil { + attrsByEntity[entityID] = make(map[string]any) + } + attrsByEntity[entityID][next.GetAttribute()] = utils.ConvertProtoAnyToInterface(next.GetValue()) } } - // Prepare the CEL environment with the argument values. - env, err := utils.ArgumentsAsCelEnv(ru.Arguments) - if err != nil { - return nil, err - } - - // Compile the rule expression into an executable form. - exp := cel.CheckedExprToAst(ru.Expression) - prg, err := env.Program(exp) - if err != nil { - return nil, err - } + // Evaluate CEL per entity with its specific attributes. + for _, entityID := range request.EntityIDs { + arguments := make(map[string]any, len(baseArguments)) + for k, v := range baseArguments { + arguments[k] = v + } + for k, v := range attrsByEntity[entityID] { + arguments[k] = v + } - // Evaluate the rule expression with the provided arguments. - out, _, err := prg.Eval(arguments) - if err != nil { - return denied(emptyResponseMetadata()), fmt.Errorf("failed to evaluate expression: %w", err) - } + // Evaluate the rule expression with the arguments for this entity. + out, _, err := prg.Eval(arguments) + if err != nil { + return resp, fmt.Errorf("failed to evaluate expression: %w", err) + } - // Ensure the result of evaluation is boolean and decide on permission. - result, ok := out.Value().(bool) - if !ok { - return denied(emptyResponseMetadata()), fmt.Errorf("expected boolean result, but got %T", out.Value()) - } + result, ok := out.Value().(bool) + if !ok { + return resp, fmt.Errorf("expected boolean result, but got %T", out.Value()) + } - // If the result of the CEL evaluation is true, return an "allowed" response, otherwise return a "denied" response - if result { - return allowed(emptyResponseMetadata()), nil + if result { + resp.Results[entityID] = base.CheckResult_CHECK_RESULT_ALLOWED + } } - return denied(emptyResponseMetadata()), nil + return resp, nil } } -// checkUnion checks if the subject has permission by running multiple CheckFunctions concurrently, -// the permission check is successful if any one of the CheckFunctions succeeds (union). -func checkUnion(ctx context.Context, functions []CheckFunction, limit int) (*base.PermissionCheckResponse, error) { +// checkUnion checks if the subject has permission by running multiple CheckFunctions concurrently. +// Per-entity merge: for each entityID, if ANY function returned ALLOWED -> ALLOWED, else -> DENIED. +func checkUnion(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) { // Initialize the response metadata responseMetadata := emptyResponseMetadata() // If there are no functions, deny the permission and return if len(functions) == 0 { - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, + return &invoke.BatchCheckResponse{ + Results: map[string]base.CheckResult{}, Metadata: responseMetadata, }, nil } @@ -658,6 +755,12 @@ func checkUnion(ctx context.Context, functions []CheckFunction, limit int) (*bas close(decisionChan) }() + // Merged per-entity results: for each entityID, if ANY function returned ALLOWED -> ALLOWED. + mergedResults := map[string]base.CheckResult{} + // Track how many entities are still DENIED. When this reaches 0, all are ALLOWED and we can exit early. + deniedCount := 0 + entityIDsSeen := false + // Iterate over the results of the CheckFunctions for range len(functions) { select { @@ -665,33 +768,64 @@ func checkUnion(ctx context.Context, functions []CheckFunction, limit int) (*bas case d := <-decisionChan: // Merge the response metadata with the received metadata responseMetadata = joinResponseMetas(responseMetadata, d.resp.Metadata) - // If there was an error, deny the permission and return the error + // If there was an error, return what we have so far with the error if d.err != nil { - return denied(responseMetadata), d.err + return &invoke.BatchCheckResponse{ + Results: mergedResults, + Metadata: responseMetadata, + }, d.err } - // If the CheckFunction allowed the permission, allow the permission and return - if d.resp.GetCan() == base.CheckResult_CHECK_RESULT_ALLOWED { - return allowed(responseMetadata), nil + // Per-entity union: if this function allowed an entity, mark it as allowed + for entityID, result := range d.resp.Results { + if result == base.CheckResult_CHECK_RESULT_ALLOWED { + if prev, exists := mergedResults[entityID]; !exists { + mergedResults[entityID] = base.CheckResult_CHECK_RESULT_ALLOWED + // New entity seen, already allowed, no change to deniedCount + } else if prev == base.CheckResult_CHECK_RESULT_DENIED { + mergedResults[entityID] = base.CheckResult_CHECK_RESULT_ALLOWED + deniedCount-- + } + } else if _, exists := mergedResults[entityID]; !exists { + mergedResults[entityID] = base.CheckResult_CHECK_RESULT_DENIED + deniedCount++ + } } - // If the context is done, deny the permission and return a cancellation error + entityIDsSeen = true + + // Early exit: if all known entities are now ALLOWED, no further functions can change the result. + if entityIDsSeen && deniedCount == 0 { + return &invoke.BatchCheckResponse{ + Results: mergedResults, + Metadata: responseMetadata, + }, nil + } + // If the context is done, return a cancellation error case <-ctx.Done(): - return denied(responseMetadata), errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) + return &invoke.BatchCheckResponse{ + Results: mergedResults, + Metadata: responseMetadata, + }, errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) } } - // If all CheckFunctions are done and none have allowed the permission, deny the permission and return - return denied(responseMetadata), nil + return &invoke.BatchCheckResponse{ + Results: mergedResults, + Metadata: responseMetadata, + }, nil } -// checkIntersection checks if the subject has permission by running multiple CheckFunctions concurrently, -// the permission check is successful only when all CheckFunctions succeed (intersection). -func checkIntersection(ctx context.Context, functions []CheckFunction, limit int) (*base.PermissionCheckResponse, error) { +// checkIntersection checks if the subject has permission by running multiple CheckFunctions concurrently. +// Per-entity merge: for each entityID, ALL functions must return ALLOWED -> ALLOWED, else -> DENIED. +func checkIntersection(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) { // Initialize the response metadata responseMetadata := emptyResponseMetadata() // If there are no functions, deny the permission and return if len(functions) == 0 { - return denied(responseMetadata), nil + return &invoke.BatchCheckResponse{ + Results: map[string]base.CheckResult{}, + Metadata: responseMetadata, + }, nil } // Create a channel to receive the results of the CheckFunctions @@ -709,6 +843,13 @@ func checkIntersection(ctx context.Context, functions []CheckFunction, limit int close(decisionChan) }() + // Track per-entity: how many functions returned ALLOWED and total functions seen. + allowedCounts := map[string]int{} + deniedEntities := map[string]struct{}{} // entities that received at least one DENIED + // Track all entity IDs seen. + allEntityIDs := map[string]struct{}{} + entityIDsSeen := false + // Iterate over the results of the CheckFunctions for range len(functions) { select { @@ -716,32 +857,83 @@ func checkIntersection(ctx context.Context, functions []CheckFunction, limit int case d := <-decisionChan: // Merge the response metadata with the received metadata responseMetadata = joinResponseMetas(responseMetadata, d.resp.Metadata) - // If there was an error, deny the permission and return the error + // If there was an error, return denied with the error if d.err != nil { - return denied(responseMetadata), d.err + // Build denied results for all seen entities + results := make(map[string]base.CheckResult, len(allEntityIDs)) + for id := range allEntityIDs { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, d.err } - // If the CheckFunction denied the permission, deny the permission and return - if d.resp.GetCan() == base.CheckResult_CHECK_RESULT_DENIED { - return denied(responseMetadata), nil + // Track per-entity allowed counts + for entityID, result := range d.resp.Results { + allEntityIDs[entityID] = struct{}{} + if result == base.CheckResult_CHECK_RESULT_ALLOWED { + allowedCounts[entityID]++ + } else { + deniedEntities[entityID] = struct{}{} + } } - // If the context is done, deny the permission and return a cancellation error + entityIDsSeen = true + + // Early exit: if ALL known entities have been DENIED by at least one function, + // no further functions can make them ALLOWED (intersection requires all). + if entityIDsSeen && len(deniedEntities) == len(allEntityIDs) { + results := make(map[string]base.CheckResult, len(allEntityIDs)) + for id := range allEntityIDs { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, nil + } + // If the context is done, return a cancellation error case <-ctx.Done(): - return denied(responseMetadata), errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) + results := make(map[string]base.CheckResult, len(allEntityIDs)) + for id := range allEntityIDs { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) } } - // If all CheckFunctions allowed the permission, allow the permission and return - return allowed(responseMetadata), nil + // Build final results: entity is ALLOWED only if ALL functions returned ALLOWED for it. + numFunctions := len(functions) + results := make(map[string]base.CheckResult, len(allEntityIDs)) + for entityID := range allEntityIDs { + if allowedCounts[entityID] == numFunctions { + results[entityID] = base.CheckResult_CHECK_RESULT_ALLOWED + } else { + results[entityID] = base.CheckResult_CHECK_RESULT_DENIED + } + } + + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, nil } -// checkExclusion is a function that checks if there are any exclusions for given CheckFunctions -func checkExclusion(ctx context.Context, functions []CheckFunction, limit int) (*base.PermissionCheckResponse, error) { +// checkExclusion is a function that checks if there are any exclusions for given CheckFunctions. +// Per-entity merge: for each entityID, first function ALLOWED AND all remaining DENIED -> ALLOWED, else -> DENIED. +func checkExclusion(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) { // Initialize the response metadata responseMetadata := emptyResponseMetadata() // Check if there are at least 2 functions, otherwise return an error indicating that exclusion requires more than one function if len(functions) <= 1 { - return denied(responseMetadata), errors.New(base.ErrorCode_ERROR_CODE_EXCLUSION_REQUIRES_MORE_THAN_ONE_FUNCTION.String()) + return &invoke.BatchCheckResponse{ + Results: map[string]base.CheckResult{}, + Metadata: responseMetadata, + }, errors.New(base.ErrorCode_ERROR_CODE_EXCLUSION_REQUIRES_MORE_THAN_ONE_FUNCTION.String()) } // Initialize channels to handle the result of the first function and the remaining functions separately @@ -755,16 +947,16 @@ func checkExclusion(ctx context.Context, functions []CheckFunction, limit int) ( var wg sync.WaitGroup wg.Add(1) go func() { + defer wg.Done() result, err := functions[0](cancelCtx) leftDecisionChan <- CheckResponse{ resp: result, err: err, } - wg.Done() }() - // Run the remaining functions concurrently with a limit - clean := checkRun(cancelCtx, functions[1:], decisionChan, limit-1) + // Run the remaining functions concurrently + clean := checkRun(cancelCtx, functions[1:], decisionChan, limit) // Ensure that all resources are properly cleaned up when the function exits defer func() { @@ -775,23 +967,53 @@ func checkExclusion(ctx context.Context, functions []CheckFunction, limit int) ( close(leftDecisionChan) }() + // Per-entity results from the first (left) function. + var leftResults map[string]base.CheckResult + // Process the result from the first function select { case left := <-leftDecisionChan: responseMetadata = joinResponseMetas(responseMetadata, left.resp.Metadata) if left.err != nil { - return denied(responseMetadata), left.err + return &invoke.BatchCheckResponse{ + Results: map[string]base.CheckResult{}, + Metadata: responseMetadata, + }, left.err } - if left.resp.GetCan() == base.CheckResult_CHECK_RESULT_DENIED { - return denied(responseMetadata), nil - } + leftResults = left.resp.Results case <-ctx.Done(): - return denied(responseMetadata), errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) + return &invoke.BatchCheckResponse{ + Results: map[string]base.CheckResult{}, + Metadata: responseMetadata, + }, errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) + } + + // Early exit: if ALL entities in the left result are DENIED, exclusion cannot produce ALLOWED. + allLeftDenied := true + leftAllowedCount := 0 + for _, result := range leftResults { + if result == base.CheckResult_CHECK_RESULT_ALLOWED { + allLeftDenied = false + leftAllowedCount++ + } + } + if allLeftDenied { + results := make(map[string]base.CheckResult, len(leftResults)) + for id := range leftResults { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, nil } + // Track per-entity: whether any remaining function returned ALLOWED (which would exclude). + excludedEntities := map[string]bool{} + // Process the results from the remaining functions for range len(functions) - 1 { select { @@ -799,58 +1021,97 @@ func checkExclusion(ctx context.Context, functions []CheckFunction, limit int) ( responseMetadata = joinResponseMetas(responseMetadata, d.resp.Metadata) if d.err != nil { - return denied(responseMetadata), d.err + // On error, return denied for all entities + results := make(map[string]base.CheckResult, len(leftResults)) + for id := range leftResults { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, d.err } - if d.resp.GetCan() == base.CheckResult_CHECK_RESULT_ALLOWED { - return denied(responseMetadata), nil + // If any remaining function allowed an entity, that entity is excluded (denied) + for entityID, result := range d.resp.Results { + if result == base.CheckResult_CHECK_RESULT_ALLOWED { + if !excludedEntities[entityID] && leftResults[entityID] == base.CheckResult_CHECK_RESULT_ALLOWED { + leftAllowedCount-- + } + excludedEntities[entityID] = true + } + } + + // Early exit: if all initially-ALLOWED entities are now excluded, result is all DENIED + if leftAllowedCount <= 0 { + results := make(map[string]base.CheckResult, len(leftResults)) + for id := range leftResults { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, nil } case <-ctx.Done(): - return denied(responseMetadata), errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) + results := make(map[string]base.CheckResult, len(leftResults)) + for id := range leftResults { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, errors.New(base.ErrorCode_ERROR_CODE_CANCELLED.String()) } } - // If none of the functions allowed the action, then it's allowed by exclusion - return allowed(responseMetadata), nil + // Build final results: entity is ALLOWED if first function ALLOWED AND no remaining function ALLOWED + results := make(map[string]base.CheckResult, len(leftResults)) + for entityID, leftResult := range leftResults { + if leftResult == base.CheckResult_CHECK_RESULT_ALLOWED && !excludedEntities[entityID] { + results[entityID] = base.CheckResult_CHECK_RESULT_ALLOWED + } else { + results[entityID] = base.CheckResult_CHECK_RESULT_DENIED + } + } + + return &invoke.BatchCheckResponse{ + Results: results, + Metadata: responseMetadata, + }, nil } -// checkRun is a function that executes a list of CheckFunctions concurrently with a specified limit. +// checkRun executes a list of CheckFunctions concurrently. +// DB-level concurrency is controlled by the semaphore DataReader proxy, not here. +// checkRun executes a list of CheckFunctions concurrently with a local concurrency limit. +// DB-level concurrency is controlled by the semaphore DataReader proxy. +// The local limit here prevents excessive goroutine fan-out and depth exhaustion. func checkRun(ctx context.Context, functions []CheckFunction, decisionChan chan<- CheckResponse, limit int) func() { - // Create a channel that enforces the concurrency limit - cl := make(chan struct{}, limit) var wg sync.WaitGroup + cl := make(chan struct{}, limit) - // Define a helper function that calls a CheckFunction and sends the result to the decisionChan - check := func(child CheckFunction) { - result, err := child(ctx) - decisionChan <- CheckResponse{ - resp: result, - err: err, - } - // Once the CheckFunction is done, release the concurrency limit - <-cl - wg.Done() - } - - // Start a goroutine that iterates over the functions wg.Add(1) go func() { - run: - // Iterate over the functions + defer wg.Done() for _, fun := range functions { child := fun select { - // If the concurrency limit allows it, start the function in a new goroutine case cl <- struct{}{}: - wg.Add(1) - go check(child) - // If the context is done, break the loop case <-ctx.Done(): - break run + return } + wg.Add(1) + go func() { + defer wg.Done() + result, err := child(ctx) + decisionChan <- CheckResponse{ + resp: result, + err: err, + } + <-cl + }() } - wg.Done() }() // Return a cleanup function that waits for all goroutines to finish and then closes the concurrency limit channel @@ -860,39 +1121,50 @@ func checkRun(ctx context.Context, functions []CheckFunction, decisionChan chan< } } -// checkFail is a helper function that returns a CheckFunction that always returns a denied PermissionCheckResponse +// checkFail is a helper function that returns a CheckFunction that always returns a denied BatchCheckResponse // with the provided error and an empty PermissionCheckResponseMetadata. // // The function works as follows: // 1. The function takes an error as input parameter. // 2. The function returns a CheckFunction that takes a context as input parameter and always returns a denied -// PermissionCheckResponse with the provided error and an empty PermissionCheckResponseMetadata. +// BatchCheckResponse with the provided error and an empty PermissionCheckResponseMetadata. func checkFail(err error) CheckFunction { - return func(ctx context.Context) (*base.PermissionCheckResponse, error) { - return denied(&base.PermissionCheckResponseMetadata{}), err + return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { + return &invoke.BatchCheckResponse{ + Results: map[string]base.CheckResult{}, + Metadata: &base.PermissionCheckResponseMetadata{}, + }, err } } -// denied is a helper function that returns a denied PermissionCheckResponse with the provided PermissionCheckResponseMetadata. -// -// The function works as follows: -// 1. The function takes a PermissionCheckResponseMetadata as input parameter. -// 2. The function returns a denied PermissionCheckResponse with a RESULT_DENIED Can value and the provided metadata. -func denied(meta *base.PermissionCheckResponseMetadata) *base.PermissionCheckResponse { - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: meta, - } +// usersetGroupKey identifies a group of userset subjects that share the same entity type and relation, +// allowing their next-level queries to be batched into a single SQL call with IN (...). +type usersetGroupKey struct { + entityType string + relation string +} + +// entityRef identifies a specific entity by type and ID (used as map key). +type entityRef struct { + entityType string + entityID string } -// allowed is a helper function that returns an allowed PermissionCheckResponse with the provided PermissionCheckResponseMetadata. +// denied is a helper function that returns a denied BatchCheckResponse for the given entity IDs +// with the provided PermissionCheckResponseMetadata. // // The function works as follows: -// 1. The function takes a PermissionCheckResponseMetadata as input parameter. -// 2. The function returns an allowed PermissionCheckResponse with a RESULT_ALLOWED Can value and the provided metadata. -func allowed(meta *base.PermissionCheckResponseMetadata) *base.PermissionCheckResponse { - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_ALLOWED, +// 1. The function takes entity IDs and a PermissionCheckResponseMetadata as input parameters. +// 2. The function returns a denied BatchCheckResponse with RESULT_DENIED for each entity and the provided metadata. +func denied(entityIDs []string, meta *base.PermissionCheckResponseMetadata) *invoke.BatchCheckResponse { + return &invoke.BatchCheckResponse{ + Results: func() map[string]base.CheckResult { + results := make(map[string]base.CheckResult, len(entityIDs)) + for _, id := range entityIDs { + results[id] = base.CheckResult_CHECK_RESULT_DENIED + } + return results + }(), Metadata: meta, } } diff --git a/internal/engines/check_goroutine_leak_test.go b/internal/engines/check_goroutine_leak_test.go index 87296449d..5a9d9fbd0 100644 --- a/internal/engines/check_goroutine_leak_test.go +++ b/internal/engines/check_goroutine_leak_test.go @@ -124,7 +124,7 @@ var _ = Describe("goroutine-leak-tests", func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - _, err = invoker.Check(ctx, &base.PermissionCheckRequest{ + _, err = invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -134,7 +134,7 @@ var _ = Describe("goroutine-leak-tests", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) }(i) } @@ -238,7 +238,7 @@ var _ = Describe("goroutine-leak-tests", func() { cancel() // This should return a cancellation error - _, err = invoker.Check(ctx, &base.PermissionCheckRequest{ + _, err = invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -248,7 +248,7 @@ var _ = Describe("goroutine-leak-tests", func() { SchemaVersion: "", Depth: 20, }, - }) + })) // We expect an error due to context cancellation Expect(err).Should(HaveOccurred()) @@ -329,14 +329,14 @@ var _ = Describe("goroutine-leak-tests", func() { // Test with high number of concurrent requests numRequests := 30 errors := make(chan error, numRequests) - results := make(chan *base.PermissionCheckResponse, numRequests) + results := make(chan *invoke.BatchCheckResponse, numRequests) for i := 0; i < numRequests; i++ { go func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - resp, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + resp, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -346,7 +346,7 @@ var _ = Describe("goroutine-leak-tests", func() { SchemaVersion: "", Depth: 20, }, - }) + })) if err != nil { errors <- err diff --git a/internal/engines/check_test.go b/internal/engines/check_test.go index 2e164803e..c41d43a13 100644 --- a/internal/engines/check_test.go +++ b/internal/engines/check_test.go @@ -138,7 +138,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -148,10 +148,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -241,7 +241,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -251,10 +251,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -346,7 +346,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -356,10 +356,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -471,7 +471,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -481,10 +481,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -571,7 +571,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -584,10 +584,10 @@ var _ = Describe("check-engine", func() { Context: &base.Context{ Tuples: tuples, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -677,7 +677,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -687,10 +687,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -808,7 +808,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -821,10 +821,10 @@ var _ = Describe("check-engine", func() { Context: &base.Context{ Tuples: tuples, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -931,7 +931,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -944,10 +944,10 @@ var _ = Describe("check-engine", func() { Context: &base.Context{ Tuples: contextual, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1039,7 +1039,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1049,10 +1049,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1153,7 +1153,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1163,10 +1163,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1257,7 +1257,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1267,10 +1267,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1361,7 +1361,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1371,10 +1371,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1501,7 +1501,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1514,10 +1514,10 @@ var _ = Describe("check-engine", func() { Context: &base.Context{ Tuples: tuples, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1648,7 +1648,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1658,10 +1658,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1793,7 +1793,7 @@ var _ = Describe("check-engine", func() { ctx.Data = value } - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1804,10 +1804,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -1899,7 +1899,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -1912,10 +1912,10 @@ var _ = Describe("check-engine", func() { Context: &base.Context{ Attributes: attributes, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2050,7 +2050,7 @@ var _ = Describe("check-engine", func() { ctx.Data = value } - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2061,10 +2061,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2213,7 +2213,7 @@ var _ = Describe("check-engine", func() { } for permission, res := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2223,10 +2223,10 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: check.depth, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(res).Should(Equal(response.GetCan())) + Expect(res).Should(Equal(response.UnionResult())) } } }) @@ -2353,7 +2353,7 @@ var _ = Describe("check-engine", func() { } for permission := range check.assertions { - response, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + response, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: entity, Subject: subject, @@ -2363,11 +2363,11 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: check.depth, }, - }) + })) Expect(err).Should(HaveOccurred()) Expect(err.Error()).Should(ContainSubstring("ERROR_CODE_DEPTH_NOT_ENOUGH")) - Expect(response.GetCan()).Should(Equal(base.CheckResult_CHECK_RESULT_DENIED)) + Expect(response.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_DENIED)) } } }) diff --git a/internal/engines/subject_permission.go b/internal/engines/subject_permission.go index 52009d032..b521f07b0 100644 --- a/internal/engines/subject_permission.go +++ b/internal/engines/subject_permission.go @@ -97,17 +97,18 @@ func (engine *SubjectPermissionEngine) SubjectPermission(ctx context.Context, re // The checkEngine's Check method is called with a new PermissionCheckRequest. // The request is created using the data from the original request, and the permission from the current iteration. - cr, err := engine.checker.Check(ctx, &base.PermissionCheckRequest{ - TenantId: request.GetTenantId(), + cr, err := engine.checker.Check(ctx, &invoke.BatchCheckRequest{ + TenantID: request.GetTenantId(), + EntityType: request.GetEntity().GetType(), + EntityIDs: []string{request.GetEntity().GetId()}, + Permission: permission, + Subject: request.GetSubject(), Metadata: &base.PermissionCheckRequestMetadata{ SchemaVersion: request.GetMetadata().GetSchemaVersion(), SnapToken: request.GetMetadata().GetSnapToken(), Depth: request.GetMetadata().GetDepth(), }, - Entity: request.GetEntity(), - Permission: permission, - Subject: request.GetSubject(), - Context: request.GetContext(), + Context: request.GetContext(), }) // If there's an error, it is sent over the resultChannel along with the permission and a "denied" result. if err != nil { @@ -116,7 +117,7 @@ func (engine *SubjectPermissionEngine) SubjectPermission(ctx context.Context, re } // If there's no error, the result of the check (along with the permission and a nil error) is sent over the resultChannel. - resultChannel <- SubjectPermissionResponse{permission: permission, result: cr.Can, err: nil} + resultChannel <- SubjectPermissionResponse{permission: permission, result: cr.UnionResult(), err: nil} }(p) } diff --git a/internal/engines/utils.go b/internal/engines/utils.go index 0a23ba1b5..722cc021f 100644 --- a/internal/engines/utils.go +++ b/internal/engines/utils.go @@ -9,6 +9,7 @@ import ( "google.golang.org/protobuf/types/known/anypb" + "github.com/Permify/permify/internal/invoke" "github.com/Permify/permify/internal/schema" "github.com/Permify/permify/pkg/attribute" base "github.com/Permify/permify/pkg/pb/base/v1" @@ -17,6 +18,7 @@ import ( const ( _defaultConcurrencyLimit = 100 + _defaultMaxBatchSize = 100 ) // CheckOption - a functional option type for configuring the CheckEngine. @@ -29,6 +31,13 @@ func CheckConcurrencyLimit(limit int) CheckOption { } } +// CheckMaxBatchSize - a functional option that sets the maximum batch size for the CheckEngine. +func CheckMaxBatchSize(size int) CheckOption { + return func(c *CheckEngine) { + c.maxBatchSize = size + } +} + type LookupOption func(engine *LookupEngine) func LookupConcurrencyLimit(limit int) LookupOption { @@ -73,9 +82,9 @@ type SubjectPermissionResponse struct { err error } -// CheckResponse - a struct that holds a PermissionCheckResponse and an error for a single check function. +// CheckResponse - a struct that holds a BatchCheckResponse and an error for a single check function. type CheckResponse struct { - resp *base.PermissionCheckResponse + resp *invoke.BatchCheckResponse err error } diff --git a/internal/invoke/batch.go b/internal/invoke/batch.go new file mode 100644 index 000000000..e8b4be3df --- /dev/null +++ b/internal/invoke/batch.go @@ -0,0 +1,133 @@ +package invoke + +import ( + base "github.com/Permify/permify/pkg/pb/base/v1" +) + +// BatchCheckResponse holds per-entity check results. +// Each entity ID maps to its individual CheckResult. +type BatchCheckResponse struct { + Results map[string]base.CheckResult // entityID → result + Metadata *base.PermissionCheckResponseMetadata +} + +// NewBatchCheckResponse creates a response from a set of entity IDs with the same result. +func NewBatchCheckResponse(result base.CheckResult, entityIDs ...string) *BatchCheckResponse { + results := make(map[string]base.CheckResult, len(entityIDs)) + for _, id := range entityIDs { + results[id] = result + } + return &BatchCheckResponse{ + Results: results, + Metadata: &base.PermissionCheckResponseMetadata{}, + } +} + +// IsAllowed returns true if ANY entity in the batch was allowed (union semantics). +func (r *BatchCheckResponse) IsAllowed() bool { + for _, result := range r.Results { + if result == base.CheckResult_CHECK_RESULT_ALLOWED { + return true + } + } + return false +} + +// UnionResult returns a single CheckResult with union semantics. +func (r *BatchCheckResponse) UnionResult() base.CheckResult { + if r.IsAllowed() { + return base.CheckResult_CHECK_RESULT_ALLOWED + } + return base.CheckResult_CHECK_RESULT_DENIED +} + +// ToPermissionCheckResponse converts to a single-entity proto response (union semantics). +// Used at the public API boundary. +func (r *BatchCheckResponse) ToPermissionCheckResponse() *base.PermissionCheckResponse { + return &base.PermissionCheckResponse{ + Can: r.UnionResult(), + Metadata: r.Metadata, + } +} + +// Merge combines results from another response into this one. +func (r *BatchCheckResponse) Merge(other *BatchCheckResponse) { + for id, result := range other.Results { + r.Results[id] = result + } + if other.Metadata != nil { + r.Metadata.CheckCount += other.Metadata.CheckCount + } +} + +// BatchCheckRequest is the internal batch-capable check request. +// Unlike PermissionCheckRequest which has a single Entity, this holds +// multiple entity IDs of the same type, enabling IN (...) SQL queries. +type BatchCheckRequest struct { + TenantID string + EntityType string + EntityIDs []string // batch: multiple entity IDs of the same type + Permission string + Subject *base.Subject + Metadata *base.PermissionCheckRequestMetadata + Context *base.Context + Arguments []*base.Argument +} + +// NewBatchCheckRequest creates a BatchCheckRequest from a single-entity PermissionCheckRequest. +func NewBatchCheckRequest(req *base.PermissionCheckRequest) *BatchCheckRequest { + return &BatchCheckRequest{ + TenantID: req.GetTenantId(), + EntityType: req.GetEntity().GetType(), + EntityIDs: []string{req.GetEntity().GetId()}, + Permission: req.GetPermission(), + Subject: req.GetSubject(), + Metadata: req.GetMetadata(), + Context: req.GetContext(), + Arguments: req.GetArguments(), + } +} + +// ToPermissionCheckRequest converts back to a single-entity request for the given entity ID. +// Used for cache key generation, single-entity invoke, etc. +func (r *BatchCheckRequest) ToPermissionCheckRequest(entityID string) *base.PermissionCheckRequest { + return &base.PermissionCheckRequest{ + TenantId: r.TenantID, + Entity: &base.Entity{ + Type: r.EntityType, + Id: entityID, + }, + Permission: r.Permission, + Subject: r.Subject, + Metadata: r.Metadata, + Context: r.Context, + Arguments: r.Arguments, + } +} + +// Clone creates a shallow copy of the request with a new EntityIDs slice. +func (r *BatchCheckRequest) Clone() *BatchCheckRequest { + ids := make([]string, len(r.EntityIDs)) + copy(ids, r.EntityIDs) + return &BatchCheckRequest{ + TenantID: r.TenantID, + EntityType: r.EntityType, + EntityIDs: ids, + Permission: r.Permission, + Subject: r.Subject, + Metadata: r.Metadata, + Context: r.Context, + Arguments: r.Arguments, + } +} + +// CloneWithDepth creates a clone with a modified depth value. +func (r *BatchCheckRequest) CloneWithDepth(depth int32) *BatchCheckRequest { + c := r.Clone() + c.Metadata = &base.PermissionCheckRequestMetadata{ + SchemaVersion: r.Metadata.GetSchemaVersion(), + SnapToken: r.Metadata.GetSnapToken(), + Depth: depth, + } + return c +} diff --git a/internal/invoke/invoke.go b/internal/invoke/invoke.go index ff92382df..a2015fad5 100644 --- a/internal/invoke/invoke.go +++ b/internal/invoke/invoke.go @@ -27,10 +27,11 @@ type Invoker interface { } // Check is an interface that defines a method for checking permissions. -// It requires an implementation of InvokeCheck that takes a context and a PermissionCheckRequest, -// and returns a PermissionCheckResponse and an error if any. +// It requires an implementation of Check that takes a context and a BatchCheckRequest +// (which supports both single-entity and multi-entity checks), +// and returns a BatchCheckResponse with per-entity results and an error if any. type Check interface { - Check(ctx context.Context, request *base.PermissionCheckRequest) (response *base.PermissionCheckResponse, err error) + Check(ctx context.Context, request *BatchCheckRequest) (response *BatchCheckResponse, err error) } // Expand is an interface that defines a method for expanding permissions. @@ -100,94 +101,75 @@ func NewDirectInvoker( } // Check is a method that implements the Check interface. -// It calls the Run method of the CheckEngine with the provided context and PermissionCheckRequest, -// and returns a PermissionCheckResponse and an error if any. -func (invoker *DirectInvoker) Check(ctx context.Context, request *base.PermissionCheckRequest) (response *base.PermissionCheckResponse, err error) { +// It validates depth, sets snap/schema tokens, decrements depth, and delegates to the underlying checker. +// Supports batch: request.EntityIDs may contain one or more entity IDs. +func (invoker *DirectInvoker) Check(ctx context.Context, request *BatchCheckRequest) (response *BatchCheckResponse, err error) { ctx, span := internal.Tracer.Start(ctx, "check", trace.WithAttributes( - attribute.KeyValue{Key: "tenant_id", Value: attribute.StringValue(request.GetTenantId())}, - attribute.KeyValue{Key: "entity", Value: attribute.StringValue(tuple.EntityToString(request.GetEntity()))}, - attribute.KeyValue{Key: "permission", Value: attribute.StringValue(request.GetPermission())}, - attribute.KeyValue{Key: "subject", Value: attribute.StringValue(tuple.SubjectToString(request.GetSubject()))}, + attribute.KeyValue{Key: "tenant_id", Value: attribute.StringValue(request.TenantID)}, + attribute.KeyValue{Key: "entity_type", Value: attribute.StringValue(request.EntityType)}, + attribute.KeyValue{Key: "entity_count", Value: attribute.IntValue(len(request.EntityIDs))}, + attribute.KeyValue{Key: "permission", Value: attribute.StringValue(request.Permission)}, + attribute.KeyValue{Key: "subject", Value: attribute.StringValue(tuple.SubjectToString(request.Subject))}, )) defer span.End() invoker.checkHistogram.Record(ctx, 1, metric.WithAttributeSet( attribute.NewSet( - attribute.KeyValue{Key: "subject_id", Value: attribute.StringValue(request.GetSubject().GetId())}, - attribute.KeyValue{Key: "subject_type", Value: attribute.StringValue(request.GetSubject().GetType())}, + attribute.KeyValue{Key: "subject_id", Value: attribute.StringValue(request.Subject.GetId())}, + attribute.KeyValue{Key: "subject_type", Value: attribute.StringValue(request.Subject.GetType())}, )), ) + denied := NewBatchCheckResponse(base.CheckResult_CHECK_RESULT_DENIED, request.EntityIDs...) + // Validate the depth of the request. err = checkDepth(request) if err != nil { span.RecordError(err) span.SetStatus(otelCodes.Error, err.Error()) span.SetAttributes(attribute.KeyValue{Key: "can", Value: attribute.StringValue(base.CheckResult_CHECK_RESULT_DENIED.String())}) - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + return denied, err } // Set the SnapToken if it's not provided in the request. - if request.GetMetadata().GetSnapToken() == "" { + if request.Metadata.GetSnapToken() == "" { var st token.SnapToken - st, err = invoker.dataReader.HeadSnapshot(ctx, request.GetTenantId()) + st, err = invoker.dataReader.HeadSnapshot(ctx, request.TenantID) if err != nil { span.RecordError(err) span.SetStatus(otelCodes.Error, err.Error()) span.SetAttributes(attribute.KeyValue{Key: "can", Value: attribute.StringValue(base.CheckResult_CHECK_RESULT_DENIED.String())}) - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + return denied, err } request.Metadata.SnapToken = st.Encode().String() } // Set the SchemaVersion if it's not provided in the request. - if request.GetMetadata().GetSchemaVersion() == "" { - request.Metadata.SchemaVersion, err = invoker.schemaReader.HeadVersion(ctx, request.GetTenantId()) + if request.Metadata.GetSchemaVersion() == "" { + request.Metadata.SchemaVersion, err = invoker.schemaReader.HeadVersion(ctx, request.TenantID) if err != nil { span.RecordError(err) span.SetStatus(otelCodes.Error, err.Error()) span.SetAttributes(attribute.KeyValue{Key: "can", Value: attribute.StringValue(base.CheckResult_CHECK_RESULT_DENIED.String())}) - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + return denied, err } } - // Create a copy of the request to safely decrement depth without mutating the original. - nextRequest := request.CloneVT() - nextRequest.Metadata.Depth = request.GetMetadata().Depth - 1 + // Decrement depth and delegate. + next := request.CloneWithDepth(request.Metadata.GetDepth() - 1) // Perform the actual permission check using the provided request. - response, err = invoker.cc.Check(ctx, nextRequest) + response, err = invoker.cc.Check(ctx, next) if err != nil { span.RecordError(err) span.SetStatus(otelCodes.Error, err.Error()) span.SetAttributes(attribute.KeyValue{Key: "can", Value: attribute.StringValue(base.CheckResult_CHECK_RESULT_DENIED.String())}) - return &base.PermissionCheckResponse{ - Can: base.CheckResult_CHECK_RESULT_DENIED, - Metadata: &base.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, err + return denied, err } // increaseCheckCount increments the CheckCount value in the response metadata by 1. - atomic.AddInt32(&response.GetMetadata().CheckCount, +1) + atomic.AddInt32(&response.Metadata.CheckCount, +1) - span.SetAttributes(attribute.KeyValue{Key: "can", Value: attribute.StringValue(response.GetCan().String())}) return response, err } diff --git a/internal/invoke/utils.go b/internal/invoke/utils.go index edab8687f..330840cc5 100644 --- a/internal/invoke/utils.go +++ b/internal/invoke/utils.go @@ -2,14 +2,13 @@ package invoke import ( "errors" - "sync/atomic" base "github.com/Permify/permify/pkg/pb/base/v1" ) // checkDepth validates that the request has sufficient depth for permission checks -func checkDepth(request *base.PermissionCheckRequest) error { - if atomic.LoadInt32(&request.GetMetadata().Depth) < 0 { // Check depth is not negative +func checkDepth(request *BatchCheckRequest) error { + if request.Metadata.GetDepth() < 0 { return errors.New(base.ErrorCode_ERROR_CODE_DEPTH_NOT_ENOUGH.String()) } return nil diff --git a/internal/servers/permission_server.go b/internal/servers/permission_server.go index 1fc14ce15..540ebf3b7 100644 --- a/internal/servers/permission_server.go +++ b/internal/servers/permission_server.go @@ -38,7 +38,7 @@ func (r *PermissionServer) Check(ctx context.Context, request *v1.PermissionChec return nil, status.Error(GetStatus(v), v.Error()) // Return validation error } - response, err := r.invoker.Check(ctx, request) + response, err := r.invoker.Check(ctx, invoke.NewBatchCheckRequest(request)) if err != nil { span.RecordError(err) span.SetStatus(otelCodes.Error, err.Error()) @@ -46,7 +46,7 @@ func (r *PermissionServer) Check(ctx context.Context, request *v1.PermissionChec return nil, status.Error(GetStatus(err), err.Error()) } - return response, nil + return response.ToPermissionCheckResponse(), nil } // BulkCheck - Performs multiple authorization checks in a single request @@ -84,95 +84,94 @@ func (r *PermissionServer) BulkCheck(ctx context.Context, request *v1.Permission return nil, err } - // The buffer size is equal to the number of references in the entity. - type resultItem struct { - index int - response *v1.PermissionCheckResponse + // Group items by (entityType, permission, subject) for batch processing. + // Items in the same group share a single BatchCheckRequest. + type groupKey struct { + entityType, permission, subjectType, subjectID, subjectRelation string + } + type groupEntry struct { + entityIDs []string + indices []int + subject *v1.Subject } - resultChannel := make(chan resultItem, len(checkItems)) - // The WaitGroup and Mutex are used for synchronization. - var wg sync.WaitGroup - var mutex sync.Mutex + results := make([]*v1.PermissionCheckResponse, len(checkItems)) + groups := map[groupKey]*groupEntry{} - // Process each check request - for i, checkRequestItem := range checkItems { - wg.Add(1) + for i, item := range checkItems { + if v := item.Validate(); v != nil { + results[i] = &v1.PermissionCheckResponse{ + Can: v1.CheckResult_CHECK_RESULT_DENIED, + Metadata: &v1.PermissionCheckResponseMetadata{}, + } + continue + } - go func(index int, checkRequestItem *v1.PermissionBulkCheckRequestItem) { - defer wg.Done() + subj := item.GetSubject() + key := groupKey{ + entityType: item.GetEntity().GetType(), + permission: item.GetPermission(), + subjectType: subj.GetType(), + subjectID: subj.GetId(), + subjectRelation: subj.GetRelation(), + } + g, ok := groups[key] + if !ok { + g = &groupEntry{subject: subj} + groups[key] = g + } + g.entityIDs = append(g.entityIDs, item.GetEntity().GetId()) + g.indices = append(g.indices, i) + } - // Validate individual request - v := checkRequestItem.Validate() - if v != nil { - resultChannel <- resultItem{ - index: index, - response: &v1.PermissionCheckResponse{ - Can: v1.CheckResult_CHECK_RESULT_DENIED, - Metadata: &v1.PermissionCheckResponseMetadata{ - CheckCount: 0, - }, - }, - } - return - } + // Execute each group as a batch check, concurrently. + var wg sync.WaitGroup + for key, g := range groups { + wg.Add(1) + go func(key groupKey, g *groupEntry) { + defer wg.Done() - // Perform the check using existing Check function - checkRequest := &v1.PermissionCheckRequest{ - TenantId: request.GetTenantId(), - Subject: checkRequestItem.GetSubject(), - Entity: checkRequestItem.GetEntity(), - Permission: checkRequestItem.GetPermission(), + batchReq := &invoke.BatchCheckRequest{ + TenantID: request.GetTenantId(), + EntityType: key.entityType, + EntityIDs: g.entityIDs, + Permission: key.permission, + Subject: g.subject, Metadata: request.GetMetadata(), Context: request.GetContext(), Arguments: request.GetArguments(), } - response, err := r.invoker.Check(ctx, checkRequest) + + resp, err := r.invoker.Check(ctx, batchReq) if err != nil { - // Log error but don't fail the entire bulk operation - slog.ErrorContext(ctx, "check failed in bulk operation", "error", err.Error(), "index", index) - resultChannel <- resultItem{ - index: index, - response: &v1.PermissionCheckResponse{ - Can: v1.CheckResult_CHECK_RESULT_DENIED, - Metadata: &v1.PermissionCheckResponseMetadata{ - CheckCount: 1, - }, - }, + slog.ErrorContext(ctx, "batch check failed in bulk operation", "error", err.Error()) + for _, idx := range g.indices { + results[idx] = &v1.PermissionCheckResponse{ + Can: v1.CheckResult_CHECK_RESULT_DENIED, + Metadata: &v1.PermissionCheckResponseMetadata{CheckCount: 1}, + } } return } - resultChannel <- resultItem{index: index, response: &v1.PermissionCheckResponse{ - Can: response.GetCan(), - Metadata: response.GetMetadata(), - }} - }(i, checkRequestItem) + // Map batch results back to original indices. + for j, entityID := range g.entityIDs { + result := v1.CheckResult_CHECK_RESULT_DENIED + if r, ok := resp.Results[entityID]; ok { + result = r + } + results[g.indices[j]] = &v1.PermissionCheckResponse{ + Can: result, + Metadata: resp.Metadata, + } + } + }(key, g) } + wg.Wait() - // Once the function returns, we wait for all goroutines to finish, then close the resultChannel. - defer func() { - wg.Wait() - close(resultChannel) - }() - - // We read the responses from the resultChannel. - // We expect as many responses as there are references in the entity. - results := make([]*v1.PermissionCheckResponse, len(request.GetItems())) - for range checkItems { - select { - // If we receive a response from the resultChannel, we check for errors. - case response := <-resultChannel: - // If there's no error, we add the result to our response's Results map. - // We use a mutex to safely update the map since multiple goroutines may be writing to it concurrently. - mutex.Lock() - results[response.index] = response.response - mutex.Unlock() - - // If the context is done (i.e., canceled or deadline exceeded), we return an empty response and an error. - case <-ctx.Done(): - return emptyResp, errors.New(v1.ErrorCode_ERROR_CODE_CANCELLED.String()) - } + // Check for context cancellation. + if ctx.Err() != nil { + return emptyResp, errors.New(v1.ErrorCode_ERROR_CODE_CANCELLED.String()) } return &v1.PermissionBulkCheckResponse{ diff --git a/internal/storage/postgres/gc/gc_test.go b/internal/storage/postgres/gc/gc_test.go index 85bb4414e..26d56066d 100644 --- a/internal/storage/postgres/gc/gc_test.go +++ b/internal/storage/postgres/gc/gc_test.go @@ -99,7 +99,7 @@ var _ = Describe("GarbageCollector", func() { checkEngine.SetInvoker(invoker) - checkRes1, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkRes1, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -115,9 +115,9 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "b56661f8-7be6-4342-a4c0-918ee04e5983", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkRes1.Can).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(checkRes1.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) // Step 5: Delete the data _, err = dataWriter.Delete(ctx, tenantID, &base.TupleFilter{ @@ -129,7 +129,7 @@ var _ = Describe("GarbageCollector", func() { Expect(err).ShouldNot(HaveOccurred()) // Step 6: Perform a permission check (expected to be invalid) - checkRes2, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkRes2, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -145,9 +145,9 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "b56661f8-7be6-4342-a4c0-918ee04e5983", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkRes2.Can).Should(Equal(base.CheckResult_CHECK_RESULT_DENIED)) + Expect(checkRes2.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_DENIED)) // Step 7: Insert the same data again tup2, err := tuple.Tuple("organisation:1#member@user:b56661f8-7be6-4342-a4c0-918ee04e5983") @@ -157,7 +157,7 @@ var _ = Describe("GarbageCollector", func() { Expect(err).ShouldNot(HaveOccurred()) // Step 8: Perform a permission check (expected to be valid) - checkRes3, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkRes3, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -173,9 +173,9 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "b56661f8-7be6-4342-a4c0-918ee04e5983", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkRes3.Can).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(checkRes3.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) // Step 9: Run the garbage collector time.Sleep(5 * time.Second) // Pause for 5 seconds @@ -183,7 +183,7 @@ var _ = Describe("GarbageCollector", func() { Expect(err).ShouldNot(HaveOccurred()) // Step 10: Perform a permission check after GC (expected to be valid) - checkRes4, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkRes4, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -199,9 +199,9 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "b56661f8-7be6-4342-a4c0-918ee04e5983", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkRes4.Can).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(checkRes4.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) }) It("should perform tenant-aware garbage collection correctly", func() { @@ -252,7 +252,7 @@ var _ = Describe("GarbageCollector", func() { checkEngine.SetInvoker(invoker) // Check tenant A - checkResA, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkResA, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -268,12 +268,12 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "user-a", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkResA.Can).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(checkResA.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) // Check tenant B - checkResB, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkResB, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -289,9 +289,9 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "user-b", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkResB.Can).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(checkResB.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) // Step 5: Delete data for tenant A only _, err = dataWriter.Delete(ctx, tenantA, &base.TupleFilter{ @@ -308,7 +308,7 @@ var _ = Describe("GarbageCollector", func() { Expect(err).ShouldNot(HaveOccurred()) // Step 7: Verify tenant A's permission is denied (data was deleted) - checkResA2, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkResA2, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -324,12 +324,12 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "user-a", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkResA2.Can).Should(Equal(base.CheckResult_CHECK_RESULT_DENIED)) + Expect(checkResA2.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_DENIED)) // Step 8: Verify tenant B's permission is still allowed (data was not affected by GC) - checkResB2, err := invoker.Check(ctx, &base.PermissionCheckRequest{ + checkResB2, err := invoker.Check(ctx, invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ Metadata: &base.PermissionCheckRequestMetadata{ SnapToken: "", SchemaVersion: "", @@ -345,9 +345,9 @@ var _ = Describe("GarbageCollector", func() { Type: "user", Id: "user-b", }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(checkResB2.Can).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(checkResB2.UnionResult()).Should(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) }) }) diff --git a/pkg/balancer/balancer.go b/pkg/balancer/balancer.go index 0a06cc7e5..6491edb50 100644 --- a/pkg/balancer/balancer.go +++ b/pkg/balancer/balancer.go @@ -4,9 +4,9 @@ import ( "errors" "fmt" "log/slog" + "sync/atomic" "google.golang.org/grpc/balancer" - "google.golang.org/grpc/balancer/base" "google.golang.org/grpc/connectivity" "google.golang.org/grpc/resolver" @@ -20,8 +20,9 @@ type Balancer struct { // The ClientConn to communicate with the gRPC client. clientConn ClientConnWrapper - // Current picker used to select SubConns for requests. - picker balancer.Picker + // picker stores the consistent-hash NodePicker for application code (lock-free reads). + // Nil when nodes are unavailable. + picker atomic.Pointer[picker] // Evaluates connectivity state transitions for SubConns. connectivityEvaluator *balancer.ConnectivityStateEvaluator @@ -52,18 +53,14 @@ func (b *Balancer) ResolverError(err error) { b.lastResolverError = err if b.addressSubConns.Len() == 0 { b.state = connectivity.TransientFailure - b.picker = base.NewErrPicker(errors.Join(b.lastConnectionError, b.lastResolverError)) + b.picker.Store(nil) } if b.state != connectivity.TransientFailure { return } - // Update the balancer state and picker. - b.clientConn.UpdateState(balancer.State{ - ConnectivityState: b.state, - Picker: b.picker, - }) + b.updateGRPCState() } func (b *Balancer) UpdateClientConnState(s balancer.ClientConnState) error { @@ -99,8 +96,8 @@ func (b *Balancer) UpdateClientConnState(s balancer.ClientConnState) error { // Check if the consistent hashing configuration exists. if b.consistent == nil { slog.Error("No consistent hashing configuration found") - b.picker = base.NewErrPicker(errors.Join(b.lastConnectionError, b.lastResolverError)) - b.clientConn.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker}) + b.picker.Store(nil) + b.updateGRPCState() return fmt.Errorf("no consistent hashing configuration found") } @@ -167,29 +164,20 @@ func (b *Balancer) UpdateClientConnState(s balancer.ClientConnState) error { return balancer.ErrBadResolverState } - // Update the picker based on the current balancer state. + // Update the application picker. if b.state == connectivity.TransientFailure { - slog.Warn("Transient failure detected, using error picker") - b.picker = base.NewErrPicker(errors.Join(b.lastConnectionError, b.lastResolverError)) + slog.Warn("Transient failure detected") + b.picker.Store(nil) } else { width := b.config.PickerWidth if width < 1 { width = 1 } - slog.Info("Creating new picker", - slog.Int("width", width), - ) - b.picker = &picker{ - consistent: b.consistent, - width: width, - } + slog.Info("Creating new picker", slog.Int("width", width)) + b.picker.Store(&picker{consistent: b.consistent, width: width}) } - // Update the ClientConn state with the new picker. - slog.Info("Updating ClientConn state", - slog.String("connectivity_state", b.state.String()), - ) - b.clientConn.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker}) + b.updateGRPCState() return nil } @@ -241,10 +229,19 @@ func (b *Balancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubCon } b.state = b.connectivityEvaluator.RecordTransition(oldS, s) + b.updateGRPCState() +} + +// updateGRPCState pushes the current connectivity state to gRPC. +// Always uses subConnPicker — application code pre-computes SubConn in context. +func (b *Balancer) updateGRPCState() { slog.Info("Updating ClientConn state", slog.String("connectivity_state", b.state.String()), ) - b.clientConn.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker}) + b.clientConn.UpdateState(balancer.State{ + ConnectivityState: b.state, + Picker: &subConnPicker{}, + }) } func (b *Balancer) Close() {} diff --git a/pkg/balancer/balancer_test.go b/pkg/balancer/balancer_test.go index 39c0adc40..471b30ad4 100644 --- a/pkg/balancer/balancer_test.go +++ b/pkg/balancer/balancer_test.go @@ -4,7 +4,6 @@ import ( "errors" "testing" - "google.golang.org/grpc/balancer/base" estats "google.golang.org/grpc/experimental/stats" "github.com/cespare/xxhash/v2" @@ -259,7 +258,7 @@ var _ = Describe("Balancer", func() { It("should handle resolver errors with no SubConns", func() { b.ResolverError(errors.New("resolver failure")) Expect(b.state).To(Equal(connectivity.TransientFailure)) - Expect(b.picker).To(Equal(base.NewErrPicker(errors.Join(b.lastConnectionError, b.lastResolverError)))) + Expect(b.picker.Load()).To(BeNil()) }) }) diff --git a/pkg/balancer/builder.go b/pkg/balancer/builder.go index 2e6c38356..2b7327d41 100644 --- a/pkg/balancer/builder.go +++ b/pkg/balancer/builder.go @@ -4,10 +4,10 @@ import ( "encoding/json" "fmt" "sync" + "sync/atomic" "golang.org/x/exp/slog" "google.golang.org/grpc/balancer" - "google.golang.org/grpc/balancer/base" "google.golang.org/grpc/connectivity" "google.golang.org/grpc/resolver" "google.golang.org/grpc/serviceconfig" @@ -18,10 +18,10 @@ import ( // contextKey is a custom type for context keys to avoid collisions type contextKey string -// Package-level constants for the balancer name and consistent hash key. +// Package-level constants for the balancer name and context keys. const ( - Name = "consistenthashing" // Name of the balancer. - Key = contextKey("consistenthashkey") // Key for the consistent hash. + Name = "consistenthashing" // Name of the balancer. + SubConnKey = contextKey("subconn") // Context key for pre-computed SubConn. ) // Config represents the configuration for the consistent hashing balancer. @@ -70,9 +70,19 @@ func (c *Config) ServiceConfigJSON() (string, error) { return string(jsonData), nil } +var defaultBuilder atomic.Pointer[builder] + // NewBuilder initializes a new builder with the given hashing function. +// The builder is also stored as the package default, accessible via GetBuilder. func NewBuilder(fn consistent.Hasher) Builder { - return &builder{hasher: fn} + b := &builder{hasher: fn} + defaultBuilder.Store(b) + return b +} + +// GetBuilder returns the builder created by NewBuilder, or nil. +func GetBuilder() Builder { + return defaultBuilder.Load() } // ConsistentMember represents a member in the consistent hashing ring. @@ -86,15 +96,20 @@ func (s ConsistentMember) String() string { return s.name } // builder is responsible for creating and configuring the consistent hashing balancer. type builder struct { - sync.Mutex // Mutex for thread-safe updates to the builder. - hasher consistent.Hasher // Hashing function for the consistent hash ring. - config Config // Current balancer configuration. + sync.Mutex // Mutex for thread-safe updates to config. + hasher consistent.Hasher // Hashing function for the consistent hash ring. + config Config // Current balancer configuration. + bal atomic.Pointer[Balancer] // Reference to the active gRPC balancer (lock-free). } // Builder defines the interface for the consistent hashing balancer builder. type Builder interface { balancer.Builder // Interface for building balancers. balancer.ConfigParser // Interface for parsing balancer configurations. + + // Picker returns the current NodePicker for routing keys to SubConns. + // Returns nil if the balancer is not ready. + Picker() NodePicker } // Name returns the name of the balancer. @@ -103,19 +118,29 @@ func (b *builder) Name() string { return Name } // Build creates a new instance of the consistent hashing balancer. func (b *builder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer { // Initialize a new balancer with default values. + // picker starts as nil (no SubConns yet); gRPC gets an errPicker via UpdateState below. bal := &Balancer{ clientConn: cc, addressSubConns: resolver.NewAddressMap(), subConnStates: make(map[balancer.SubConn]connectivity.State), connectivityEvaluator: &balancer.ConnectivityStateEvaluator{}, - state: connectivity.Connecting, // Initial state. + state: connectivity.Connecting, hasher: b.hasher, - picker: base.NewErrPicker(balancer.ErrNoSubConnAvailable), // Default picker with no SubConns available. } + b.bal.Store(bal) + return bal } +// Picker returns the current NodePicker. Nil if the balancer is not ready. +func (b *builder) Picker() NodePicker { + if bal := b.bal.Load(); bal != nil { + return bal.picker.Load() + } + return nil +} + // ParseConfig parses the balancer configuration from the provided JSON. func (b *builder) ParseConfig(rm json.RawMessage) (serviceconfig.LoadBalancingConfig, error) { var cfg Config diff --git a/pkg/balancer/picker.go b/pkg/balancer/picker.go index fcf0a17d1..e6e08e1eb 100644 --- a/pkg/balancer/picker.go +++ b/pkg/balancer/picker.go @@ -11,6 +11,24 @@ import ( "github.com/Permify/permify/pkg/consistent" ) +// subConnPicker is a trivial gRPC picker: reads a pre-computed SubConn from context. +// Installed once in gRPC's balancer state; never needs to be replaced. +type subConnPicker struct{} + +func (p *subConnPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { + sc, ok := info.Ctx.Value(SubConnKey).(balancer.SubConn) + if !ok || sc == nil { + return balancer.PickResult{}, fmt.Errorf("no SubConn in context") + } + return balancer.PickResult{SubConn: sc}, nil +} + +// NodePicker resolves a routing key to a target SubConn. +type NodePicker interface { + Pick(key []byte) (balancer.SubConn, error) +} + +// picker implements NodePicker using consistent hashing. type picker struct { consistent *consistent.Consistent width int @@ -35,24 +53,14 @@ var randomIndex = func(max int) int { return int(n.Int64()) } -func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { - // Safely extract the key from the context - keyValue := info.Ctx.Value(Key) - if keyValue == nil { - return balancer.PickResult{}, fmt.Errorf("context key missing") - } - key, ok := keyValue.([]byte) - if !ok { - return balancer.PickResult{}, fmt.Errorf("context key is not of type []byte") - } - - // Retrieve the closest N members +// Pick computes the target SubConn for a routing key using consistent hashing. +func (p *picker) Pick(key []byte) (balancer.SubConn, error) { members, err := p.consistent.ClosestN(key, p.width) if err != nil { - return balancer.PickResult{}, fmt.Errorf("failed to get closest members: %w", err) + return nil, fmt.Errorf("failed to get closest members: %w", err) } if len(members) == 0 { - return balancer.PickResult{}, fmt.Errorf("no available members") + return nil, fmt.Errorf("no available members") } // Randomly pick one member if width > 1 @@ -64,9 +72,9 @@ func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { // Assert the member type chosen, ok := members[index].(ConsistentMember) if !ok { - return balancer.PickResult{}, fmt.Errorf("invalid member type: expected subConnMember") + return nil, fmt.Errorf("invalid member type: expected ConsistentMember") } // Return the chosen connection - return balancer.PickResult{SubConn: chosen.SubConn}, nil + return chosen.SubConn, nil } diff --git a/pkg/balancer/picker_test.go b/pkg/balancer/picker_test.go index 51ea23fb7..fcfe19440 100644 --- a/pkg/balancer/picker_test.go +++ b/pkg/balancer/picker_test.go @@ -46,108 +46,81 @@ var _ = Describe("Picker and Consistent Hashing", func() { } }) - Describe("Picker Logic", func() { - var ( - p *picker - testCtx context.Context - ) + Describe("subConnPicker (pass-through)", func() { + var p *subConnPicker BeforeEach(func() { - // Initialize picker with consistent hashing and a width of 2 - p = &picker{ - consistent: c, - width: 2, - } - // Set up context with a valid key - testCtx = context.WithValue(context.Background(), Key, []byte("test-key")) + p = &subConnPicker{} }) - It("should pick a member successfully", func() { - members, err := c.ClosestN([]byte("test-key"), 2) + It("should return SubConn from context", func() { + sc := &mockSubConnWrapper{} + ctx := context.WithValue(context.Background(), SubConnKey, balancer.SubConn(sc)) + result, err := p.Pick(balancer.PickInfo{Ctx: ctx}) Expect(err).To(BeNil()) - Expect(len(members)).To(BeNumerically(">", 0)) - Expect(members[0].(ConsistentMember).String()).To(Equal("member1")) + Expect(result.SubConn).To(Equal(sc)) }) - It("should return an error if the context key is missing", func() { + It("should return error if no SubConn in context", func() { result, err := p.Pick(balancer.PickInfo{Ctx: context.Background()}) - Expect(err).To(MatchError("context key missing")) + Expect(err).To(MatchError("no SubConn in context")) Expect(result.SubConn).To(BeNil()) }) - It("should return an error if no members are available", func() { - // Remove all members - for _, m := range members { - c.Remove(m.String()) - } - result, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) - Expect(err).To(MatchError("failed to get closest members: not enough members to satisfy the request")) + It("should return error if SubConn is nil in context", func() { + ctx := context.WithValue(context.Background(), SubConnKey, nil) + result, err := p.Pick(balancer.PickInfo{Ctx: ctx}) + Expect(err).To(MatchError("no SubConn in context")) Expect(result.SubConn).To(BeNil()) }) + }) - It("should handle context key with wrong type", func() { - wrongCtx := context.WithValue(context.Background(), Key, "wrong-type") - result, err := p.Pick(balancer.PickInfo{Ctx: wrongCtx}) - Expect(err).To(MatchError("context key is not of type []byte")) - Expect(result.SubConn).To(BeNil()) - }) + Describe("Pick", func() { + var p *picker - It("should handle empty key", func() { - emptyCtx := context.WithValue(context.Background(), Key, []byte{}) - result, err := p.Pick(balancer.PickInfo{Ctx: emptyCtx}) - Expect(err).To(BeNil()) - Expect(result.SubConn).ToNot(BeNil()) + BeforeEach(func() { + p = &picker{consistent: c, width: 2} }) - It("should handle picker with width of 1", func() { - p.width = 1 - result, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) + It("should locate a member successfully", func() { + sc, err := p.Pick([]byte("test-key")) Expect(err).To(BeNil()) - Expect(result.SubConn).ToNot(BeNil()) + Expect(sc).ToNot(BeNil()) }) - It("should handle picker with width larger than available members", func() { - p.width = 10 - _, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) - // If not enough members, should error - if len(members) < 10 { - Expect(err).ToNot(BeNil()) - } else { - Expect(err).To(BeNil()) + It("should return an error if no members are available", func() { + for _, m := range members { + c.Remove(m.String()) } + _, err := p.Pick([]byte("test-key")) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("failed to get closest members")) }) - It("should handle picker with zero width", func() { - p.width = 0 - result, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) + It("should handle empty key", func() { + sc, err := p.Pick([]byte{}) Expect(err).To(BeNil()) - Expect(result.SubConn).ToNot(BeNil()) + Expect(sc).ToNot(BeNil()) }) - It("should handle picker with negative width", func() { - p.width = -1 - result, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) + It("should handle width of 1", func() { + p.width = 1 + sc, err := p.Pick([]byte("test-key")) Expect(err).To(BeNil()) - Expect(result.SubConn).ToNot(BeNil()) + Expect(sc).ToNot(BeNil()) }) - It("should consistently pick the same member for the same key", func() { - result1, err1 := p.Pick(balancer.PickInfo{Ctx: testCtx}) - result2, err2 := p.Pick(balancer.PickInfo{Ctx: testCtx}) - Expect(err1).To(BeNil()) - Expect(err2).To(BeNil()) - Expect(result1.SubConn).To(Equal(result2.SubConn)) + It("should handle width larger than available members", func() { + p.width = 10 + _, err := p.Pick([]byte("test-key")) + Expect(err).ToNot(BeNil()) }) - It("should pick different members for different keys", func() { - ctx1 := context.WithValue(context.Background(), Key, []byte("key1")) - ctx2 := context.WithValue(context.Background(), Key, []byte("key2")) - result1, err1 := p.Pick(balancer.PickInfo{Ctx: ctx1}) - result2, err2 := p.Pick(balancer.PickInfo{Ctx: ctx2}) - Expect(err1).To(BeNil()) - Expect(err2).To(BeNil()) - Expect(result1.SubConn).ToNot(BeNil()) - Expect(result2.SubConn).ToNot(BeNil()) + It("should handle width of zero", func() { + p.width = 0 + sc, err := p.Pick([]byte("test-key")) + Expect(err).To(BeNil()) + Expect(sc).ToNot(BeNil()) }) It("should handle very long keys", func() { @@ -155,10 +128,9 @@ var _ = Describe("Picker and Consistent Hashing", func() { for i := range longKey { longKey[i] = byte(i % 256) } - longCtx := context.WithValue(context.Background(), Key, longKey) - result, err := p.Pick(balancer.PickInfo{Ctx: longCtx}) + sc, err := p.Pick(longKey) Expect(err).To(BeNil()) - Expect(result.SubConn).ToNot(BeNil()) + Expect(sc).ToNot(BeNil()) }) It("should handle special characters in keys", func() { @@ -192,10 +164,9 @@ var _ = Describe("Picker and Consistent Hashing", func() { []byte("key/with/forward/slash"), } for _, key := range specialKeys { - specialCtx := context.WithValue(context.Background(), Key, key) - result, err := p.Pick(balancer.PickInfo{Ctx: specialCtx}) + sc, err := p.Pick(key) Expect(err).To(BeNil(), "Should handle key: %s", string(key)) - Expect(result.SubConn).ToNot(BeNil(), "Should return SubConn for key: %s", string(key)) + Expect(sc).ToNot(BeNil(), "Should return SubConn for key: %s", string(key)) } }) @@ -213,20 +184,16 @@ var _ = Describe("Picker and Consistent Hashing", func() { []byte("key with தமிழ்"), } for _, key := range unicodeKeys { - unicodeCtx := context.WithValue(context.Background(), Key, key) - result, err := p.Pick(balancer.PickInfo{Ctx: unicodeCtx}) + sc, err := p.Pick(key) Expect(err).To(BeNil(), "Should handle unicode key: %s", string(key)) - Expect(result.SubConn).ToNot(BeNil(), "Should return SubConn for unicode key: %s", string(key)) + Expect(sc).ToNot(BeNil(), "Should return SubConn for unicode key: %s", string(key)) } }) }) Describe("Consistent Hashing Behavior", func() { It("should distribute keys evenly across members", func() { - p := &picker{ - consistent: c, - width: 1, - } + p := &picker{consistent: c, width: 1} // Map mockSubConnWrapper pointer to name subConnToName := map[balancer.SubConn]string{} @@ -239,10 +206,9 @@ var _ = Describe("Picker and Consistent Hashing", func() { for i := 0; i < keyCount; i++ { key := []byte(fmt.Sprintf("key-%d", i)) - ctx := context.WithValue(context.Background(), Key, key) - result, err := p.Pick(balancer.PickInfo{Ctx: ctx}) + sc, err := p.Pick(key) Expect(err).To(BeNil()) - pickedName := subConnToName[result.SubConn] + pickedName := subConnToName[sc] memberCounts[pickedName]++ } @@ -256,38 +222,30 @@ var _ = Describe("Picker and Consistent Hashing", func() { }) It("should handle member removal gracefully", func() { - p := &picker{ - consistent: c, - width: 2, - } - testCtx := context.WithValue(context.Background(), Key, []byte("test-key")) - result1, err1 := p.Pick(balancer.PickInfo{Ctx: testCtx}) + p := &picker{consistent: c, width: 2} + sc1, err1 := p.Pick([]byte("test-key")) Expect(err1).To(BeNil()) - Expect(result1.SubConn).ToNot(BeNil()) + Expect(sc1).ToNot(BeNil()) c.Remove("member1") - result2, err2 := p.Pick(balancer.PickInfo{Ctx: testCtx}) + sc2, err2 := p.Pick([]byte("test-key")) if len(members)-1 < 2 { Expect(err2).ToNot(BeNil()) } else { Expect(err2).To(BeNil()) - Expect(result2.SubConn).ToNot(BeNil()) + Expect(sc2).ToNot(BeNil()) } }) It("should handle member addition gracefully", func() { - p := &picker{ - consistent: c, - width: 2, - } - testCtx := context.WithValue(context.Background(), Key, []byte("test-key")) - result1, err1 := p.Pick(balancer.PickInfo{Ctx: testCtx}) + p := &picker{consistent: c, width: 2} + sc1, err1 := p.Pick([]byte("test-key")) Expect(err1).To(BeNil()) - Expect(result1.SubConn).ToNot(BeNil()) + Expect(sc1).ToNot(BeNil()) newMember := ConsistentMember{SubConn: &mockSubConnWrapper{}, name: "member4"} c.Add(newMember) - result2, err2 := p.Pick(balancer.PickInfo{Ctx: testCtx}) + sc2, err2 := p.Pick([]byte("test-key")) Expect(err2).To(BeNil()) - Expect(result2.SubConn).ToNot(BeNil()) + Expect(sc2).ToNot(BeNil()) }) }) @@ -301,46 +259,25 @@ var _ = Describe("Picker and Consistent Hashing", func() { Load: 1.0, }) - p := &picker{ - consistent: brokenC, - width: 2, - } - - // Create a test context - testCtx := context.WithValue(context.Background(), Key, []byte("test-key")) - - // Try to pick - should handle the error gracefully - result, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) + p := &picker{consistent: brokenC, width: 2} + _, err := p.Pick([]byte("test-key")) Expect(err).To(HaveOccurred()) - Expect(result.SubConn).To(BeNil()) }) It("should handle nil consistent hashing", func() { - p := &picker{ - consistent: nil, - width: 2, - } - - // Create a test context - testCtx := context.WithValue(context.Background(), Key, []byte("test-key")) - - // This should panic or return an error + p := &picker{consistent: nil, width: 2} Expect(func() { - p.Pick(balancer.PickInfo{Ctx: testCtx}) + p.Pick([]byte("test-key")) }).To(Panic()) }) }) - Describe("Picker Configuration", func() { + Describe("picker Configuration", func() { It("should work with different width configurations", func() { widths := []int{1, 2, 3, 5, 10} for _, width := range widths { - p := &picker{ - consistent: c, - width: width, - } - testCtx := context.WithValue(context.Background(), Key, []byte("test-key")) - _, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) + p := &picker{consistent: c, width: width} + _, err := p.Pick([]byte("test-key")) if width > len(members) { Expect(err).ToNot(BeNil(), "Should error with width %d", width) } else { @@ -352,12 +289,8 @@ var _ = Describe("Picker and Consistent Hashing", func() { It("should handle edge case width values", func() { edgeWidths := []int{0, -1, -100, 1000, 999999} for _, width := range edgeWidths { - p := &picker{ - consistent: c, - width: width, - } - testCtx := context.WithValue(context.Background(), Key, []byte("test-key")) - _, err := p.Pick(balancer.PickInfo{Ctx: testCtx}) + p := &picker{consistent: c, width: width} + _, err := p.Pick([]byte("test-key")) if width > len(members) { Expect(err).ToNot(BeNil(), "Should error with edge width %d", width) } else { diff --git a/pkg/cmd/serve.go b/pkg/cmd/serve.go index 77f40955c..054aff6f3 100644 --- a/pkg/cmd/serve.go +++ b/pkg/cmd/serve.go @@ -24,6 +24,7 @@ import ( cacheproxy "github.com/Permify/permify/internal/storage/proxies/cache" cbproxy "github.com/Permify/permify/internal/storage/proxies/circuitbreaker" sfproxy "github.com/Permify/permify/internal/storage/proxies/singleflight" + consistentbalancer "github.com/Permify/permify/pkg/balancer" "github.com/Permify/permify/pkg/cmd/flags" PQDatabase "github.com/Permify/permify/pkg/database/postgres" @@ -442,7 +443,10 @@ func serve() func(cmd *cobra.Command, args []string) error { } // Initialize the engines using the key manager, schema reader, and relationship reader - checkEngine := engines.NewCheckEngine(schemaReader, dataReader, engines.CheckConcurrencyLimit(cfg.Service.Permission.ConcurrencyLimit)) + checkEngine := engines.NewCheckEngine(schemaReader, dataReader, + engines.CheckConcurrencyLimit(cfg.Service.Permission.ConcurrencyLimit), + engines.CheckMaxBatchSize(cfg.Service.Permission.BulkLimit), + ) expandEngine := engines.NewExpandEngine(schemaReader, dataReader) // Declare a variable `checker` of type `invoke.Check`. @@ -464,6 +468,7 @@ func serve() func(cmd *cobra.Command, args []string) error { ctx, checker, schemaReader, + consistentbalancer.GetBuilder(), cfg.Server.NameOverride, &cfg.Distributed, &cfg.Server.GRPC, diff --git a/pkg/cmd/validate.go b/pkg/cmd/validate.go index 3c246f99b..7cce2eca9 100644 --- a/pkg/cmd/validate.go +++ b/pkg/cmd/validate.go @@ -15,6 +15,7 @@ import ( "github.com/rs/xid" "github.com/spf13/cobra" + "github.com/Permify/permify/internal/invoke" "github.com/Permify/permify/internal/storage" serverValidation "github.com/Permify/permify/internal/validation" "github.com/Permify/permify/pkg/attribute" @@ -287,17 +288,18 @@ func validate() func(cmd *cobra.Command, args []string) error { } // Perform a permission check based on the context, entity, permission, and subject - res, err := dev.Container.Invoker.Check(ctx, &base.PermissionCheckRequest{ - TenantId: "t1", - Context: cont, + res, err := dev.Container.Invoker.Check(ctx, &invoke.BatchCheckRequest{ + TenantID: "t1", + EntityType: entity.GetType(), + EntityIDs: []string{entity.GetId()}, + Permission: permission, + Subject: subject, Metadata: &base.PermissionCheckRequestMetadata{ SchemaVersion: version, SnapToken: token.NewNoopToken().Encode().String(), Depth: depth, }, - Entity: entity, - Permission: permission, - Subject: subject, + Context: cont, }) if err != nil { list.Add(fmt.Sprintf("%s -> %s", query, err.Error())) @@ -306,13 +308,13 @@ func validate() func(cmd *cobra.Command, args []string) error { } // If the check result matches the expected result, log a success message - if res.Can == exp { + if res.UnionResult() == exp { color.Success.Print(" success:") fmt.Printf(" %s \n", query) } else { // If the check result does not match the expected result, log a failure message color.Danger.Printf(" fail: %s ->", query) - if res.Can == base.CheckResult_CHECK_RESULT_ALLOWED { + if res.UnionResult() == base.CheckResult_CHECK_RESULT_ALLOWED { color.Danger.Println(" expected: DENIED actual: ALLOWED ") list.Add(fmt.Sprintf("%s -> expected: DENIED actual: ALLOWED ", query)) } else { diff --git a/pkg/development/development.go b/pkg/development/development.go index 425b91cfb..97af3031c 100644 --- a/pkg/development/development.go +++ b/pkg/development/development.go @@ -328,17 +328,18 @@ func (c *Development) RunWithShape(ctx context.Context, shape *file.Shape) (erro } // A Permission Check is made for the current entity, permission and subject - res, err := c.Container.Invoker.Check(ctx, &v1.PermissionCheckRequest{ - TenantId: "t1", + res, err := c.Container.Invoker.Check(ctx, &invoke.BatchCheckRequest{ + TenantID: "t1", + EntityType: entity.GetType(), + EntityIDs: []string{entity.GetId()}, + Permission: permission, + Subject: subject, Metadata: &v1.PermissionCheckRequestMetadata{ SchemaVersion: version, SnapToken: token.NewNoopToken().Encode().String(), Depth: 100, }, - Context: cont, - Entity: entity, - Permission: permission, - Subject: subject, + Context: cont, }) if err != nil { errors = append(errors, Error{ @@ -352,7 +353,7 @@ func (c *Development) RunWithShape(ctx context.Context, shape *file.Shape) (erro query := tuple.SubjectToString(subject) + " " + permission + " " + tuple.EntityToString(entity) // Check if the permission check result matches the expected result - if res.Can != exp { + if res.UnionResult() != exp { var expectedStr, actualStr string if exp == v1.CheckResult_CHECK_RESULT_ALLOWED { expectedStr = "true" @@ -360,7 +361,7 @@ func (c *Development) RunWithShape(ctx context.Context, shape *file.Shape) (erro expectedStr = "false" } - if res.Can == v1.CheckResult_CHECK_RESULT_ALLOWED { + if res.UnionResult() == v1.CheckResult_CHECK_RESULT_ALLOWED { actualStr = "true" } else { actualStr = "false" From 3d1f2f2c991052c859e77e78a8267534a5f2ba14 Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Mon, 20 Jul 2026 15:15:42 +0200 Subject: [PATCH 04/13] perf: use concurrency limit to limit per-request db queries --- internal/engines/expand.go | 4 +- internal/engines/subject_permission.go | 19 +--- internal/engines/utils.go | 10 -- internal/invoke/concurrency.go | 33 +++++++ internal/servers/permission_server.go | 19 +++- internal/servers/server.go | 7 +- .../storage/proxies/semaphore/data_reader.go | 95 +++++++++++++++++++ pkg/cmd/serve.go | 5 +- 8 files changed, 157 insertions(+), 35 deletions(-) create mode 100644 internal/invoke/concurrency.go create mode 100644 internal/storage/proxies/semaphore/data_reader.go diff --git a/internal/engines/expand.go b/internal/engines/expand.go index a789fafbe..24b778b0c 100644 --- a/internal/engines/expand.go +++ b/internal/engines/expand.go @@ -24,9 +24,7 @@ type ExpandEngine struct { dataReader storage.DataReader } -// NewExpandEngine - This function creates a new instance of ExpandEngine by taking a SchemaReader and a RelationshipReader as -// parameters and returning a pointer to the created instance. The SchemaReader is used to read schema definitions, while the -// RelationshipReader is used to read relationship definitions. +// NewExpandEngine creates a new instance of ExpandEngine. func NewExpandEngine(sr storage.SchemaReader, rr storage.DataReader) *ExpandEngine { return &ExpandEngine{ schemaReader: sr, diff --git a/internal/engines/subject_permission.go b/internal/engines/subject_permission.go index b521f07b0..df361dc8a 100644 --- a/internal/engines/subject_permission.go +++ b/internal/engines/subject_permission.go @@ -16,24 +16,13 @@ type SubjectPermissionEngine struct { checker invoke.Check // schemaReader is responsible for reading schema information schemaReader storage.SchemaReader - // concurrencyLimit is the maximum number of concurrent permission checks allowed - concurrencyLimit int } -func NewSubjectPermission(checker invoke.Check, sr storage.SchemaReader, opts ...SubjectPermissionOption) *SubjectPermissionEngine { - // Initialize a CheckEngine with default concurrency limit and provided parameters - engine := &SubjectPermissionEngine{ - checker: checker, - schemaReader: sr, - concurrencyLimit: _defaultConcurrencyLimit, +func NewSubjectPermission(checker invoke.Check, sr storage.SchemaReader) *SubjectPermissionEngine { + return &SubjectPermissionEngine{ + checker: checker, + schemaReader: sr, } - - // Apply provided options to configure the CheckEngine - for _, opt := range opts { - opt(engine) - } - - return engine } // SubjectPermission is a method on the SubjectPermissionEngine struct. diff --git a/internal/engines/utils.go b/internal/engines/utils.go index 722cc021f..215fe42b3 100644 --- a/internal/engines/utils.go +++ b/internal/engines/utils.go @@ -56,16 +56,6 @@ func SubjectFilterConcurrencyLimit(limit int) SubjectFilterOption { } } -// SubjectPermissionOption - a functional option type for configuring the SubjectPermissionEngine. -type SubjectPermissionOption func(engine *SubjectPermissionEngine) - -// SubjectPermissionConcurrencyLimit - a functional option that sets the concurrency limit for the SubjectPermissionEngine. -func SubjectPermissionConcurrencyLimit(limit int) SubjectPermissionOption { - return func(c *SubjectPermissionEngine) { - c.concurrencyLimit = limit - } -} - // joinResponseMetas - a helper function that merges multiple PermissionCheckResponseMetadata structs into one. func joinResponseMetas(meta ...*base.PermissionCheckResponseMetadata) *base.PermissionCheckResponseMetadata { response := &base.PermissionCheckResponseMetadata{} diff --git a/internal/invoke/concurrency.go b/internal/invoke/concurrency.go new file mode 100644 index 000000000..1e7e7e953 --- /dev/null +++ b/internal/invoke/concurrency.go @@ -0,0 +1,33 @@ +package invoke + +import ( + "context" + "math" + + "golang.org/x/sync/semaphore" +) + +// DefaultConcurrencyLimit is the default max concurrent DB operations per request. +const DefaultConcurrencyLimit = 100 + +type concurrencySemaphoreKeyType struct{} + +var concurrencySemaphoreKey = concurrencySemaphoreKeyType{} + +// noopSemaphore is a semaphore that never blocks — used as fallback +// when no request-scoped semaphore was set (e.g. in tests or dev mode). +var noopSemaphore = semaphore.NewWeighted(math.MaxInt64) + +// WithConcurrencySemaphore stores a request-scoped semaphore in the context. +func WithConcurrencySemaphore(ctx context.Context, sem *semaphore.Weighted) context.Context { + return context.WithValue(ctx, concurrencySemaphoreKey, sem) +} + +// ConcurrencySemaphoreFromContext retrieves the request-scoped semaphore. +// Returns a no-op semaphore if none was set — never returns nil. +func ConcurrencySemaphoreFromContext(ctx context.Context) *semaphore.Weighted { + if sem, ok := ctx.Value(concurrencySemaphoreKey).(*semaphore.Weighted); ok && sem != nil { + return sem + } + return noopSemaphore +} diff --git a/internal/servers/permission_server.go b/internal/servers/permission_server.go index 540ebf3b7..261cf0b37 100644 --- a/internal/servers/permission_server.go +++ b/internal/servers/permission_server.go @@ -7,6 +7,7 @@ import ( "sync" otelCodes "go.opentelemetry.io/otel/codes" + "golang.org/x/sync/semaphore" "google.golang.org/grpc/status" "github.com/Permify/permify/internal" @@ -18,13 +19,18 @@ import ( type PermissionServer struct { v1.UnimplementedPermissionServer - invoker invoke.Invoker + invoker invoke.Invoker + concurrencyLimit int } // NewPermissionServer - Creates new Permission Server -func NewPermissionServer(i invoke.Invoker) *PermissionServer { +func NewPermissionServer(i invoke.Invoker, concurrencyLimit int) *PermissionServer { + if concurrencyLimit <= 0 { + concurrencyLimit = invoke.DefaultConcurrencyLimit + } return &PermissionServer{ - invoker: i, + invoker: i, + concurrencyLimit: concurrencyLimit, } } @@ -32,6 +38,7 @@ func NewPermissionServer(i invoke.Invoker) *PermissionServer { func (r *PermissionServer) Check(ctx context.Context, request *v1.PermissionCheckRequest) (*v1.PermissionCheckResponse, error) { ctx, span := internal.Tracer.Start(ctx, "permissions.check") defer span.End() + ctx = invoke.WithConcurrencySemaphore(ctx, semaphore.NewWeighted(int64(r.concurrencyLimit))) v := request.Validate() if v != nil { @@ -58,6 +65,7 @@ func (r *PermissionServer) BulkCheck(ctx context.Context, request *v1.Permission ctx, span := internal.Tracer.Start(ctx, "permissions.bulk-check") defer span.End() + ctx = invoke.WithConcurrencySemaphore(ctx, semaphore.NewWeighted(int64(r.concurrencyLimit))) // Validate tenant_id if request.GetTenantId() == "" { @@ -183,6 +191,7 @@ func (r *PermissionServer) BulkCheck(ctx context.Context, request *v1.Permission func (r *PermissionServer) Expand(ctx context.Context, request *v1.PermissionExpandRequest) (*v1.PermissionExpandResponse, error) { ctx, span := internal.Tracer.Start(ctx, "permissions.expand") defer span.End() + ctx = invoke.WithConcurrencySemaphore(ctx, semaphore.NewWeighted(int64(r.concurrencyLimit))) v := request.Validate() if v != nil { @@ -204,6 +213,7 @@ func (r *PermissionServer) Expand(ctx context.Context, request *v1.PermissionExp func (r *PermissionServer) LookupEntity(ctx context.Context, request *v1.PermissionLookupEntityRequest) (*v1.PermissionLookupEntityResponse, error) { ctx, span := internal.Tracer.Start(ctx, "permissions.lookup-entity") defer span.End() + ctx = invoke.WithConcurrencySemaphore(ctx, semaphore.NewWeighted(int64(r.concurrencyLimit))) v := request.Validate() if v != nil { @@ -225,6 +235,7 @@ func (r *PermissionServer) LookupEntity(ctx context.Context, request *v1.Permiss func (r *PermissionServer) LookupEntityStream(request *v1.PermissionLookupEntityRequest, server v1.Permission_LookupEntityStreamServer) error { ctx, span := internal.Tracer.Start(server.Context(), "permissions.lookup-entity-stream") defer span.End() + ctx = invoke.WithConcurrencySemaphore(ctx, semaphore.NewWeighted(int64(r.concurrencyLimit))) v := request.Validate() if v != nil { @@ -246,6 +257,7 @@ func (r *PermissionServer) LookupEntityStream(request *v1.PermissionLookupEntity func (r *PermissionServer) LookupSubject(ctx context.Context, request *v1.PermissionLookupSubjectRequest) (*v1.PermissionLookupSubjectResponse, error) { ctx, span := internal.Tracer.Start(ctx, "permissions.lookup-subject") defer span.End() + ctx = invoke.WithConcurrencySemaphore(ctx, semaphore.NewWeighted(int64(r.concurrencyLimit))) v := request.Validate() if v != nil { @@ -267,6 +279,7 @@ func (r *PermissionServer) LookupSubject(ctx context.Context, request *v1.Permis func (r *PermissionServer) SubjectPermission(ctx context.Context, request *v1.PermissionSubjectPermissionRequest) (*v1.PermissionSubjectPermissionResponse, error) { ctx, span := internal.Tracer.Start(ctx, "permissions.subject-permission") defer span.End() + ctx = invoke.WithConcurrencySemaphore(ctx, semaphore.NewWeighted(int64(r.concurrencyLimit))) v := request.Validate() if v != nil { diff --git a/internal/servers/server.go b/internal/servers/server.go index 6c7fa525f..b1422dadc 100644 --- a/internal/servers/server.go +++ b/internal/servers/server.go @@ -63,6 +63,9 @@ type Container struct { TW storage.TenantWriter W storage.Watcher + + // ConcurrencyLimit for permission checks + ConcurrencyLimit int } // NewContainer is a constructor for the Container struct. @@ -172,7 +175,7 @@ func (s *Container) Run( grpcServer := grpc.NewServer(opts...) // Register various gRPC services to the server. - grpcV1.RegisterPermissionServer(grpcServer, NewPermissionServer(s.Invoker)) + grpcV1.RegisterPermissionServer(grpcServer, NewPermissionServer(s.Invoker, s.ConcurrencyLimit)) grpcV1.RegisterSchemaServer(grpcServer, NewSchemaServer(s.SW, s.SR)) grpcV1.RegisterDataServer(grpcServer, NewDataServer(s.DR, s.DW, s.BR, s.SR)) grpcV1.RegisterBundleServer(grpcServer, NewBundleServer(s.BR, s.BW)) @@ -185,7 +188,7 @@ func (s *Container) Run( // Create another gRPC server, presumably for invoking permissions. invokeServer := grpc.NewServer(opts...) - grpcV1.RegisterPermissionServer(invokeServer, NewPermissionServer(localInvoker)) + grpcV1.RegisterPermissionServer(invokeServer, NewPermissionServer(localInvoker, s.ConcurrencyLimit)) // Register health check and reflection services for the invokeServer. health.RegisterHealthServer(invokeServer, NewHealthServer()) // Register health server for invoker diff --git a/internal/storage/proxies/semaphore/data_reader.go b/internal/storage/proxies/semaphore/data_reader.go new file mode 100644 index 000000000..68f5150ac --- /dev/null +++ b/internal/storage/proxies/semaphore/data_reader.go @@ -0,0 +1,95 @@ +package semaphore + +import ( + "context" + + "github.com/Permify/permify/internal/invoke" + "github.com/Permify/permify/internal/storage" + "github.com/Permify/permify/pkg/database" + base "github.com/Permify/permify/pkg/pb/base/v1" + "github.com/Permify/permify/pkg/token" +) + +// DataReader wraps a storage.DataReader with per-request semaphore protection. +// Each DB query acquires a slot from the request-scoped semaphore (stored in context), +// preventing connection pool exhaustion from concurrent goroutines. +type DataReader struct { + delegate storage.DataReader +} + +// NewDataReader creates a semaphore-protected DataReader. +func NewDataReader(delegate storage.DataReader) *DataReader { + return &DataReader{delegate: delegate} +} + +func (r *DataReader) QueryRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, snap string, pagination database.CursorPagination) (*database.TupleIterator, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, err + } + defer sem.Release(1) + return r.delegate.QueryRelationships(ctx, tenantID, filter, snap, pagination) +} + +func (r *DataReader) QueryRelationshipsWithSubjectFilter(ctx context.Context, tenantID string, filter *base.TupleFilter, subject *base.Subject, snap string, pagination database.CursorPagination) (*database.TupleIterator, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, err + } + defer sem.Release(1) + return r.delegate.QueryRelationshipsWithSubjectFilter(ctx, tenantID, filter, subject, snap, pagination) +} + +func (r *DataReader) ReadRelationships(ctx context.Context, tenantID string, filter *base.TupleFilter, snap string, pagination database.Pagination) (*database.TupleCollection, database.EncodedContinuousToken, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, nil, err + } + defer sem.Release(1) + return r.delegate.ReadRelationships(ctx, tenantID, filter, snap, pagination) +} + +func (r *DataReader) QuerySingleAttribute(ctx context.Context, tenantID string, filter *base.AttributeFilter, snap string) (*base.Attribute, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, err + } + defer sem.Release(1) + return r.delegate.QuerySingleAttribute(ctx, tenantID, filter, snap) +} + +func (r *DataReader) QueryAttributes(ctx context.Context, tenantID string, filter *base.AttributeFilter, snap string, pagination database.CursorPagination) (*database.AttributeIterator, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, err + } + defer sem.Release(1) + return r.delegate.QueryAttributes(ctx, tenantID, filter, snap, pagination) +} + +func (r *DataReader) ReadAttributes(ctx context.Context, tenantID string, filter *base.AttributeFilter, snap string, pagination database.Pagination) (*database.AttributeCollection, database.EncodedContinuousToken, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, nil, err + } + defer sem.Release(1) + return r.delegate.ReadAttributes(ctx, tenantID, filter, snap, pagination) +} + +func (r *DataReader) QueryUniqueSubjectReferences(ctx context.Context, tenantID string, subjectReference *base.RelationReference, excluded []string, snap string, pagination database.Pagination) ([]string, database.EncodedContinuousToken, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, nil, err + } + defer sem.Release(1) + return r.delegate.QueryUniqueSubjectReferences(ctx, tenantID, subjectReference, excluded, snap, pagination) +} + +func (r *DataReader) HeadSnapshot(ctx context.Context, tenantID string) (token.SnapToken, error) { + sem := invoke.ConcurrencySemaphoreFromContext(ctx) + if err := sem.Acquire(ctx, 1); err != nil { + return nil, err + } + defer sem.Release(1) + return r.delegate.HeadSnapshot(ctx, tenantID) +} diff --git a/pkg/cmd/serve.go b/pkg/cmd/serve.go index 054aff6f3..3887d7fd9 100644 --- a/pkg/cmd/serve.go +++ b/pkg/cmd/serve.go @@ -23,6 +23,7 @@ import ( "github.com/Permify/permify/internal/storage/postgres/gc" cacheproxy "github.com/Permify/permify/internal/storage/proxies/cache" cbproxy "github.com/Permify/permify/internal/storage/proxies/circuitbreaker" + semproxy "github.com/Permify/permify/internal/storage/proxies/semaphore" sfproxy "github.com/Permify/permify/internal/storage/proxies/singleflight" consistentbalancer "github.com/Permify/permify/pkg/balancer" "github.com/Permify/permify/pkg/cmd/flags" @@ -415,6 +416,7 @@ func serve() func(cmd *cobra.Command, args []string) error { schemaReader = cacheproxy.NewSchemaReader(schemaReader, schemaCache) dataReader = sfproxy.NewDataReader(dataReader) + dataReader = semproxy.NewDataReader(dataReader) schemaReader = sfproxy.NewSchemaReader(schemaReader) // Check if circuit breaker should be enabled for services @@ -501,8 +503,6 @@ func serve() func(cmd *cobra.Command, args []string) error { subjectPermissionEngine := engines.NewSubjectPermission( checker, schemaReader, - // Set concurrency limit for the subject permission checks. - engines.SubjectPermissionConcurrencyLimit(cfg.Service.Permission.ConcurrencyLimit), ) // Create a new invoker that is used to directly call various functions or engines. @@ -542,6 +542,7 @@ func serve() func(cmd *cobra.Command, args []string) error { tenantWriter, watcher, ) + container.ConcurrencyLimit = cfg.Service.Permission.ConcurrencyLimit // Create an error group with the provided context var g *errgroup.Group From 72a5de66cdde4315fc502eab625da8b1d376a13b Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Mon, 20 Jul 2026 18:46:35 +0200 Subject: [PATCH 05/13] fix: review fixes --- internal/engines/balancer/balancer.go | 4 +--- pkg/balancer/balancer.go | 2 +- pkg/balancer/picker.go | 4 +++- pkg/cmd/serve.go | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/internal/engines/balancer/balancer.go b/internal/engines/balancer/balancer.go index c498709d7..d6c998479 100644 --- a/internal/engines/balancer/balancer.go +++ b/internal/engines/balancer/balancer.go @@ -220,9 +220,7 @@ func (c *Balancer) Check(ctx context.Context, request *invoke.BatchCheckRequest) } for _, r := range results { if r.err != nil { - slog.ErrorContext(ctx, "node group check failed", "error", r.err.Error()) - // Mark all entities in failed group as denied. - continue + return deniedResp, r.err } if r.resp != nil { for entityID, result := range r.resp.Results { diff --git a/pkg/balancer/balancer.go b/pkg/balancer/balancer.go index 6491edb50..310bda9e4 100644 --- a/pkg/balancer/balancer.go +++ b/pkg/balancer/balancer.go @@ -240,7 +240,7 @@ func (b *Balancer) updateGRPCState() { ) b.clientConn.UpdateState(balancer.State{ ConnectivityState: b.state, - Picker: &subConnPicker{}, + Picker: defaultSubConnPicker, }) } diff --git a/pkg/balancer/picker.go b/pkg/balancer/picker.go index e6e08e1eb..b590200ff 100644 --- a/pkg/balancer/picker.go +++ b/pkg/balancer/picker.go @@ -12,9 +12,11 @@ import ( ) // subConnPicker is a trivial gRPC picker: reads a pre-computed SubConn from context. -// Installed once in gRPC's balancer state; never needs to be replaced. +// Stateless — a single instance is reused for all UpdateState calls. type subConnPicker struct{} +var defaultSubConnPicker = &subConnPicker{} + func (p *subConnPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { sc, ok := info.Ctx.Value(SubConnKey).(balancer.SubConn) if !ok || sc == nil { diff --git a/pkg/cmd/serve.go b/pkg/cmd/serve.go index 3887d7fd9..d7115c8e6 100644 --- a/pkg/cmd/serve.go +++ b/pkg/cmd/serve.go @@ -415,8 +415,8 @@ func serve() func(cmd *cobra.Command, args []string) error { // Add caching to the schema reader using a decorator schemaReader = cacheproxy.NewSchemaReader(schemaReader, schemaCache) - dataReader = sfproxy.NewDataReader(dataReader) dataReader = semproxy.NewDataReader(dataReader) + dataReader = sfproxy.NewDataReader(dataReader) schemaReader = sfproxy.NewSchemaReader(schemaReader) // Check if circuit breaker should be enabled for services From 5e8d4b3128a02158fb8b9de4db9e86061e8683dc Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Tue, 21 Jul 2026 14:42:27 +0200 Subject: [PATCH 06/13] perf: batch entity lookup --- internal/engines/bulk.go | 81 +++++++++++- internal/engines/entity_filter.go | 210 +++++++++++++++++------------- internal/engines/expand.go | 4 +- internal/engines/lookup.go | 15 ++- internal/engines/utils.go | 6 + pkg/cmd/serve.go | 3 +- 6 files changed, 217 insertions(+), 102 deletions(-) diff --git a/internal/engines/bulk.go b/internal/engines/bulk.go index f01a70d24..0c3193f0a 100644 --- a/internal/engines/bulk.go +++ b/internal/engines/bulk.go @@ -105,6 +105,8 @@ type executionState struct { mu sync.Mutex // results stores the results of all requests in their original order results []base.CheckResult + // requests holds the sorted requests (stored once, reused by callbackWithToken) + requests []BulkCheckerRequest // processedIndex tracks the next result to be processed in order processedIndex int // successCount tracks the number of successful permission checks @@ -273,10 +275,83 @@ func (bc *BulkChecker) ExecuteRequests(size uint32) error { // Main execution en // Initialize execution state for tracking progress bc.executionState = &executionState{ - results: make([]base.CheckResult, len(requests)), - limit: int64(size), + results: make([]base.CheckResult, len(requests)), + requests: requests, + limit: int64(size), + } + + // For entity-type checks: all requests share the same (type, permission, subject). + // Batch them into a single Check call instead of N individual calls. + if bc.typ == BulkCheckerTypeEntity { + return bc.executeBatchEntity(requests, size) + } + + // For subject-type checks: each request has a different subject, process concurrently. + return bc.executeConcurrent(requests, size) +} + +// executeBatchEntity batches all entity IDs into a single Check call. +// All requests share the same (type, permission, subject) — they only differ in entity ID. +func (bc *BulkChecker) executeBatchEntity(requests []BulkCheckerRequest, size uint32) error { + // Separate pre-computed results from those needing a check. + var entityIDs []string + var needCheckIndices []int + + for i, req := range requests { + if req.Result != base.CheckResult_CHECK_RESULT_UNSPECIFIED { + bc.executionState.results[i] = req.Result + } else { + entityIDs = append(entityIDs, req.Request.GetEntity().GetId()) + needCheckIndices = append(needCheckIndices, i) + } } + // Batch check all entities at once. + if len(entityIDs) > 0 { + template := requests[needCheckIndices[0]].Request + batchResp, err := bc.checker.Check(bc.ctx, &invoke.BatchCheckRequest{ + TenantID: template.GetTenantId(), + EntityType: template.GetEntity().GetType(), + EntityIDs: entityIDs, + Permission: template.GetPermission(), + Subject: template.GetSubject(), + Metadata: template.GetMetadata(), + Context: template.GetContext(), + Arguments: template.GetArguments(), + }) + if err != nil { + if isContextError(err) { + return nil + } + return fmt.Errorf("bulk execution failed: %w", err) + } + + for j, idx := range needCheckIndices { + if result, ok := batchResp.Results[entityIDs[j]]; ok { + bc.executionState.results[idx] = result + } else { + bc.executionState.results[idx] = base.CheckResult_CHECK_RESULT_DENIED + } + } + } + + // Process results in order, invoke callback for ALLOWED. + for i := range requests { + result := bc.executionState.results[i] + if result == base.CheckResult_CHECK_RESULT_ALLOWED { + if atomic.LoadInt64(&bc.executionState.successCount) >= int64(size) { + break + } + atomic.AddInt64(&bc.executionState.successCount, 1) + bc.callbackWithToken(i) + } + } + return nil +} + +// executeConcurrent processes requests concurrently (used for subject-type checks +// where each request has a different subject and can't be batched). +func (bc *BulkChecker) executeConcurrent(requests []BulkCheckerRequest, size uint32) error { // Create execution context with cancellation for graceful shutdown execCtx, execCancel := context.WithCancel(bc.ctx) defer execCancel() @@ -434,7 +509,7 @@ func (bc *BulkChecker) processResult(index int, result base.CheckResult) error { // Parameters: // - index: The index of the result in the sorted list func (bc *BulkChecker) callbackWithToken(index int) { - requests := bc.getSortedRequests() + requests := bc.executionState.requests // Validate index bounds if index >= len(requests) { diff --git a/internal/engines/entity_filter.go b/internal/engines/entity_filter.go index ef0af4cfb..7566c44ac 100644 --- a/internal/engines/entity_filter.go +++ b/internal/engines/entity_filter.go @@ -24,40 +24,56 @@ type EntityFilter struct { dataReader storage.DataReader graph *schema.LinkedSchemaGraph + + // maxBatchSize limits how many found entities are processed per batch in recursive calls + maxBatchSize int } // NewEntityFilter creates a new EntityFilter engine -func NewEntityFilter(dataReader storage.DataReader, sch *base.SchemaDefinition) *EntityFilter { +func NewEntityFilter(dataReader storage.DataReader, sch *base.SchemaDefinition, maxBatchSize int) *EntityFilter { + if maxBatchSize <= 0 { + maxBatchSize = _defaultMaxBatchSize + } return &EntityFilter{ - dataReader: dataReader, - graph: schema.NewLinkedGraph(sch), + dataReader: dataReader, + graph: schema.NewLinkedGraph(sch), + maxBatchSize: maxBatchSize, } } // EntityFilter is a method of the EntityFilterEngine struct. It executes a permission request for linked entities. +// subjectIDs allows batching: multiple subject IDs with the same (Type, Relation) are queried in one DB call. +// Pass nil to use request.GetSubject().GetId() as the single subject ID. func (engine *EntityFilter) EntityFilter( ctx context.Context, // A context used for tracing and cancellation. request *base.PermissionEntityFilterRequest, // A permission request for linked entities. + subjectIDs []string, // Batch subject IDs (nil = use request.Subject.Id). visits *VisitsMap, // A map that keeps track of visited entities to avoid infinite loops. publisher *BulkEntityPublisher, // A custom publisher that publishes results in bulk. ) (err error) { // Returns an error if one occurs during execution. + if len(subjectIDs) == 0 { + subjectIDs = []string{request.GetSubject().GetId()} + } + // Check if direct result if request.GetEntrance().GetType() == request.GetSubject().GetType() && request.GetEntrance().GetValue() == request.GetSubject().GetRelation() { - found := &base.Entity{ - Type: request.GetSubject().GetType(), - Id: request.GetSubject().GetId(), - } + for _, id := range subjectIDs { + found := &base.Entity{ + Type: request.GetSubject().GetType(), + Id: id, + } - if !visits.AddPublished(found) { // If the entity and relation has already been visited. - return nil - } + if !visits.AddPublished(found) { // If the entity and relation has already been visited. + continue + } - // If the entity reference is the same as the subject, publish the result directly and return. - publisher.Publish(found, &base.PermissionCheckRequestMetadata{ - SnapToken: request.GetMetadata().GetSnapToken(), - SchemaVersion: request.GetMetadata().GetSchemaVersion(), - Depth: request.GetMetadata().GetDepth(), - }, request.GetContext(), base.CheckResult_CHECK_RESULT_UNSPECIFIED) + // If the entity reference is the same as the subject, publish the result directly and return. + publisher.Publish(found, &base.PermissionCheckRequestMetadata{ + SnapToken: request.GetMetadata().GetSnapToken(), + SchemaVersion: request.GetMetadata().GetSchemaVersion(), + Depth: request.GetMetadata().GetDepth(), + }, request.GetContext(), base.CheckResult_CHECK_RESULT_UNSPECIFIED) + } } // Retrieve linked entrances @@ -86,18 +102,12 @@ func (engine *EntityFilter) EntityFilter( // Switch on the kind of linked entrance. switch entrance.LinkedEntranceKind() { case schema.RelationLinkedEntrance: // If the linked entrance is a relation entrance. - err = engine.relationEntrance(cont, request, entrance, visits, g, publisher) // Call the relation entrance method. + err = engine.relationEntrance(cont, request, entrance, subjectIDs, visits, g, publisher) // Call the relation entrance method. if err != nil { return err } case schema.ComputedUserSetLinkedEntrance: // If the linked entrance is a computed user set entrance. - err = engine.lt(cont, request, &base.EntityAndRelation{ // Call the run method with a new entity and relation. - Entity: &base.Entity{ - Type: entrance.TargetEntrance.GetType(), - Id: request.GetSubject().GetId(), - }, - Relation: entrance.TargetEntrance.GetValue(), - }, visits, g, publisher) + err = engine.processFoundEntities(cont, request, entrance.TargetEntrance.GetType(), entrance.TargetEntrance.GetValue(), subjectIDs, visits, g, publisher) if err != nil { return err } @@ -107,7 +117,7 @@ func (engine *EntityFilter) EntityFilter( return err } case schema.TupleToUserSetLinkedEntrance: // If the linked entrance is a tuple to user set entrance. - err = engine.tupleToUserSetEntrance(cont, request, entrance, visits, g, publisher) // Call the tuple to user set entrance method. + err = engine.tupleToUserSetEntrance(cont, request, entrance, subjectIDs, visits, g, publisher) // Call the tuple to user set entrance method. if err != nil { return err } @@ -378,10 +388,12 @@ func (engine *EntityFilter) expandRecursiveRelation( } // relationEntrance is a method of the EntityFilterEngine struct. It handles relation entrances. +// Uses batch subject IDs in SubjectFilter for a single DB query instead of per-entity queries. func (engine *EntityFilter) relationEntrance( ctx context.Context, // A context used for tracing and cancellation. request *base.PermissionEntityFilterRequest, // A permission request for linked entities. entrance *schema.LinkedEntrance, // A linked entrance. + subjectIDs []string, // Batch subject IDs for SubjectFilter. visits *VisitsMap, // A map that keeps track of visited entities to avoid infinite loops. g *errgroup.Group, // An errgroup used for executing goroutines. publisher *BulkEntityPublisher, // A custom publisher that publishes results in bulk. @@ -408,7 +420,7 @@ func (engine *EntityFilter) relationEntrance( Relation: entrance.TargetEntrance.GetValue(), Subject: &base.SubjectFilter{ Type: request.GetSubject().GetType(), - Ids: []string{request.GetSubject().GetId()}, + Ids: subjectIDs, // Batch: all subject IDs in one query Relation: request.GetSubject().GetRelation(), }, } @@ -446,26 +458,29 @@ func (engine *EntityFilter) relationEntrance( // NewUniqueTupleIterator() ensures that the iterator only returns unique tuples. it := database.NewUniqueTupleIterator(rit, cti) - for it.HasNext() { // Loop over each relationship. - // Get the next tuple's subject. - current, ok := it.GetNext() - if !ok { - break + // Process results in chunks to avoid large intermediate slices. + // All results share the same entity type and relation (from the filter). + for it.HasNext() { + var chunkType, chunkRelation string + chunk := make([]string, 0, engine.maxBatchSize) + for it.HasNext() && len(chunk) < engine.maxBatchSize { + current, ok := it.GetNext() + if !ok { + break + } + chunkType = current.GetEntity().GetType() + chunkRelation = current.GetRelation() + chunk = append(chunk, current.GetEntity().GetId()) + } + if err := engine.processFoundEntities(ctx, request, chunkType, chunkRelation, chunk, visits, g, publisher); err != nil { + return err } - g.Go(func() error { - return engine.lt(ctx, request, &base.EntityAndRelation{ // Call the run method with a new entity and relation. - Entity: &base.Entity{ - Type: current.GetEntity().GetType(), - Id: current.GetEntity().GetId(), - }, - Relation: current.GetRelation(), - }, visits, g, publisher) - }) } return nil } // tupleToUserSetEntrance is a method of the EntityFilterEngine struct. It handles tuple to user set entrances. +// Uses batch subject IDs in SubjectFilter for a single DB query instead of per-entity queries. func (engine *EntityFilter) tupleToUserSetEntrance( // A context used for tracing and cancellation. ctx context.Context, @@ -473,6 +488,8 @@ func (engine *EntityFilter) tupleToUserSetEntrance( request *base.PermissionEntityFilterRequest, // A linked entrance. entrance *schema.LinkedEntrance, + // Batch subject IDs for SubjectFilter. + subjectIDs []string, // A map that keeps track of visited entities to avoid infinite loops. visits *VisitsMap, // An errgroup used for executing goroutines. @@ -502,7 +519,7 @@ func (engine *EntityFilter) tupleToUserSetEntrance( Relation: entrance.TupleSetRelation, // Query for relationships that match the tuple set relation. Subject: &base.SubjectFilter{ Type: request.GetSubject().GetType(), - Ids: []string{request.GetSubject().GetId()}, + Ids: subjectIDs, // Batch: all subject IDs in one query Relation: "", }, } @@ -540,82 +557,93 @@ func (engine *EntityFilter) tupleToUserSetEntrance( // NewUniqueTupleIterator() ensures that the iterator only returns unique tuples. it := database.NewUniqueTupleIterator(rit, cti) - for it.HasNext() { // Loop over each relationship. - // Get the next tuple's subject. - current, ok := it.GetNext() - if !ok { - break + // Process results in chunks to avoid large intermediate slices. + // All results share the same entity type and relation (from the entrance). + entType := entrance.TargetEntrance.GetType() + entRel := entrance.TargetEntrance.GetValue() + for it.HasNext() { + chunk := make([]string, 0, engine.maxBatchSize) + for it.HasNext() && len(chunk) < engine.maxBatchSize { + current, ok := it.GetNext() + if !ok { + break + } + chunk = append(chunk, current.GetEntity().GetId()) + } + if err := engine.processFoundEntities(ctx, request, entType, entRel, chunk, visits, g, publisher); err != nil { + return err } - g.Go(func() error { - return engine.lt(ctx, request, &base.EntityAndRelation{ // Call the run method with a new entity and relation. - Entity: &base.Entity{ - Type: entrance.TargetEntrance.GetType(), - Id: current.GetEntity().GetId(), - }, - Relation: entrance.TargetEntrance.GetValue(), - }, visits, g, publisher) - }) } return nil } -// run is a method of the EntityFilterEngine struct. It executes the linked entity engine for a given request. -func (engine *EntityFilter) lt( - ctx context.Context, // A context used for tracing and cancellation. - request *base.PermissionEntityFilterRequest, // A permission request for linked entities. - found *base.EntityAndRelation, // An entity and relation that was previously found. - visits *VisitsMap, // A map that keeps track of visited entities to avoid infinite loops. - g *errgroup.Group, // An errgroup used for executing goroutines. - publisher *BulkEntityPublisher, // A custom publisher that publishes results in bulk. -) error { // Returns an error if one occurs during execution. - if !visits.AddER(found.GetEntity(), found.GetRelation()) { // If the entity and relation has already been visited. +// processFoundEntities handles batch processing of found entities at a single graph level. +// All founds share the same (entityType, relation). It checks LinkedEntrances once, +// publishes direct matches, and recursively calls EntityFilter with batch subject IDs. +func (engine *EntityFilter) processFoundEntities( + ctx context.Context, + request *base.PermissionEntityFilterRequest, + entityType string, + relation string, + entityIds []string, + visits *VisitsMap, + g *errgroup.Group, + publisher *BulkEntityPublisher, +) error { + // Visit check each entity. + var filtered []string + for _, id := range entityIds { + if visits.AddER(&base.Entity{Type: entityType, Id: id}, relation) { + filtered = append(filtered, id) + } + } + if len(filtered) == 0 { return nil } - var err error - - // Retrieve linked entrances - var entrances []*schema.LinkedEntrance - entrances, err = engine.graph.LinkedEntrances( + // Compute LinkedEntrances once (depends on Type+Relation, not Id). + entrances, err := engine.graph.LinkedEntrances( request.GetEntrance(), - &base.Entrance{ - Type: request.GetSubject().GetType(), - Value: request.GetSubject().GetRelation(), - }, - ) // Retrieve the linked entrances for the request. + &base.Entrance{Type: entityType, Value: relation}, + ) if err != nil { return err } - if entrances == nil { // If there are no linked entrances for the request. - if found.GetEntity().GetType() == request.GetEntrance().GetType() && found.GetRelation() == request.GetEntrance().GetValue() { // Check if the found entity matches the requested entity reference. - if !visits.AddPublished(found.GetEntity()) { // If the entity and relation has already been visited. - return nil + if entrances == nil { + // Direct match: publish all entities. + if entityType == request.GetEntrance().GetType() && relation == request.GetEntrance().GetValue() { + for _, id := range filtered { + entity := &base.Entity{Type: entityType, Id: id} + if !visits.AddPublished(entity) { + continue + } + publisher.Publish(entity, &base.PermissionCheckRequestMetadata{ + SnapToken: request.GetMetadata().GetSnapToken(), + SchemaVersion: request.GetMetadata().GetSchemaVersion(), + Depth: request.GetMetadata().GetDepth(), + }, request.GetContext(), base.CheckResult_CHECK_RESULT_UNSPECIFIED) } - publisher.Publish(found.GetEntity(), &base.PermissionCheckRequestMetadata{ // Publish the found entity with the permission check metadata. - SnapToken: request.GetMetadata().GetSnapToken(), - SchemaVersion: request.GetMetadata().GetSchemaVersion(), - Depth: request.GetMetadata().GetDepth(), - }, request.GetContext(), base.CheckResult_CHECK_RESULT_UNSPECIFIED) - return nil } - return nil // Otherwise, return without publishing any results. + return nil } + // Needs recursion: ONE EntityFilter call with all IDs as batch subject IDs. + // Input is already chunked by the caller (relationEntrance/tupleToUserSetEntrance). g.Go(func() error { - return engine.EntityFilter(ctx, &base.PermissionEntityFilterRequest{ // Call the Run method recursively with a new permission request. + return engine.EntityFilter(ctx, &base.PermissionEntityFilterRequest{ TenantId: request.GetTenantId(), Entrance: request.GetEntrance(), Subject: &base.Subject{ - Type: found.GetEntity().GetType(), - Id: found.GetEntity().GetId(), - Relation: found.GetRelation(), + Type: entityType, + Id: filtered[0], // Representative ID for Subject field + Relation: relation, }, Scope: request.GetScope(), Metadata: request.GetMetadata(), Context: request.GetContext(), Cursor: request.GetCursor(), - }, visits, publisher) + }, filtered, visits, publisher) }) return nil } diff --git a/internal/engines/expand.go b/internal/engines/expand.go index 24b778b0c..a789fafbe 100644 --- a/internal/engines/expand.go +++ b/internal/engines/expand.go @@ -24,7 +24,9 @@ type ExpandEngine struct { dataReader storage.DataReader } -// NewExpandEngine creates a new instance of ExpandEngine. +// NewExpandEngine - This function creates a new instance of ExpandEngine by taking a SchemaReader and a RelationshipReader as +// parameters and returning a pointer to the created instance. The SchemaReader is used to read schema definitions, while the +// RelationshipReader is used to read relationship definitions. func NewExpandEngine(sr storage.SchemaReader, rr storage.DataReader) *ExpandEngine { return &ExpandEngine{ schemaReader: sr, diff --git a/internal/engines/lookup.go b/internal/engines/lookup.go index 254af94d9..f0a31d6c2 100644 --- a/internal/engines/lookup.go +++ b/internal/engines/lookup.go @@ -26,6 +26,8 @@ type LookupEngine struct { schemaMap sync.Map // concurrencyLimit is the maximum number of concurrent permission checks allowed concurrencyLimit int + // maxBatchSize is the maximum number of subject IDs per batch query in EntityFilter + maxBatchSize int } func NewLookupEngine( @@ -40,6 +42,7 @@ func NewLookupEngine( dataReader: dataReader, schemaMap: sync.Map{}, concurrencyLimit: _defaultConcurrencyLimit, + maxBatchSize: _defaultMaxBatchSize, } // options @@ -75,7 +78,7 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm // Create configuration for BulkChecker config := BulkCheckerConfig{ ConcurrencyLimit: engine.concurrencyLimit, - BufferSize: 1000, + BufferSize: engine.maxBatchSize, } // Create and start BulkChecker. It performs permission checks in parallel. @@ -99,7 +102,7 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm visits := &VisitsMap{} // Perform an entity filter operation based on the permission request - err = NewEntityFilter(engine.dataReader, sc).EntityFilter(ctx, &base.PermissionEntityFilterRequest{ + err = NewEntityFilter(engine.dataReader, sc, engine.maxBatchSize).EntityFilter(ctx, &base.PermissionEntityFilterRequest{ TenantId: request.GetTenantId(), Metadata: &base.PermissionEntityFilterRequestMetadata{ SnapToken: request.GetMetadata().GetSnapToken(), @@ -114,7 +117,7 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm Context: request.GetContext(), Scope: request.GetScope(), Cursor: request.GetContinuousToken(), - }, visits, publisher) + }, nil, visits, publisher) if err != nil { return nil, err } @@ -156,7 +159,7 @@ func (engine *LookupEngine) LookupEntityStream(ctx context.Context, request *bas // Create configuration for BulkChecker config := BulkCheckerConfig{ ConcurrencyLimit: engine.concurrencyLimit, - BufferSize: 1000, + BufferSize: engine.maxBatchSize, } // Create and start BulkChecker. It performs permission checks concurrently. @@ -179,7 +182,7 @@ func (engine *LookupEngine) LookupEntityStream(ctx context.Context, request *bas visits := &VisitsMap{} // Perform an entity filter operation based on the permission request - err = NewEntityFilter(engine.dataReader, sc).EntityFilter(ctx, &base.PermissionEntityFilterRequest{ + err = NewEntityFilter(engine.dataReader, sc, engine.maxBatchSize).EntityFilter(ctx, &base.PermissionEntityFilterRequest{ TenantId: request.GetTenantId(), Metadata: &base.PermissionEntityFilterRequestMetadata{ SnapToken: request.GetMetadata().GetSnapToken(), @@ -193,7 +196,7 @@ func (engine *LookupEngine) LookupEntityStream(ctx context.Context, request *bas Subject: request.GetSubject(), Context: request.GetContext(), Cursor: request.GetContinuousToken(), - }, visits, publisher) + }, nil, visits, publisher) if err != nil { return err } diff --git a/internal/engines/utils.go b/internal/engines/utils.go index 215fe42b3..999518bf5 100644 --- a/internal/engines/utils.go +++ b/internal/engines/utils.go @@ -46,6 +46,12 @@ func LookupConcurrencyLimit(limit int) LookupOption { } } +func LookupMaxBatchSize(size int) LookupOption { + return func(c *LookupEngine) { + c.maxBatchSize = size + } +} + // SubjectFilterOption - a functional option type for configuring the LookupSubjectEngine. type SubjectFilterOption func(engine *SubjectFilter) diff --git a/pkg/cmd/serve.go b/pkg/cmd/serve.go index d7115c8e6..1f70bf6d1 100644 --- a/pkg/cmd/serve.go +++ b/pkg/cmd/serve.go @@ -496,7 +496,8 @@ func serve() func(cmd *cobra.Command, args []string) error { schemaReader, dataReader, // Set concurrency limit based on the configuration. - engines.LookupConcurrencyLimit(cfg.Service.Permission.BulkLimit), + engines.LookupConcurrencyLimit(cfg.Service.Permission.ConcurrencyLimit), + engines.LookupMaxBatchSize(cfg.Service.Permission.BulkLimit), ) // Initialize the subjectPermissionEngine, responsible for handling subject permissions. From a9afa41257d3d5d5f2b2d515060c6284db3ab68d Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Tue, 21 Jul 2026 15:23:36 +0200 Subject: [PATCH 07/13] perf: batch entity lookup (early termination) --- internal/engines/bulk.go | 118 +++++++++++++++++++- internal/engines/entity_filter.go | 6 ++ internal/engines/lookup.go | 65 +++++++++-- internal/engines/lookup_test.go | 172 ++++++++++++++++++++++++++++++ 4 files changed, 350 insertions(+), 11 deletions(-) diff --git a/internal/engines/bulk.go b/internal/engines/bulk.go index 0c3193f0a..40aee5fdf 100644 --- a/internal/engines/bulk.go +++ b/internal/engines/bulk.go @@ -49,6 +49,9 @@ type BulkCheckerConfig struct { // BufferSize defines the size of the internal request buffer. // This should be set based on expected request volume to avoid blocking. BufferSize int + // Streaming disables background request collection. When true, ExecuteStreamingRequests + // reads directly from the request channel, enabling early termination. + Streaming bool } // DefaultBulkCheckerConfig returns a sensible default configuration @@ -165,8 +168,14 @@ func NewBulkChecker(ctx context.Context, checker invoke.Check, typ BulkCheckerTy collectionDone: make(chan struct{}), } - // Start the background request collection goroutine - go bc.collectRequests() + // Start the background request collection goroutine (not needed in streaming mode). + if config.Streaming { + // No collector goroutine — ExecuteStreamingRequests reads directly from the channel. + // Close collectionDone immediately so StopCollectingRequests doesn't block. + close(bc.collectionDone) + } else { + go bc.collectRequests() + } return bc, nil } @@ -247,6 +256,111 @@ func (bc *BulkChecker) sortRequests(requests []BulkCheckerRequest) { } } +// ExecuteStreamingRequests processes candidates as they arrive from EntityFilter, +// checking them in batches without waiting for all candidates to be collected. +// When size ALLOWED results are found, it cancels the context to stop EntityFilter. +// No ordering is applied — results are returned in arrival order. +// +// NOTE: This method reads directly from requestChan. The background collectRequests +// goroutine must be stopped first by calling StopCollectingRequests or cancelling context +// before the producer (EntityFilter) starts. In practice, the producer goroutine is started +// AFTER this method begins reading, so collectRequests sees a closed channel or cancelled context. +func (bc *BulkChecker) ExecuteStreamingRequests(size uint32) error { + // size=0 means no early termination — process all candidates. + var successCount int64 + + for { + // Collect a chunk from the channel. + chunk := make([]BulkCheckerRequest, 0, bc.config.BufferSize) + collect: + for len(chunk) < bc.config.BufferSize { + select { + case req, ok := <-bc.requestChan: + if !ok { + break collect // Channel closed, process remaining chunk + } + chunk = append(chunk, req) + case <-bc.ctx.Done(): + return nil + default: + if len(chunk) > 0 { + break collect // No more pending, process what we have + } + // Nothing yet, block-wait for at least one + select { + case req, ok := <-bc.requestChan: + if !ok { + break collect + } + chunk = append(chunk, req) + case <-bc.ctx.Done(): + return nil + } + } + } + + if len(chunk) == 0 { + return nil // Channel closed, nothing left + } + + // Separate pre-computed from needs-check. + var entityIDs []string + var needCheckIndices []int + results := make([]base.CheckResult, len(chunk)) + + for i, req := range chunk { + if req.Result != base.CheckResult_CHECK_RESULT_UNSPECIFIED { + results[i] = req.Result + } else { + entityIDs = append(entityIDs, req.Request.GetEntity().GetId()) + needCheckIndices = append(needCheckIndices, i) + } + } + + // Batch check. + if len(entityIDs) > 0 { + template := chunk[needCheckIndices[0]].Request + batchResp, err := bc.checker.Check(bc.ctx, &invoke.BatchCheckRequest{ + TenantID: template.GetTenantId(), + EntityType: template.GetEntity().GetType(), + EntityIDs: entityIDs, + Permission: template.GetPermission(), + Subject: template.GetSubject(), + Metadata: template.GetMetadata(), + Context: template.GetContext(), + Arguments: template.GetArguments(), + }) + if err != nil { + if isContextError(err) { + return nil + } + return fmt.Errorf("streaming bulk execution failed: %w", err) + } + + for j, idx := range needCheckIndices { + if result, ok := batchResp.Results[entityIDs[j]]; ok { + results[idx] = result + } else { + results[idx] = base.CheckResult_CHECK_RESULT_DENIED + } + } + } + + // Process results, callback for ALLOWED. + for i, req := range chunk { + if results[i] == base.CheckResult_CHECK_RESULT_ALLOWED { + id := req.Request.GetEntity().GetId() + bc.callback(id, "") + successCount++ + if size > 0 && successCount >= int64(size) { + bc.cancel() // Stop EntityFilter + return nil + } + } + } + } +} + // ExecuteRequests processes requests concurrently with comprehensive error handling and resource management. // This method is the main entry point for bulk permission checking. It: // 1. Stops collecting new requests diff --git a/internal/engines/entity_filter.go b/internal/engines/entity_filter.go index 7566c44ac..56a8e63cf 100644 --- a/internal/engines/entity_filter.go +++ b/internal/engines/entity_filter.go @@ -185,6 +185,12 @@ func (engine *EntityFilter) attributeEntrance( it := database.NewUniqueAttributeIterator(rit, cti) + // Only publish entities of the target type (the type we're looking up). + // Attribute entrances on intermediate types are not candidates. + if entrance.TargetEntrance.GetType() != request.GetEntrance().GetType() { + return nil + } + var attributeEntityIDs []string attributeEntityIDSet := make(map[string]struct{}) diff --git a/internal/engines/lookup.go b/internal/engines/lookup.go index f0a31d6c2..e463a1ee3 100644 --- a/internal/engines/lookup.go +++ b/internal/engines/lookup.go @@ -8,6 +8,8 @@ import ( "strings" "sync" + "google.golang.org/grpc/metadata" + "github.com/Permify/permify/internal/invoke" "github.com/Permify/permify/internal/storage" "github.com/Permify/permify/internal/storage/context/utils" @@ -75,10 +77,13 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm ct = token } + streaming := skipOrdering(ctx) + // Create configuration for BulkChecker config := BulkCheckerConfig{ ConcurrencyLimit: engine.concurrencyLimit, BufferSize: engine.maxBatchSize, + Streaming: streaming, } // Create and start BulkChecker. It performs permission checks in parallel. @@ -101,8 +106,7 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm // Create a map to keep track of visited entities visits := &VisitsMap{} - // Perform an entity filter operation based on the permission request - err = NewEntityFilter(engine.dataReader, sc, engine.maxBatchSize).EntityFilter(ctx, &base.PermissionEntityFilterRequest{ + filterRequest := &base.PermissionEntityFilterRequest{ TenantId: request.GetTenantId(), Metadata: &base.PermissionEntityFilterRequestMetadata{ SnapToken: request.GetMetadata().GetSnapToken(), @@ -117,15 +121,48 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm Context: request.GetContext(), Scope: request.GetScope(), Cursor: request.GetContinuousToken(), - }, nil, visits, publisher) - if err != nil { - return nil, err } - // At this point, the BulkChecker has collected and sorted requests - err = checker.ExecuteRequests(size) // Execute the collected requests in parallel - if err != nil { - return nil, err + ef := NewEntityFilter(engine.dataReader, sc, engine.maxBatchSize) + + if streaming { + // Streaming mode (x-permify-skip-ordering): EntityFilter and permission checks + // run in parallel. Candidates are checked in batches as they arrive, and processing + // stops as soon as enough ALLOWED results are found. + filterErrCh := make(chan error, 1) + go func() { + filterErrCh <- ef.EntityFilter(ctx, filterRequest, nil, visits, publisher) + checker.StopCollectingRequests() + }() + + err = checker.ExecuteStreamingRequests(size) + if err != nil { + return nil, err + } + + // Check EntityFilter error — non-blocking since it may still be running + // (context was cancelled for early termination, EntityFilter will finish eventually). + select { + case filterErr := <-filterErrCh: + if filterErr != nil && !isContextError(filterErr) { + return nil, filterErr + } + default: + // EntityFilter still running — that's fine, it will exit via cancelled context. + } + + ct = "" + } else { + // Standard mode: collect all candidates, sort, then check. + err = ef.EntityFilter(ctx, filterRequest, nil, visits, publisher) + if err != nil { + return nil, err + } + + err = checker.ExecuteRequests(size) + if err != nil { + return nil, err + } } // Return response containing allowed entity IDs @@ -313,3 +350,13 @@ func (engine *LookupEngine) readSchema(ctx context.Context, tenantID, schemaVers // Return the freshly read schema. return sch, nil } + +// skipOrdering checks gRPC metadata for the x-permify-skip-ordering flag. +func skipOrdering(ctx context.Context) bool { + md, ok := metadata.FromIncomingContext(ctx) + if !ok { + return false + } + vals := md.Get("x-permify-skip-ordering") + return len(vals) > 0 && vals[0] == "true" +} diff --git a/internal/engines/lookup_test.go b/internal/engines/lookup_test.go index adb0a74fd..897056f6c 100644 --- a/internal/engines/lookup_test.go +++ b/internal/engines/lookup_test.go @@ -6,6 +6,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "google.golang.org/grpc/metadata" "google.golang.org/protobuf/types/known/structpb" "github.com/Permify/permify/internal/config" @@ -6190,6 +6191,177 @@ entity group_perms { }) }) + Context("Streaming Mode (skip-ordering)", func() { + It("Drive Sample: streaming returns same results as standard mode", func() { + db, err := factories.DatabaseFactory( + config.Database{ + Engine: "memory", + }, + ) + + Expect(err).ShouldNot(HaveOccurred()) + + conf, err := newSchema(driveSchemaEntityFilter) + Expect(err).ShouldNot(HaveOccurred()) + + schemaWriter := factories.SchemaWriterFactory(db) + err = schemaWriter.WriteSchema(context.Background(), conf) + Expect(err).ShouldNot(HaveOccurred()) + + schemaReader := factories.SchemaReaderFactory(db) + dataReader := factories.DataReaderFactory(db) + dataWriter := factories.DataWriterFactory(db) + + checkEngine := NewCheckEngine(schemaReader, dataReader) + + lookupEngine := NewLookupEngine( + checkEngine, + schemaReader, + dataReader, + ) + + invoker := invoke.NewDirectInvoker( + schemaReader, + dataReader, + checkEngine, + nil, + lookupEngine, + nil, + ) + + checkEngine.SetInvoker(invoker) + + var tuples []*base.Tuple + for _, relationship := range []string{ + "doc:1#owner@user:2", + "doc:1#folder@user:3", + "doc:2#owner@user:1", + "doc:3#owner@user:1", + } { + t, err := tuple.Tuple(relationship) + Expect(err).ShouldNot(HaveOccurred()) + tuples = append(tuples, t) + } + + _, err = dataWriter.Write(context.Background(), "t1", database.NewTupleCollection(tuples...), database.NewAttributeCollection()) + Expect(err).ShouldNot(HaveOccurred()) + + // Standard mode + standardResp, err := invoker.LookupEntity(context.Background(), &base.PermissionLookupEntityRequest{ + TenantId: "t1", + EntityType: "doc", + Subject: &base.Subject{ + Type: "user", + Id: "1", + }, + Permission: "read", + Metadata: &base.PermissionLookupEntityRequestMetadata{ + SnapToken: token.NewNoopToken().Encode().String(), + SchemaVersion: "", + Depth: 100, + }, + }) + Expect(err).ShouldNot(HaveOccurred()) + + // Streaming mode via gRPC metadata + streamingCtx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-permify-skip-ordering", "true")) + streamingResp, err := invoker.LookupEntity(streamingCtx, &base.PermissionLookupEntityRequest{ + TenantId: "t1", + EntityType: "doc", + Subject: &base.Subject{ + Type: "user", + Id: "1", + }, + Permission: "read", + Metadata: &base.PermissionLookupEntityRequestMetadata{ + SnapToken: token.NewNoopToken().Encode().String(), + SchemaVersion: "", + Depth: 100, + }, + }) + Expect(err).ShouldNot(HaveOccurred()) + + // Same entity IDs (order may differ) + Expect(streamingResp.GetEntityIds()).Should(ConsistOf(standardResp.GetEntityIds())) + // Streaming returns unordered token + Expect(streamingResp.GetContinuousToken()).Should(Equal("")) + }) + + It("Drive Sample: streaming with page_size=1 returns exactly 1 result", func() { + db, err := factories.DatabaseFactory( + config.Database{ + Engine: "memory", + }, + ) + + Expect(err).ShouldNot(HaveOccurred()) + + conf, err := newSchema(driveSchemaEntityFilter) + Expect(err).ShouldNot(HaveOccurred()) + + schemaWriter := factories.SchemaWriterFactory(db) + err = schemaWriter.WriteSchema(context.Background(), conf) + Expect(err).ShouldNot(HaveOccurred()) + + schemaReader := factories.SchemaReaderFactory(db) + dataReader := factories.DataReaderFactory(db) + dataWriter := factories.DataWriterFactory(db) + + checkEngine := NewCheckEngine(schemaReader, dataReader) + + lookupEngine := NewLookupEngine( + checkEngine, + schemaReader, + dataReader, + ) + + invoker := invoke.NewDirectInvoker( + schemaReader, + dataReader, + checkEngine, + nil, + lookupEngine, + nil, + ) + + checkEngine.SetInvoker(invoker) + + var tuples []*base.Tuple + for _, relationship := range []string{ + "doc:1#owner@user:1", + "doc:2#owner@user:1", + "doc:3#owner@user:1", + } { + t, err := tuple.Tuple(relationship) + Expect(err).ShouldNot(HaveOccurred()) + tuples = append(tuples, t) + } + + _, err = dataWriter.Write(context.Background(), "t1", database.NewTupleCollection(tuples...), database.NewAttributeCollection()) + Expect(err).ShouldNot(HaveOccurred()) + + streamingCtx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-permify-skip-ordering", "true")) + resp, err := invoker.LookupEntity(streamingCtx, &base.PermissionLookupEntityRequest{ + TenantId: "t1", + EntityType: "doc", + Subject: &base.Subject{ + Type: "user", + Id: "1", + }, + Permission: "read", + Metadata: &base.PermissionLookupEntityRequestMetadata{ + SnapToken: token.NewNoopToken().Encode().String(), + SchemaVersion: "", + Depth: 100, + }, + PageSize: 1, + }) + Expect(err).ShouldNot(HaveOccurred()) + Expect(resp.GetEntityIds()).Should(HaveLen(1)) + Expect(resp.GetEntityIds()[0]).Should(BeElementOf("1", "2", "3")) + }) + }) + Context("Recursive Attribute Lookup", func() { It("should include same-type recursive attribute permissions", func() { schema := ` From fd92637f5a94eb85fe9a8d1ebf383a7f0626fc60 Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Tue, 21 Jul 2026 17:50:32 +0200 Subject: [PATCH 08/13] fix: review fixes --- internal/engines/bulk.go | 6 ++++++ internal/engines/entity_filter.go | 12 ++++++------ internal/engines/lookup.go | 4 +++- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/internal/engines/bulk.go b/internal/engines/bulk.go index 40aee5fdf..736c512a9 100644 --- a/internal/engines/bulk.go +++ b/internal/engines/bulk.go @@ -651,6 +651,12 @@ func (bc *BulkChecker) callbackWithToken(index int) { bc.callback(id, ct) } +// Context returns the BulkChecker's cancellable context. +// Use this to run producers (e.g. EntityFilter) so they stop when the checker cancels. +func (bc *BulkChecker) Context() context.Context { + return bc.ctx +} + // Close properly cleans up resources and cancels all operations. // This method should be called when the BulkChecker is no longer needed // to ensure proper resource cleanup and prevent goroutine leaks. diff --git a/internal/engines/entity_filter.go b/internal/engines/entity_filter.go index 56a8e63cf..b786dc63a 100644 --- a/internal/engines/entity_filter.go +++ b/internal/engines/entity_filter.go @@ -142,6 +142,12 @@ func (engine *EntityFilter) attributeEntrance( visits *VisitsMap, // A map that keeps track of visited entities to avoid infinite loops. publisher *BulkEntityPublisher, // A custom publisher that publishes results in bulk. ) error { // Returns an error if one occurs during execution. + // Only publish entities of the target type (the type we're looking up). + // Attribute entrances on intermediate types are not candidates — skip DB queries. + if entrance.TargetEntrance.GetType() != request.GetEntrance().GetType() { + return nil + } + // attributeEntrance only handles direct attribute access if !visits.AddEA(entrance.TargetEntrance.GetType(), entrance.TargetEntrance.GetValue()) { return nil @@ -185,12 +191,6 @@ func (engine *EntityFilter) attributeEntrance( it := database.NewUniqueAttributeIterator(rit, cti) - // Only publish entities of the target type (the type we're looking up). - // Attribute entrances on intermediate types are not candidates. - if entrance.TargetEntrance.GetType() != request.GetEntrance().GetType() { - return nil - } - var attributeEntityIDs []string attributeEntityIDSet := make(map[string]struct{}) diff --git a/internal/engines/lookup.go b/internal/engines/lookup.go index e463a1ee3..ab269b6be 100644 --- a/internal/engines/lookup.go +++ b/internal/engines/lookup.go @@ -129,9 +129,11 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm // Streaming mode (x-permify-skip-ordering): EntityFilter and permission checks // run in parallel. Candidates are checked in batches as they arrive, and processing // stops as soon as enough ALLOWED results are found. + // Use checker's context so bc.cancel() stops EntityFilter too. + checkerCtx := checker.Context() filterErrCh := make(chan error, 1) go func() { - filterErrCh <- ef.EntityFilter(ctx, filterRequest, nil, visits, publisher) + filterErrCh <- ef.EntityFilter(checkerCtx, filterRequest, nil, visits, publisher) checker.StopCollectingRequests() }() From 547a44298cf6e440026ab8d852e8e5191fd5ba8b Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Wed, 22 Jul 2026 13:41:21 +0200 Subject: [PATCH 09/13] fix: panic: send on closed channel --- internal/engines/bulk.go | 54 +++++++++++++++++++++++--------------- internal/engines/lookup.go | 2 +- 2 files changed, 34 insertions(+), 22 deletions(-) diff --git a/internal/engines/bulk.go b/internal/engines/bulk.go index 736c512a9..dce627e6a 100644 --- a/internal/engines/bulk.go +++ b/internal/engines/bulk.go @@ -94,6 +94,9 @@ type BulkChecker struct { executionState *executionState // collectionDone signals when request collection has completed collectionDone chan struct{} + // producerDone signals when the producer (EntityFilter) has finished sending requests. + // Used in streaming mode instead of closing requestChan to avoid send-on-closed-channel races. + producerDone chan struct{} // Callback for processing results // callback is invoked for each successful permission check with the entity/subject ID and continuous token @@ -166,6 +169,7 @@ func NewBulkChecker(ctx context.Context, checker invoke.Check, typ BulkCheckerTy requests: make([]BulkCheckerRequest, 0, config.BufferSize), callback: callback, collectionDone: make(chan struct{}), + producerDone: make(chan struct{}), } // Start the background request collection goroutine (not needed in streaming mode). @@ -271,36 +275,34 @@ func (bc *BulkChecker) ExecuteStreamingRequests(size uint32) error { for { // Collect a chunk from the channel. + // We never close requestChan (to avoid send-on-closed-channel races). + // Instead, producerDone signals that no more items will be sent. chunk := make([]BulkCheckerRequest, 0, bc.config.BufferSize) - collect: + producerFinished := false + + // Block-wait for at least one item, producer done, or context cancel. + select { + case req := <-bc.requestChan: + chunk = append(chunk, req) + case <-bc.producerDone: + producerFinished = true + case <-bc.ctx.Done(): + return nil + } + + // Drain all immediately available items from the channel (non-blocking). + draining: for len(chunk) < bc.config.BufferSize { select { - case req, ok := <-bc.requestChan: - if !ok { - break collect // Channel closed, process remaining chunk - } + case req := <-bc.requestChan: chunk = append(chunk, req) - case <-bc.ctx.Done(): - return nil default: - if len(chunk) > 0 { - break collect // No more pending, process what we have - } - // Nothing yet, block-wait for at least one - select { - case req, ok := <-bc.requestChan: - if !ok { - break collect - } - chunk = append(chunk, req) - case <-bc.ctx.Done(): - return nil - } + break draining } } if len(chunk) == 0 { - return nil // Channel closed, nothing left + return nil // Producer done, nothing left } // Separate pre-computed from needs-check. @@ -358,6 +360,10 @@ func (bc *BulkChecker) ExecuteStreamingRequests(size uint32) error { } } } + + if producerFinished { + return nil + } } } @@ -651,6 +657,12 @@ func (bc *BulkChecker) callbackWithToken(index int) { bc.callback(id, ct) } +// SignalProducerDone signals that the producer has finished sending requests. +// Used in streaming mode — ExecuteStreamingRequests will drain remaining items and exit. +func (bc *BulkChecker) SignalProducerDone() { + close(bc.producerDone) +} + // Context returns the BulkChecker's cancellable context. // Use this to run producers (e.g. EntityFilter) so they stop when the checker cancels. func (bc *BulkChecker) Context() context.Context { diff --git a/internal/engines/lookup.go b/internal/engines/lookup.go index ab269b6be..460ad6d8d 100644 --- a/internal/engines/lookup.go +++ b/internal/engines/lookup.go @@ -134,7 +134,7 @@ func (engine *LookupEngine) LookupEntity(ctx context.Context, request *base.Perm filterErrCh := make(chan error, 1) go func() { filterErrCh <- ef.EntityFilter(checkerCtx, filterRequest, nil, visits, publisher) - checker.StopCollectingRequests() + checker.SignalProducerDone() }() err = checker.ExecuteStreamingRequests(size) From 6664cfc273908d242ac968f3ec17c37020232c9e Mon Sep 17 00:00:00 2001 From: Maxim Manuylov Date: Wed, 22 Jul 2026 14:30:56 +0200 Subject: [PATCH 10/13] fix: review fixes --- internal/engines/bulk.go | 1 + 1 file changed, 1 insertion(+) diff --git a/internal/engines/bulk.go b/internal/engines/bulk.go index dce627e6a..7736f5c42 100644 --- a/internal/engines/bulk.go +++ b/internal/engines/bulk.go @@ -660,6 +660,7 @@ func (bc *BulkChecker) callbackWithToken(index int) { // SignalProducerDone signals that the producer has finished sending requests. // Used in streaming mode — ExecuteStreamingRequests will drain remaining items and exit. func (bc *BulkChecker) SignalProducerDone() { + defer func() { _ = recover() }() close(bc.producerDone) } From 69d7ae203cdc5ff1372685bf11538a19e15284e9 Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Mon, 3 Aug 2026 19:07:09 +0200 Subject: [PATCH 11/13] fix: test compilation --- internal/engines/check_test.go | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/internal/engines/check_test.go b/internal/engines/check_test.go index c41d43a13..d2976ac0a 100644 --- a/internal/engines/check_test.go +++ b/internal/engines/check_test.go @@ -2426,7 +2426,7 @@ var _ = Describe("check-engine", func() { ) Expect(err).ShouldNot(HaveOccurred()) - resp, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + resp, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: &base.Entity{Type: "resource", Id: "r1"}, Permission: "view", @@ -2436,9 +2436,9 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(resp.GetCan()).To(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(resp.UnionResult()).To(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) }) It("should allow cross-type recursive attribute permissions", func() { @@ -2504,7 +2504,7 @@ var _ = Describe("check-engine", func() { ) Expect(err).ShouldNot(HaveOccurred()) - resp, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + resp, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: &base.Entity{Type: "resource", Id: "r1"}, Permission: "view", @@ -2514,9 +2514,9 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(resp.GetCan()).To(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(resp.UnionResult()).To(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) }) It("should allow mixed-entrance recursive attribute permissions", func() { @@ -2574,7 +2574,7 @@ var _ = Describe("check-engine", func() { ) Expect(err).ShouldNot(HaveOccurred()) - resp, err := invoker.Check(context.Background(), &base.PermissionCheckRequest{ + resp, err := invoker.Check(context.Background(), invoke.NewBatchCheckRequest(&base.PermissionCheckRequest{ TenantId: "t1", Entity: &base.Entity{Type: "resource", Id: "zc"}, Permission: "view", @@ -2584,9 +2584,9 @@ var _ = Describe("check-engine", func() { SchemaVersion: "", Depth: 20, }, - }) + })) Expect(err).ShouldNot(HaveOccurred()) - Expect(resp.GetCan()).To(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) + Expect(resp.UnionResult()).To(Equal(base.CheckResult_CHECK_RESULT_ALLOWED)) }) }) }) From 8e07b4fb18dde47d9c744e1adf088b0f1ba811a2 Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Mon, 3 Aug 2026 19:53:15 +0200 Subject: [PATCH 12/13] fix: test compilation --- internal/servers/server_behavior_test.go | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/internal/servers/server_behavior_test.go b/internal/servers/server_behavior_test.go index 9c91fc50c..307a08e1b 100644 --- a/internal/servers/server_behavior_test.go +++ b/internal/servers/server_behavior_test.go @@ -17,6 +17,7 @@ import ( health "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/status" + "github.com/Permify/permify/internal/invoke" "github.com/Permify/permify/internal/storage" "github.com/Permify/permify/pkg/database" v1 "github.com/Permify/permify/pkg/pb/base/v1" @@ -106,24 +107,19 @@ func (f *fakeBundleStore) Delete(_ context.Context, tenantID, name string) error type fakePermissionInvoker struct { err error - checkReq *v1.PermissionCheckRequest + checkReq *invoke.BatchCheckRequest expandReq *v1.PermissionExpandRequest lookupEntityReq *v1.PermissionLookupEntityRequest lookupSubjectReq *v1.PermissionLookupSubjectRequest subjectPermissionReq *v1.PermissionSubjectPermissionRequest } -func (f *fakePermissionInvoker) Check(_ context.Context, request *v1.PermissionCheckRequest) (*v1.PermissionCheckResponse, error) { +func (f *fakePermissionInvoker) Check(_ context.Context, request *invoke.BatchCheckRequest) (*invoke.BatchCheckResponse, error) { f.checkReq = request if f.err != nil { return nil, f.err } - return &v1.PermissionCheckResponse{ - Can: v1.CheckResult_CHECK_RESULT_ALLOWED, - Metadata: &v1.PermissionCheckResponseMetadata{ - CheckCount: 1, - }, - }, nil + return invoke.NewBatchCheckResponse(v1.CheckResult_CHECK_RESULT_ALLOWED, request.EntityIDs...), nil } func (f *fakePermissionInvoker) Expand(_ context.Context, request *v1.PermissionExpandRequest) (*v1.PermissionExpandResponse, error) { @@ -334,7 +330,7 @@ func TestInterceptorLogger(t *testing.T) { func TestPermissionServerPassesThroughInvoker(t *testing.T) { invoker := &fakePermissionInvoker{} - server := NewPermissionServer(invoker) + server := NewPermissionServer(invoker, 0) if server == nil { t.Fatal("expected permission server") } @@ -344,7 +340,7 @@ func TestPermissionServerPassesThroughInvoker(t *testing.T) { if err != nil { t.Fatalf("unexpected check error: %v", err) } - if invoker.checkReq != checkReq || checkResp.GetCan() != v1.CheckResult_CHECK_RESULT_ALLOWED { + if checkResp.GetCan() != v1.CheckResult_CHECK_RESULT_ALLOWED { t.Fatalf("check did not return invoker response") } @@ -388,7 +384,7 @@ func TestPermissionServerPassesThroughInvoker(t *testing.T) { func TestPermissionServerValidationAndInvokerErrors(t *testing.T) { invoker := &fakePermissionInvoker{} - server := NewPermissionServer(invoker) + server := NewPermissionServer(invoker, 0) _, err := server.Check(context.Background(), &v1.PermissionCheckRequest{}) if err == nil { From 8f47504da5db6024596ce7abba6da94f5f0f222e Mon Sep 17 00:00:00 2001 From: "maxim.manuylov" Date: Tue, 4 Aug 2026 10:47:30 +0200 Subject: [PATCH 13/13] fix: review fixes --- internal/engines/check.go | 27 +++++++++++++++--------- internal/servers/permission_server.go | 12 ++++++++--- internal/servers/server.go | 6 ++++-- internal/servers/server_behavior_test.go | 4 ++-- pkg/cmd/serve.go | 1 + 5 files changed, 33 insertions(+), 17 deletions(-) diff --git a/internal/engines/check.go b/internal/engines/check.go index 3fad02452..7f19dd8d5 100644 --- a/internal/engines/check.go +++ b/internal/engines/check.go @@ -87,8 +87,11 @@ type CheckFunction func(ctx context.Context) (*invoke.BatchCheckResponse, error) // and a slice of CheckFunctions. It combines the per-entity results of // multiple CheckFunctions according to a specific strategy and returns // a BatchCheckResponse along with an error. +// expectedEntityCount is the total number of unique entity IDs expected across +// all functions. It is used by checkUnion to gate early exit correctly when +// functions operate on disjoint entity sets. // Concurrency is controlled by a request-scoped semaphore stored in the context. -type CheckCombiner func(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) +type CheckCombiner func(ctx context.Context, functions []CheckFunction, limit int, expectedEntityCount int) (*invoke.BatchCheckResponse, error) // invoke creates a CheckFunction that invokes a batch check through the full invoke chain // (DirectInvoker -> Cache -> CheckEngine), ensuring depth tracking, caching, and tracing. @@ -175,7 +178,7 @@ func (engine *CheckEngine) check( // Otherwise, return a CheckFunction that checks a union of CheckFunctions. return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { - result, err := checkUnion(ctx, []CheckFunction{fn}, engine.concurrencyLimit) + result, err := checkUnion(ctx, []CheckFunction{fn}, engine.concurrencyLimit, len(request.EntityIDs)) if err != nil { return result, err } @@ -268,7 +271,7 @@ func (engine *CheckEngine) setChild( // (union, intersection, exclusion) on the prepared CheckFunctions. // Concurrency is controlled by the request-scoped semaphore in the context. return func(ctx context.Context) (*invoke.BatchCheckResponse, error) { - return combiner(ctx, functions, engine.concurrencyLimit) + return combiner(ctx, functions, engine.concurrencyLimit, len(request.EntityIDs)) } } @@ -354,7 +357,9 @@ func (engine *CheckEngine) checkDirectRelation(request *invoke.BatchCheckRequest // Build check functions for userset groups, chunking large groups. var checkFunctions []CheckFunction + totalExpectedEntities := 0 for key, ids := range usersetGroups { + totalExpectedEntities += len(ids) for i := 0; i < len(ids); i += engine.maxBatchSize { end := min(i+engine.maxBatchSize, len(ids)) checkFunctions = append(checkFunctions, engine.invoke(&invoke.BatchCheckRequest{ @@ -379,7 +384,7 @@ func (engine *CheckEngine) checkDirectRelation(request *invoke.BatchCheckRequest // If there are userset check functions, run them and map results back to parent entities. if len(checkFunctions) > 0 { - usersetResp, err := checkUnion(ctx, checkFunctions, engine.concurrencyLimit) + usersetResp, err := checkUnion(ctx, checkFunctions, engine.concurrencyLimit, totalExpectedEntities) if err != nil { resp.Metadata = joinResponseMetas(resp.Metadata, usersetResp.Metadata) return resp, err @@ -458,7 +463,9 @@ func (engine *CheckEngine) checkTupleToUserSet( } var checkFunctions []CheckFunction + totalExpectedEntities := 0 for entityType, ids := range subjectsByType { + totalExpectedEntities += len(ids) for i := 0; i < len(ids); i += engine.maxBatchSize { end := min(i+engine.maxBatchSize, len(ids)) checkFunctions = append(checkFunctions, engine.invoke(&invoke.BatchCheckRequest{ @@ -481,7 +488,7 @@ func (engine *CheckEngine) checkTupleToUserSet( return resp, nil } - subjectResp, err := checkUnion(ctx, checkFunctions, engine.concurrencyLimit) + subjectResp, err := checkUnion(ctx, checkFunctions, engine.concurrencyLimit, totalExpectedEntities) if err != nil { resp.Metadata = joinResponseMetas(resp.Metadata, subjectResp.Metadata) return resp, err @@ -728,7 +735,7 @@ func (engine *CheckEngine) checkDirectCall(request *invoke.BatchCheckRequest) Ch // checkUnion checks if the subject has permission by running multiple CheckFunctions concurrently. // Per-entity merge: for each entityID, if ANY function returned ALLOWED -> ALLOWED, else -> DENIED. -func checkUnion(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) { +func checkUnion(ctx context.Context, functions []CheckFunction, limit int, expectedEntityCount int) (*invoke.BatchCheckResponse, error) { // Initialize the response metadata responseMetadata := emptyResponseMetadata() @@ -792,8 +799,8 @@ func checkUnion(ctx context.Context, functions []CheckFunction, limit int) (*inv } entityIDsSeen = true - // Early exit: if all known entities are now ALLOWED, no further functions can change the result. - if entityIDsSeen && deniedCount == 0 { + // Early exit: if all expected entities are now ALLOWED, no further functions can change the result. + if entityIDsSeen && deniedCount == 0 && len(mergedResults) >= expectedEntityCount { return &invoke.BatchCheckResponse{ Results: mergedResults, Metadata: responseMetadata, @@ -816,7 +823,7 @@ func checkUnion(ctx context.Context, functions []CheckFunction, limit int) (*inv // checkIntersection checks if the subject has permission by running multiple CheckFunctions concurrently. // Per-entity merge: for each entityID, ALL functions must return ALLOWED -> ALLOWED, else -> DENIED. -func checkIntersection(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) { +func checkIntersection(ctx context.Context, functions []CheckFunction, limit int, _ int) (*invoke.BatchCheckResponse, error) { // Initialize the response metadata responseMetadata := emptyResponseMetadata() @@ -924,7 +931,7 @@ func checkIntersection(ctx context.Context, functions []CheckFunction, limit int // checkExclusion is a function that checks if there are any exclusions for given CheckFunctions. // Per-entity merge: for each entityID, first function ALLOWED AND all remaining DENIED -> ALLOWED, else -> DENIED. -func checkExclusion(ctx context.Context, functions []CheckFunction, limit int) (*invoke.BatchCheckResponse, error) { +func checkExclusion(ctx context.Context, functions []CheckFunction, limit int, _ int) (*invoke.BatchCheckResponse, error) { // Initialize the response metadata responseMetadata := emptyResponseMetadata() diff --git a/internal/servers/permission_server.go b/internal/servers/permission_server.go index 261cf0b37..9b56d883c 100644 --- a/internal/servers/permission_server.go +++ b/internal/servers/permission_server.go @@ -3,6 +3,7 @@ package servers import ( "context" "errors" + "fmt" "log/slog" "sync" @@ -21,16 +22,21 @@ type PermissionServer struct { invoker invoke.Invoker concurrencyLimit int + bulkLimit int } // NewPermissionServer - Creates new Permission Server -func NewPermissionServer(i invoke.Invoker, concurrencyLimit int) *PermissionServer { +func NewPermissionServer(i invoke.Invoker, concurrencyLimit, bulkLimit int) *PermissionServer { if concurrencyLimit <= 0 { concurrencyLimit = invoke.DefaultConcurrencyLimit } + if bulkLimit <= 0 { + bulkLimit = 100 + } return &PermissionServer{ invoker: i, concurrencyLimit: concurrencyLimit, + bulkLimit: bulkLimit, } } @@ -85,8 +91,8 @@ func (r *PermissionServer) BulkCheck(ctx context.Context, request *v1.Permission return nil, err } - if len(checkItems) > 100 { - err := status.Error(GetStatus(nil), "maximum 100 items allowed") + if len(checkItems) > r.bulkLimit { + err := status.Error(GetStatus(nil), fmt.Sprintf("maximum %d items allowed", r.bulkLimit)) span.RecordError(err) span.SetStatus(otelCodes.Error, err.Error()) return nil, err diff --git a/internal/servers/server.go b/internal/servers/server.go index b1422dadc..34873ee6f 100644 --- a/internal/servers/server.go +++ b/internal/servers/server.go @@ -66,6 +66,8 @@ type Container struct { // ConcurrencyLimit for permission checks ConcurrencyLimit int + // BulkLimit is the maximum number of items in a BulkCheck request + BulkLimit int } // NewContainer is a constructor for the Container struct. @@ -175,7 +177,7 @@ func (s *Container) Run( grpcServer := grpc.NewServer(opts...) // Register various gRPC services to the server. - grpcV1.RegisterPermissionServer(grpcServer, NewPermissionServer(s.Invoker, s.ConcurrencyLimit)) + grpcV1.RegisterPermissionServer(grpcServer, NewPermissionServer(s.Invoker, s.ConcurrencyLimit, s.BulkLimit)) grpcV1.RegisterSchemaServer(grpcServer, NewSchemaServer(s.SW, s.SR)) grpcV1.RegisterDataServer(grpcServer, NewDataServer(s.DR, s.DW, s.BR, s.SR)) grpcV1.RegisterBundleServer(grpcServer, NewBundleServer(s.BR, s.BW)) @@ -188,7 +190,7 @@ func (s *Container) Run( // Create another gRPC server, presumably for invoking permissions. invokeServer := grpc.NewServer(opts...) - grpcV1.RegisterPermissionServer(invokeServer, NewPermissionServer(localInvoker, s.ConcurrencyLimit)) + grpcV1.RegisterPermissionServer(invokeServer, NewPermissionServer(localInvoker, s.ConcurrencyLimit, s.BulkLimit)) // Register health check and reflection services for the invokeServer. health.RegisterHealthServer(invokeServer, NewHealthServer()) // Register health server for invoker diff --git a/internal/servers/server_behavior_test.go b/internal/servers/server_behavior_test.go index 307a08e1b..387be7432 100644 --- a/internal/servers/server_behavior_test.go +++ b/internal/servers/server_behavior_test.go @@ -330,7 +330,7 @@ func TestInterceptorLogger(t *testing.T) { func TestPermissionServerPassesThroughInvoker(t *testing.T) { invoker := &fakePermissionInvoker{} - server := NewPermissionServer(invoker, 0) + server := NewPermissionServer(invoker, 0, 0) if server == nil { t.Fatal("expected permission server") } @@ -384,7 +384,7 @@ func TestPermissionServerPassesThroughInvoker(t *testing.T) { func TestPermissionServerValidationAndInvokerErrors(t *testing.T) { invoker := &fakePermissionInvoker{} - server := NewPermissionServer(invoker, 0) + server := NewPermissionServer(invoker, 0, 0) _, err := server.Check(context.Background(), &v1.PermissionCheckRequest{}) if err == nil { diff --git a/pkg/cmd/serve.go b/pkg/cmd/serve.go index 1f70bf6d1..6c6b676c8 100644 --- a/pkg/cmd/serve.go +++ b/pkg/cmd/serve.go @@ -544,6 +544,7 @@ func serve() func(cmd *cobra.Command, args []string) error { watcher, ) container.ConcurrencyLimit = cfg.Service.Permission.ConcurrencyLimit + container.BulkLimit = cfg.Service.Permission.BulkLimit // Create an error group with the provided context var g *errgroup.Group