diff --git a/internal/datastore/proxy/counting.go b/internal/datastore/proxy/counting.go deleted file mode 100644 index 5fef95807..000000000 --- a/internal/datastore/proxy/counting.go +++ /dev/null @@ -1,270 +0,0 @@ -package proxy - -import ( - "context" - "sync/atomic" - - "github.com/prometheus/client_golang/prometheus" - "github.com/prometheus/client_golang/prometheus/promauto" - - "github.com/authzed/spicedb/pkg/datastore" - "github.com/authzed/spicedb/pkg/datastore/options" - core "github.com/authzed/spicedb/pkg/proto/core/v1" -) - -var ( - countingQueryRelationshipsTotal prometheus.Counter - countingReverseQueryRelationshipsTotal prometheus.Counter - countingReadNamespaceByNameTotal prometheus.Counter - countingListAllNamespacesTotal prometheus.Counter - countingLookupNamespacesWithNamesTotal prometheus.Counter -) - -func init() { - countingQueryRelationshipsTotal = promauto.NewCounter(prometheus.CounterOpts{ - Namespace: "spicedb", - Subsystem: "datastore", - Name: "query_relationships_total", - Help: "total number of QueryRelationships calls", - }) - - countingReverseQueryRelationshipsTotal = promauto.NewCounter(prometheus.CounterOpts{ - Namespace: "spicedb", - Subsystem: "datastore", - Name: "reverse_query_relationships_total", - Help: "total number of ReverseQueryRelationships calls", - }) - - countingReadNamespaceByNameTotal = promauto.NewCounter(prometheus.CounterOpts{ - Namespace: "spicedb", - Subsystem: "datastore", - Name: "read_namespace_by_name_total", - Help: "total number of ReadNamespaceByName calls", - }) - - countingListAllNamespacesTotal = promauto.NewCounter(prometheus.CounterOpts{ - Namespace: "spicedb", - Subsystem: "datastore", - Name: "list_all_namespaces_total", - Help: "total number of ListAllNamespaces calls", - }) - - countingLookupNamespacesWithNamesTotal = promauto.NewCounter(prometheus.CounterOpts{ - Namespace: "spicedb", - Subsystem: "datastore", - Name: "lookup_namespaces_with_names_total", - Help: "total number of LookupNamespacesWithNames calls", - }) -} - -// DatastoreReaderMethodCounts holds thread-safe counters for tracked Reader methods. -// This struct is typically ephemeral (per-request), and represents the delta of -// calls made during its lifetime. Use WriteMethodCounts() to add these counts to -// global Prometheus counters. -type DatastoreReaderMethodCounts struct { - queryRelationships atomic.Uint64 - reverseQueryRelationships atomic.Uint64 - readNamespaceByName atomic.Uint64 - listAllNamespaces atomic.Uint64 - lookupNamespacesWithNames atomic.Uint64 -} - -// QueryRelationships returns the count of QueryRelationships calls. -func (m *DatastoreReaderMethodCounts) QueryRelationships() uint64 { - return m.queryRelationships.Load() -} - -// ReverseQueryRelationships returns the count of ReverseQueryRelationships calls. -func (m *DatastoreReaderMethodCounts) ReverseQueryRelationships() uint64 { - return m.reverseQueryRelationships.Load() -} - -// ReadNamespaceByName returns the count of ReadNamespaceByName calls. -func (m *DatastoreReaderMethodCounts) LegacyReadNamespaceByName() uint64 { - return m.readNamespaceByName.Load() -} - -// ListAllNamespaces returns the count of ListAllNamespaces calls. -func (m *DatastoreReaderMethodCounts) LegacyListAllNamespaces() uint64 { - return m.listAllNamespaces.Load() -} - -// LookupNamespacesWithNames returns the count of LookupNamespacesWithNames calls. -func (m *DatastoreReaderMethodCounts) LegacyLookupNamespacesWithNames() uint64 { - return m.lookupNamespacesWithNames.Load() -} - -// WriteMethodCounts writes the counts to Prometheus metrics as additive counters. -// The DatastoreReaderMethodCounts struct is ephemeral (typically per-request), -// and this function adds its counts to the global Prometheus counters. -// Multiple counts can be combined by calling this function on each. -// Tests can skip calling this and just access the DatastoreReaderMethodCounts directly. -func WriteMethodCounts(counts *DatastoreReaderMethodCounts) { - countingQueryRelationshipsTotal.Add(float64(counts.QueryRelationships())) - countingReverseQueryRelationshipsTotal.Add(float64(counts.ReverseQueryRelationships())) - countingReadNamespaceByNameTotal.Add(float64(counts.LegacyReadNamespaceByName())) - countingListAllNamespacesTotal.Add(float64(counts.LegacyListAllNamespaces())) - countingLookupNamespacesWithNamesTotal.Add(float64(counts.LegacyLookupNamespacesWithNames())) -} - -// NewCountingDatastoreProxy creates a new datastore proxy that counts Reader method calls. -// Returns both the wrapped datastore and a reference to the counts object. -// The counts object can be used in tests to verify datastore usage patterns, -// or passed to WriteMethodCounts() to export counts to Prometheus. -func NewCountingDatastoreProxy(d datastore.ReadOnlyDatastore) (datastore.ReadOnlyDatastore, *DatastoreReaderMethodCounts) { - counts := &DatastoreReaderMethodCounts{} - return &countingProxy{ - delegate: d, - counts: counts, - }, counts -} - -// NewCountingDatastoreProxyForDatastore wraps a full Datastore with counting proxy. -// Like NewCountingDatastoreProxy, but returns a full Datastore interface for use with middleware. -// The ReadWriteTx method is passed through without counting (only Reader methods are counted). -func NewCountingDatastoreProxyForDatastore(d datastore.Datastore) (datastore.Datastore, *DatastoreReaderMethodCounts) { - counts := &DatastoreReaderMethodCounts{} - return &countingDatastoreProxy{ - countingProxy: countingProxy{ - delegate: d, - counts: counts, - }, - fullDatastore: d, - }, counts -} - -type countingDatastoreProxy struct { - countingProxy - fullDatastore datastore.Datastore -} - -// ReadWriteTx delegates to the underlying full datastore -func (p *countingDatastoreProxy) ReadWriteTx(ctx context.Context, fn datastore.TxUserFunc, opts ...options.RWTOptionsOption) (datastore.Revision, error) { - return p.fullDatastore.ReadWriteTx(ctx, fn, opts...) -} - -type countingProxy struct { - delegate datastore.ReadOnlyDatastore - counts *DatastoreReaderMethodCounts -} - -func (p *countingProxy) MetricsID() (string, error) { - return p.delegate.MetricsID() -} - -func (p *countingProxy) UniqueID(ctx context.Context) (string, error) { - return p.delegate.UniqueID(ctx) -} - -func (p *countingProxy) SnapshotReader(rev datastore.Revision) datastore.Reader { - delegateReader := p.delegate.SnapshotReader(rev) - return &countingReader{ - delegate: delegateReader, - counts: p.counts, - } -} - -func (p *countingProxy) OptimizedRevision(ctx context.Context) (datastore.Revision, error) { - return p.delegate.OptimizedRevision(ctx) -} - -func (p *countingProxy) HeadRevision(ctx context.Context) (datastore.Revision, error) { - return p.delegate.HeadRevision(ctx) -} - -func (p *countingProxy) CheckRevision(ctx context.Context, revision datastore.Revision) error { - return p.delegate.CheckRevision(ctx, revision) -} - -func (p *countingProxy) RevisionFromString(serialized string) (datastore.Revision, error) { - return p.delegate.RevisionFromString(serialized) -} - -func (p *countingProxy) Watch(ctx context.Context, afterRevision datastore.Revision, watchOptions datastore.WatchOptions) (<-chan datastore.RevisionChanges, <-chan error) { - return p.delegate.Watch(ctx, afterRevision, watchOptions) -} - -func (p *countingProxy) ReadyState(ctx context.Context) (datastore.ReadyState, error) { - return p.delegate.ReadyState(ctx) -} - -func (p *countingProxy) Features(ctx context.Context) (*datastore.Features, error) { - return p.delegate.Features(ctx) -} - -func (p *countingProxy) OfflineFeatures() (*datastore.Features, error) { - return p.delegate.OfflineFeatures() -} - -func (p *countingProxy) Statistics(ctx context.Context) (datastore.Stats, error) { - return p.delegate.Statistics(ctx) -} - -func (p *countingProxy) Close() error { - return p.delegate.Close() -} - -func (p *countingProxy) Unwrap() datastore.ReadOnlyDatastore { - return p.delegate -} - -type countingReader struct { - delegate datastore.Reader - counts *DatastoreReaderMethodCounts -} - -// Counted methods - increment counter then delegate - -func (r *countingReader) QueryRelationships(ctx context.Context, filter datastore.RelationshipsFilter, opts ...options.QueryOptionsOption) (datastore.RelationshipIterator, error) { - r.counts.queryRelationships.Add(1) - return r.delegate.QueryRelationships(ctx, filter, opts...) -} - -func (r *countingReader) ReverseQueryRelationships(ctx context.Context, subjectsFilter datastore.SubjectsFilter, opts ...options.ReverseQueryOptionsOption) (datastore.RelationshipIterator, error) { - r.counts.reverseQueryRelationships.Add(1) - return r.delegate.ReverseQueryRelationships(ctx, subjectsFilter, opts...) -} - -func (r *countingReader) LegacyReadNamespaceByName(ctx context.Context, nsName string) (*core.NamespaceDefinition, datastore.Revision, error) { - r.counts.readNamespaceByName.Add(1) - return r.delegate.LegacyReadNamespaceByName(ctx, nsName) -} - -func (r *countingReader) LegacyListAllNamespaces(ctx context.Context) ([]datastore.RevisionedNamespace, error) { - r.counts.listAllNamespaces.Add(1) - return r.delegate.LegacyListAllNamespaces(ctx) -} - -func (r *countingReader) LegacyLookupNamespacesWithNames(ctx context.Context, nsNames []string) ([]datastore.RevisionedNamespace, error) { - r.counts.lookupNamespacesWithNames.Add(1) - return r.delegate.LegacyLookupNamespacesWithNames(ctx, nsNames) -} - -// Passthrough methods - not counted - -func (r *countingReader) LegacyReadCaveatByName(ctx context.Context, name string) (*core.CaveatDefinition, datastore.Revision, error) { - return r.delegate.LegacyReadCaveatByName(ctx, name) -} - -func (r *countingReader) LegacyListAllCaveats(ctx context.Context) ([]datastore.RevisionedCaveat, error) { - return r.delegate.LegacyListAllCaveats(ctx) -} - -func (r *countingReader) LegacyLookupCaveatsWithNames(ctx context.Context, caveatNames []string) ([]datastore.RevisionedCaveat, error) { - return r.delegate.LegacyLookupCaveatsWithNames(ctx, caveatNames) -} - -func (r *countingReader) CountRelationships(ctx context.Context, name string) (int, error) { - return r.delegate.CountRelationships(ctx, name) -} - -func (r *countingReader) LookupCounters(ctx context.Context) ([]datastore.RelationshipCounter, error) { - return r.delegate.LookupCounters(ctx) -} - -// Type assertions -var ( - _ datastore.ReadOnlyDatastore = (*countingProxy)(nil) - _ datastore.Datastore = (*countingDatastoreProxy)(nil) - _ datastore.Reader = (*countingReader)(nil) -) diff --git a/internal/datastore/proxy/counting_test.go b/internal/datastore/proxy/counting_test.go deleted file mode 100644 index 4c95fbad8..000000000 --- a/internal/datastore/proxy/counting_test.go +++ /dev/null @@ -1,339 +0,0 @@ -package proxy - -import ( - "context" - "sync" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/mock" - "github.com/stretchr/testify/require" - - "github.com/authzed/spicedb/internal/datastore/proxy/proxy_test" - "github.com/authzed/spicedb/pkg/datastore" - core "github.com/authzed/spicedb/pkg/proto/core/v1" -) - -func TestCountingProxyBasicCounting(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - reader := &proxy_test.MockReader{} - - delegate.On("SnapshotReader", mock.Anything).Return(reader) - reader.On("QueryRelationships", mock.Anything, mock.Anything).Return(nil, nil) - reader.On("ReverseQueryRelationships", mock.Anything, mock.Anything).Return(nil, nil) - reader.On("LegacyReadNamespaceByName", mock.Anything, mock.Anything).Return(&core.NamespaceDefinition{}, datastore.NoRevision, nil) - reader.On("LegacyListAllNamespaces", mock.Anything).Return([]datastore.RevisionedNamespace{}, nil) - reader.On("LegacyLookupNamespacesWithNames", mock.Anything, mock.Anything).Return([]datastore.RevisionedNamespace{}, nil) - - ds, counts := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - - // Verify all counters start at 0 - require.Equal(uint64(0), counts.QueryRelationships()) - require.Equal(uint64(0), counts.ReverseQueryRelationships()) - require.Equal(uint64(0), counts.LegacyReadNamespaceByName()) - require.Equal(uint64(0), counts.LegacyListAllNamespaces()) - require.Equal(uint64(0), counts.LegacyLookupNamespacesWithNames()) - - r := ds.SnapshotReader(datastore.NoRevision) - - // Call each method once - _, err := r.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - require.Equal(uint64(1), counts.QueryRelationships()) - - _, err = r.ReverseQueryRelationships(ctx, datastore.SubjectsFilter{SubjectType: "user"}) - require.NoError(err) - require.Equal(uint64(1), counts.ReverseQueryRelationships()) - - _, _, err = r.LegacyReadNamespaceByName(ctx, "test") - require.NoError(err) - require.Equal(uint64(1), counts.LegacyReadNamespaceByName()) - - _, err = r.LegacyListAllNamespaces(ctx) - require.NoError(err) - require.Equal(uint64(1), counts.LegacyListAllNamespaces()) - - _, err = r.LegacyLookupNamespacesWithNames(ctx, []string{"test"}) - require.NoError(err) - require.Equal(uint64(1), counts.LegacyLookupNamespacesWithNames()) -} - -func TestCountingProxyMultipleCalls(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - reader := &proxy_test.MockReader{} - - delegate.On("SnapshotReader", mock.Anything).Return(reader) - reader.On("QueryRelationships", mock.Anything, mock.Anything).Return(nil, nil) - - ds, counts := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - r := ds.SnapshotReader(datastore.NoRevision) - - require.Equal(uint64(0), counts.QueryRelationships()) - - // Call multiple times - _, err := r.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - require.Equal(uint64(1), counts.QueryRelationships()) - - _, err = r.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - require.Equal(uint64(2), counts.QueryRelationships()) - - _, err = r.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - require.Equal(uint64(3), counts.QueryRelationships()) -} - -func TestCountingProxyCaveatMethodsNotCounted(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - reader := &proxy_test.MockReader{} - - delegate.On("SnapshotReader", mock.Anything).Return(reader) - reader.On("LegacyReadCaveatByName", "test").Return(&core.CaveatDefinition{}, datastore.NoRevision, nil) - reader.On("LegacyListAllCaveats").Return([]datastore.RevisionedCaveat{}, nil) - reader.On("LegacyLookupCaveatsWithNames", mock.Anything).Return([]datastore.RevisionedCaveat{}, nil) - - ds, counts := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - r := ds.SnapshotReader(datastore.NoRevision) - - // Call caveat methods - _, _, err := r.LegacyReadCaveatByName(ctx, "test") - require.NoError(err) - _, err = r.LegacyListAllCaveats(ctx) - require.NoError(err) - _, err = r.LegacyLookupCaveatsWithNames(ctx, []string{"test"}) - require.NoError(err) - - // Verify no counters changed - require.Equal(uint64(0), counts.QueryRelationships()) - require.Equal(uint64(0), counts.ReverseQueryRelationships()) - require.Equal(uint64(0), counts.LegacyReadNamespaceByName()) - require.Equal(uint64(0), counts.LegacyListAllNamespaces()) - require.Equal(uint64(0), counts.LegacyLookupNamespacesWithNames()) - - delegate.AssertExpectations(t) -} - -func TestCountingProxyCounterMethodsNotCounted(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - reader := &proxy_test.MockReader{} - - delegate.On("SnapshotReader", mock.Anything).Return(reader) - reader.On("CountRelationships", "counter1").Return(42, nil) - reader.On("LookupCounters").Return([]datastore.RelationshipCounter{}, nil) - - ds, counts := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - r := ds.SnapshotReader(datastore.NoRevision) - - // Call counter methods - _, err := r.CountRelationships(ctx, "counter1") - require.NoError(err) - _, err = r.LookupCounters(ctx) - require.NoError(err) - - // Verify no counters changed - require.Equal(uint64(0), counts.QueryRelationships()) - require.Equal(uint64(0), counts.ReverseQueryRelationships()) - require.Equal(uint64(0), counts.LegacyReadNamespaceByName()) - require.Equal(uint64(0), counts.LegacyListAllNamespaces()) - require.Equal(uint64(0), counts.LegacyLookupNamespacesWithNames()) - - delegate.AssertExpectations(t) -} - -func TestCountingProxyThreadSafety(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - reader := &proxy_test.MockReader{} - - delegate.On("SnapshotReader", mock.Anything).Return(reader) - reader.On("QueryRelationships", mock.Anything, mock.Anything).Return(nil, nil) - - ds, counts := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - - numGoroutines := uint64(100) - callsPerGoroutine := uint64(100) - - var wg sync.WaitGroup - for range numGoroutines { - wg.Go(func() { - r := ds.SnapshotReader(datastore.NoRevision) - for range callsPerGoroutine { - _, err := r.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - assert.NoError(t, err) - } - }) - } - - wg.Wait() - - // Should have exactly numGoroutines * callsPerGoroutine calls - expected := numGoroutines * callsPerGoroutine - require.Equal(expected, counts.QueryRelationships()) -} - -func TestCountingProxyUnwrap(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - - ds, _ := NewCountingDatastoreProxy(delegate) - - // Access unwrap directly since we wrap ReadOnlyDatastore, not Datastore - unwrapped := ds.(*countingProxy).Unwrap() - require.Equal(delegate, unwrapped) -} - -func TestCountingProxyPassthrough(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - ds, _ := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - - // Test MetricsID (mock returns hardcoded "mock") - metricsID, err := ds.MetricsID() - require.NoError(err) - require.Equal("mock", metricsID) - - // Test UniqueID (mock returns hardcoded "mockds" if CurrentUniqueID is empty) - uniqueID, err := ds.UniqueID(ctx) - require.NoError(err) - require.Equal("mockds", uniqueID) - - // Test HeadRevision - delegate.On("HeadRevision", mock.Anything).Return(datastore.NoRevision, nil).Once() - _, err = ds.HeadRevision(ctx) - require.NoError(err) - - // Test OptimizedRevision - delegate.On("OptimizedRevision", mock.Anything).Return(datastore.NoRevision, nil).Once() - _, err = ds.OptimizedRevision(ctx) - require.NoError(err) - - // Test CheckRevision - delegate.On("CheckRevision", datastore.NoRevision).Return(nil).Once() - err = ds.CheckRevision(ctx, datastore.NoRevision) - require.NoError(err) - - // Test RevisionFromString - delegate.On("RevisionFromString", "test").Return(datastore.NoRevision, nil).Once() - _, err = ds.RevisionFromString("test") - require.NoError(err) - - // Test ReadyState - delegate.On("ReadyState", mock.Anything).Return(datastore.ReadyState{IsReady: true}, nil).Once() - state, err := ds.ReadyState(ctx) - require.NoError(err) - require.True(state.IsReady) - - // Test Features - features := &datastore.Features{} - delegate.On("Features", mock.Anything).Return(features, nil).Once() - returnedFeatures, err := ds.Features(ctx) - require.NoError(err) - require.Equal(features, returnedFeatures) - - // Test OfflineFeatures - delegate.On("OfflineFeatures").Return(features, nil).Once() - returnedFeatures, err = ds.OfflineFeatures() - require.NoError(err) - require.Equal(features, returnedFeatures) - - // Test Statistics - stats := datastore.Stats{} - delegate.On("Statistics", mock.Anything).Return(stats, nil).Once() - returnedStats, err := ds.Statistics(ctx) - require.NoError(err) - require.Equal(stats, returnedStats) - - // Test Close - delegate.On("Close").Return(nil).Once() - err = ds.Close() - require.NoError(err) - - delegate.AssertExpectations(t) -} - -func TestCountingProxyMultipleReaders(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - reader1 := &proxy_test.MockReader{} - reader2 := &proxy_test.MockReader{} - - // First snapshot reader - delegate.On("SnapshotReader", datastore.NoRevision).Return(reader1).Once() - reader1.On("QueryRelationships", mock.Anything, mock.Anything).Return(nil, nil) - - // Second snapshot reader - delegate.On("SnapshotReader", mock.Anything).Return(reader2) - reader2.On("QueryRelationships", mock.Anything, mock.Anything).Return(nil, nil) - - ds, counts := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - - // Create two readers and make calls on each - r1 := ds.SnapshotReader(datastore.NoRevision) - r2 := ds.SnapshotReader(datastore.NoRevision) - - _, err := r1.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - require.Equal(uint64(1), counts.QueryRelationships()) - - _, err = r2.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - require.Equal(uint64(2), counts.QueryRelationships()) - - _, err = r1.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - require.Equal(uint64(3), counts.QueryRelationships()) -} - -func TestWriteMethodCounts(t *testing.T) { - require := require.New(t) - - delegate := &proxy_test.MockDatastore{} - reader := &proxy_test.MockReader{} - - delegate.On("SnapshotReader", mock.Anything).Return(reader) - reader.On("QueryRelationships", mock.Anything, mock.Anything).Return(nil, nil) - reader.On("LegacyReadNamespaceByName", "test").Return(&core.NamespaceDefinition{}, datastore.NoRevision, nil) - - ds, counts := NewCountingDatastoreProxy(delegate) - ctx := context.Background() - r := ds.SnapshotReader(datastore.NoRevision) - - // Make some calls - _, err := r.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - _, err = r.QueryRelationships(ctx, datastore.RelationshipsFilter{}) - require.NoError(err) - _, _, err = r.LegacyReadNamespaceByName(ctx, "test") - require.NoError(err) - - require.Equal(uint64(2), counts.QueryRelationships()) - require.Equal(uint64(1), counts.LegacyReadNamespaceByName()) - - // WriteMethodCounts should not panic and should be callable - // Note: We can't easily test the actual Prometheus counters without - // setting up a full Prometheus test environment, but we can verify - // the function executes without error - require.NotPanics(func() { - WriteMethodCounts(counts) - }) -}