diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 000000000..e7cac1c8b --- /dev/null +++ b/.gitattributes @@ -0,0 +1,4 @@ +internal/mocks/*.go linguist-generated=true +*.pb.go linguist-generated=true +*.pb.*.go linguist-generated=true +proto/internal/buf.lock linguist-generated=true \ No newline at end of file diff --git a/development/grafana/dashboards/dashboard.json b/development/grafana/dashboards/dashboard.json index d5adbe6ee..9fe868389 100644 --- a/development/grafana/dashboards/dashboard.json +++ b/development/grafana/dashboards/dashboard.json @@ -124,20 +124,15 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "disableTextWrap": false, - "editorMode": "builder", - "expr": "histogram_quantile(0.95, sum by(le, grpc_method, grpc_service) (rate(grpc_server_handling_seconds_bucket{job=\"spicedb\"}[$__rate_interval])))", - "fullMetaSearch": false, - "includeNullMetadata": true, - "instant": true, + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by(le, grpc_method, grpc_service) (rate(grpc_server_handling_seconds_bucket{job=\"spicedb\"}[$__rate_interval])))", "key": "Q-52ee0e33-980f-4a19-9821-39530de9f304-0", "legendFormat": "{{grpc_service}}/{{grpc_method}}", "range": true, - "refId": "A", - "useBackend": false + "refId": "A" } ], - "title": "p95 Latency by Method", + "title": "p99 Latency by Method", "type": "timeseries" }, { @@ -238,7 +233,7 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "description": "5xx errors", + "description": "Response codes", "fieldConfig": { "defaults": { "color": { @@ -318,107 +313,13 @@ "uid": "${DS_PROMETHEUS}" }, "editorMode": "code", - "expr": "sum by(grpc_method, grpc_code) (increase(grpc_server_handled_total{grpc_code!=\"OK\", grpc_code!=\"AlreadyExists\", grpc_code!=\"Canceled\", grpc_code!=\"InvalidArgument\", grpc_code!=\"NotFound\", grpc_code!=\"PermissionDenied\", grpc_code!=\"FailedPrecondition\", grpc_service=~\".*PermissionsService|.*SchemaService|.*ExperimentalService\"}[$__rate_interval]))", + "expr": "sum by(grpc_method, grpc_code) (increase(grpc_server_handled_total{grpc_service=~\".*PermissionsService|.*SchemaService|.*ExperimentalService\"}[$__rate_interval]))", "legendFormat": "{{grpc_method}} ({{grpc_code}})", "range": true, "refId": "A" } ], - "title": "5xx Errors", - "type": "timeseries" - }, - { - "datasource": { - "type": "prometheus", - "uid": "${DS_PROMETHEUS}" - }, - "description": "All other (non-5xx) errors", - "fieldConfig": { - "defaults": { - "color": { - "mode": "palette-classic" - }, - "custom": { - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 0, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "auto", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green", - "value": null - }, - { - "color": "red", - "value": 80 - } - ] - }, - "unit": "reqps" - }, - "overrides": [] - }, - "gridPos": { - "h": 8, - "w": 12, - "x": 12, - "y": 11 - }, - "id": 29, - "options": { - "legend": { - "calcs": [], - "displayMode": "table", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "single", - "sort": "none" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus", - "uid": "${DS_PROMETHEUS}" - }, - "editorMode": "code", - "expr": "sum by(grpc_code,grpc_method) (increase(grpc_server_handled_total{grpc_code!=\"OK\", grpc_code!=\"Canceled\", grpc_code!=\"DeadlineExceeded\", grpc_code!=\"Unavailable\", grpc_code!=\"Internal\", grpc_code!=\"Unknown\", grpc_service=~\".*PermissionsService|.*SchemaService|authzed.api.v1.ExperimentalService\"}[$__rate_interval]))", - "legendFormat": "{{grpc_method}} ({{grpc_code}})", - "range": true, - "refId": "A" - } - ], - "title": "Other Errors", + "title": "Response Codes", "type": "timeseries" }, { @@ -519,12 +420,12 @@ "uid": "${DS_PROMETHEUS}" }, "editorMode": "code", - "expr": "histogram_quantile(0.95, sum by(le, operation) (rate(spicedb_datastore_query_latency_bucket[$__rate_interval])))", + "expr": "histogram_quantile(0.99, sum by(le, operation) (rate(spicedb_datastore_query_latency_bucket[$__rate_interval])))", "range": true, "refId": "A" } ], - "title": "p95 Latency by Datastore Operation", + "title": "p99 Latency by Datastore Operation", "type": "timeseries" }, { @@ -811,7 +712,7 @@ }, "editorMode": "builder", "expr": "go_goroutines{job=\"spicedb\"}", - "legendFormat": "Goroutines (Count)", + "legendFormat": "Goroutines ({{instance}})", "range": true, "refId": "A" }, @@ -823,7 +724,7 @@ "editorMode": "builder", "expr": "irate(go_goroutines{job=\"spicedb\"}[$__rate_interval])", "hide": false, - "legendFormat": "Goroutine (Derivative)", + "legendFormat": "Goroutine Derivative ({{instance}})", "range": true, "refId": "B" } @@ -888,7 +789,24 @@ }, "unit": "s" }, - "overrides": [] + "overrides": [ + { + "matcher": { + "id": "byRegexp", + "options": "count.*" + }, + "properties": [ + { + "id": "custom.axisPlacement", + "value": "right" + }, + { + "id": "unit", + "value": "short" + } + ] + } + ] }, "gridPos": { "h": 8, @@ -915,14 +833,26 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "editorMode": "builder", - "expr": "histogram_quantile(0.95, sum by(le) (rate(go_sched_pauses_total_gc_seconds_bucket[$__rate_interval])))", - "legendFormat": "p95", + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by(le, instance) (rate(go_sched_pauses_total_gc_seconds_bucket[$__rate_interval])))", + "legendFormat": "p99 of duration ({{instance}})", "range": true, "refId": "A" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "go_gc_cycles_total_gc_cycles_total", + "hide": false, + "legendFormat": "count ({{instance}})", + "range": true, + "refId": "B" } ], - "title": "Go GC Pause Duration", + "title": "GC (garbage collection)", "type": "timeseries" }, { @@ -938,7 +868,7 @@ "custom": { "axisCenteredZero": false, "axisColorMode": "text", - "axisLabel": "", + "axisLabel": "bytes", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", @@ -981,7 +911,28 @@ }, "unit": "bytes" }, - "overrides": [] + "overrides": [ + { + "matcher": { + "id": "byRegexp", + "options": "Usage % .*" + }, + "properties": [ + { + "id": "custom.axisPlacement", + "value": "right" + }, + { + "id": "custom.axisLabel", + "value": "usage %" + }, + { + "id": "unit", + "value": "percent" + } + ] + } + ] }, "gridPos": { "h": 8, @@ -995,7 +946,7 @@ "calcs": [], "displayMode": "list", "placement": "bottom", - "showLegend": false + "showLegend": true }, "tooltip": { "mode": "single", @@ -1009,10 +960,32 @@ "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", - "expr": "go_memstats_heap_alloc_bytes{job=\"spicedb\"}", - "legendFormat": "{{label_name}}", + "expr": "go_memstats_heap_inuse_bytes{job=\"spicedb\"}", + "legendFormat": "Heap In Use ({{instance}})", "range": true, "refId": "A" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "builder", + "expr": "go_gc_gomemlimit_bytes{job=\"spicedb\"}", + "legendFormat": "GOMEMLIMIT ({{instance}})", + "range": true, + "refId": "B" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "builder", + "expr": "spicedb_memory_middleware_memory_usage_percent{job=\"spicedb\"}", + "legendFormat": "Usage % ({{instance}})", + "range": true, + "refId": "C" } ], "title": "Heap Memory Usage", @@ -1059,7 +1032,7 @@ "mode": "off" } }, - "displayName": "malloc/sec", + "displayName": "malloc/sec ({{instance}})", "mappings": [], "thresholds": { "mode": "absolute", @@ -1104,7 +1077,7 @@ }, "editorMode": "builder", "expr": "rate(go_memstats_mallocs_total{job=\"spicedb\"}[$__rate_interval])", - "legendFormat": "__auto", + "legendFormat": "{{instance}}", "range": true, "refId": "A" } @@ -1197,7 +1170,7 @@ }, "editorMode": "builder", "expr": "rate(go_sync_mutex_wait_total_seconds_total[$__rate_interval])", - "legendFormat": "__auto", + "legendFormat": "{{instance}}", "range": true, "refId": "A" } @@ -1291,12 +1264,12 @@ }, "editorMode": "builder", "expr": "rate(go_memstats_alloc_bytes_total{job=\"spicedb\"}[$__rate_interval])", - "legendFormat": "__auto", + "legendFormat": "bytes/sec ({{instance}}", "range": true, "refId": "A" } ], - "title": "Heap Allocation Rate (bytes/sec)", + "title": "Heap Allocation Rate", "type": "timeseries" } ], @@ -1329,13 +1302,13 @@ ] }, "time": { - "from": "now-30m", + "from": "now-1h", "to": "now" }, "timepicker": {}, "timezone": "", "title": "SpiceDB Dashboard", "uid": "qvsJJ_6Hk", - "version": 1, + "version": 7, "weekStart": "" } \ No newline at end of file diff --git a/development/prometheus.yaml b/development/prometheus.yaml index 6ba9da857..c34ca4fa5 100644 --- a/development/prometheus.yaml +++ b/development/prometheus.yaml @@ -1,11 +1,14 @@ --- global: - scrape_interval: "15s" - evaluation_interval: "15s" + scrape_interval: "1s" + evaluation_interval: "1s" scrape_configs: - job_name: "spicedb" static_configs: - - targets: ["spicedb:9090"] + - targets: ["spicedb-1:9090"] labels: - service: "spicedb" + service: "spicedb-1" + - targets: ["spicedb-2:9090"] + labels: + service: "spicedb-2" diff --git a/go.mod b/go.mod index c5cf2428d..8ad9871e8 100644 --- a/go.mod +++ b/go.mod @@ -112,6 +112,7 @@ require ( go.opentelemetry.io/otel/trace v1.38.0 go.uber.org/atomic v1.11.0 go.uber.org/goleak v1.3.0 + go.uber.org/mock v0.6.0 golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b golang.org/x/mod v0.28.0 golang.org/x/sync v0.17.0 @@ -138,6 +139,8 @@ tool ( github.com/golangci/golangci-lint/v2/cmd/golangci-lint // support running mage with go run mage.go github.com/magefile/mage/mage + // mocks are generated with go:generate directives. + go.uber.org/mock/mockgen // vulncheck always uses the current directory's go.mod. golang.org/x/vuln/cmd/govulncheck ) @@ -313,6 +316,7 @@ require ( github.com/klauspost/cpuid/v2 v2.2.5 // indirect github.com/kulti/thelper v0.7.1 // indirect github.com/kunwardeep/paralleltest v1.0.14 // indirect + github.com/kylelemons/godebug v1.1.0 // indirect github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect github.com/lasiar/canonicalheader v1.1.2 // indirect diff --git a/go.sum b/go.sum index e9183ce0b..d4a3c129c 100644 --- a/go.sum +++ b/go.sum @@ -2595,6 +2595,8 @@ go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwE go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= +go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= go.uber.org/multierr v1.6.0/go.mod h1:cdWPpRnG4AhwMwsgIHip0KRBQjJy5kYEpYjJxpXp9iU= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= diff --git a/internal/dispatch/dispatch.go b/internal/dispatch/dispatch.go index 4708a07c7..34d0c7786 100644 --- a/internal/dispatch/dispatch.go +++ b/internal/dispatch/dispatch.go @@ -1,3 +1,5 @@ +//go:generate go run go.uber.org/mock/mockgen -source dispatch.go -destination ../mocks/mock_dispatcher.go -package mocks Dispatcher + package dispatch import ( diff --git a/internal/middleware/memoryprotection/memory_protection.go b/internal/middleware/memoryprotection/memory_protection.go new file mode 100644 index 000000000..257ebd5e6 --- /dev/null +++ b/internal/middleware/memoryprotection/memory_protection.go @@ -0,0 +1,167 @@ +package memoryprotection + +import ( + "context" + "runtime/debug" + "strconv" + "strings" + + middleware "github.com/grpc-ecosystem/go-grpc-middleware/v2" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + log "github.com/authzed/spicedb/internal/logging" +) + +// RequestsProcessed tracks requests that were processed by this middleware. +var RequestsProcessed = promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: "spicedb", + Subsystem: "memory_middleware", + Name: "requests_processed_total", + Help: "Total requests processed by the memory protection middleware (flag --memory-protection-enabled)", +}, []string{"endpoint", "accepted"}) + +// Config holds configuration for the memory protection middleware +type Config struct { + // ThresholdPercent is the memory usage threshold for requests, in the range [0,1] + // If zero or negative, this middleware has no effect. + ThresholdPercent float64 +} + +// DefaultConfig returns reasonable default configuration for API requests +func DefaultConfig() Config { + return Config{ + ThresholdPercent: 0.90, + } +} + +// DefaultDispatchConfig returns reasonable default configuration for dispatch requests +func DefaultDispatchConfig() Config { + return Config{ + ThresholdPercent: 0.95, + } +} + +// MemoryLimitProvider gets and sets the limit of memory usage. +// In production, use DefaultMemoryLimitProvider. +// For testing, use HardCodedMemoryLimitProvider. +type MemoryLimitProvider interface { + GetInBytes() int64 + SetInBytes(int64) +} + +var ( + _ MemoryLimitProvider = (*DefaultMemoryLimitProvider)(nil) + _ MemoryLimitProvider = (*HardCodedMemoryLimitProvider)(nil) +) + +type DefaultMemoryLimitProvider struct{} + +func (p *DefaultMemoryLimitProvider) GetInBytes() int64 { + // SetMemoryLimit returns the previously set memory limit. + // A negative input does not adjust the limit, and allows for retrieval of the currently set memory limit + return debug.SetMemoryLimit(-1) +} + +func (p *DefaultMemoryLimitProvider) SetInBytes(limit int64) { + debug.SetMemoryLimit(limit) +} + +type HardCodedMemoryLimitProvider struct { + Hardcodedlimit int64 +} + +func (p *HardCodedMemoryLimitProvider) GetInBytes() int64 { + return p.Hardcodedlimit +} + +func (p *HardCodedMemoryLimitProvider) SetInBytes(limit int64) { + p.Hardcodedlimit = limit +} + +type MemoryProtectionMiddleware struct { + config Config + sampler MemorySampler +} + +// New creates a new memory admission middleware with the given sampler, which is assumed to have been started already. +func New(config Config, sampler MemorySampler, name string) *MemoryProtectionMiddleware { + am := MemoryProtectionMiddleware{ + config: config, + sampler: sampler, + } + + if am.disabled() { + log.Warn().Str("name", name).Msg("memory protection middleware disabled") + return &am + } + + log.Info(). + Str("name", name). + Float64("threshold_percent", config.ThresholdPercent). + Msg("memory protection middleware initialized") + + return &am +} + +// UnaryServerInterceptor returns a unary server interceptor that rejects incoming requests is memory usage is too high +func (am *MemoryProtectionMiddleware) UnaryServerInterceptor() grpc.UnaryServerInterceptor { + return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + if err := am.checkAdmission(info.FullMethod); err != nil { + return nil, err + } + + return handler(ctx, req) + } +} + +// StreamServerInterceptor returns a stream server interceptor that rejects incoming requests is memory usage is too high +func (am *MemoryProtectionMiddleware) StreamServerInterceptor() grpc.StreamServerInterceptor { + return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + if err := am.checkAdmission(info.FullMethod); err != nil { + return err + } + + wrapped := middleware.WrapServerStream(stream) + return handler(srv, wrapped) + } +} + +// checkAdmission returns an error if the request should be denied because memory usage is too high. +func (am *MemoryProtectionMiddleware) checkAdmission(method string) error { + if am.disabled() { + return nil + } + + accept := true + defer func() { + am.recordMetric(method, accept) + }() + + if am.sampler.GetMemoryUsagePercent() >= am.config.ThresholdPercent { + accept = false + return status.Error(codes.ResourceExhausted, "server rejected the request because memory usage is above configured threshold") + } + + return nil +} + +func (am *MemoryProtectionMiddleware) disabled() bool { + return am.config.ThresholdPercent <= 0 +} + +// recordMetric updates the RequestsProcessed metric and returns the endpoint type for the input method. +func (am *MemoryProtectionMiddleware) recordMetric(fullMethod string, accepted bool) string { + endpointType := "api" + if strings.HasPrefix(fullMethod, "/dispatch.v1.DispatchService") { + endpointType = "dispatch" + } + + acceptedStr := strconv.FormatBool(accepted) + + RequestsProcessed.WithLabelValues(endpointType, acceptedStr).Inc() + return endpointType +} diff --git a/internal/middleware/memoryprotection/memory_protection_test.go b/internal/middleware/memoryprotection/memory_protection_test.go new file mode 100644 index 000000000..9b0c7ea41 --- /dev/null +++ b/internal/middleware/memoryprotection/memory_protection_test.go @@ -0,0 +1,257 @@ +package memoryprotection + +import ( + "context" + "errors" + "io" + "math" + "testing" + + "github.com/grpc-ecosystem/go-grpc-middleware/v2/testing/testpb" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/require" + "github.com/stretchr/testify/suite" + "go.uber.org/goleak" + "go.uber.org/mock/gomock" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/authzed/spicedb/internal/middleware/memoryprotection/mocks" +) + +func TestDefaultMemoryLimitProvider(t *testing.T) { + dmlp := &DefaultMemoryLimitProvider{} + + get := dmlp.GetInBytes() + require.Equal(t, int64(math.MaxInt64), get) + + dmlp.SetInBytes(int64(100)) + get = dmlp.GetInBytes() + require.Equal(t, int64(100), get) +} + +func TestNew(t *testing.T) { + t.Cleanup(func() { + goleak.VerifyNone(t) + }) + + tests := []struct { + name string + inputConfig Config + expectDisabled bool + }{ + { + name: "reasonable config", + inputConfig: DefaultConfig(), + expectDisabled: false, + }, + { + name: "disabled via config", + inputConfig: Config{ + ThresholdPercent: 0, + }, + expectDisabled: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + lp := HardCodedMemoryLimitProvider{Hardcodedlimit: 100} + + sampler := NewMemorySamplerOnInterval(100, &lp) + t.Cleanup(sampler.Close) + + am := New(tt.inputConfig, sampler, "name") + require.NotNil(t, am) + require.Equal(t, am.disabled(), tt.expectDisabled) + + err := am.checkAdmission("some_method") + if tt.expectDisabled { + require.Nil(t, err) // if the middleware is off, every request is let through + } + }) + } +} + +func TestDefaultConfig(t *testing.T) { + config := DefaultConfig() + require.Equal(t, 0.90, config.ThresholdPercent) +} + +func TestDefaultDispatchConfig(t *testing.T) { + config := DefaultDispatchConfig() + require.Equal(t, 0.95, config.ThresholdPercent) +} + +func TestMemoryProtectionMiddleware_RecordRejection(t *testing.T) { + t.Cleanup(func() { + goleak.VerifyNone(t) + }) + + // Create a custom registry for testing + registry := prometheus.NewRegistry() + + // Register our metrics with the test registry + testRequestsProcessed := prometheus.NewCounterVec(prometheus.CounterOpts{ + Namespace: "spicedb", + Subsystem: "memory_middleware", + Name: "requests_processed_total", + Help: "Total requests processed by the memory protection middleware (flag --memory-protection-enabled)", + }, []string{"endpoint", "accepted"}) + registry.MustRegister(testRequestsProcessed) + + // Replace the global counter with our test counter for this test + originalCounter := RequestsProcessed + RequestsProcessed = testRequestsProcessed + defer func() { + RequestsProcessed = originalCounter + }() + + lp := HardCodedMemoryLimitProvider{Hardcodedlimit: 100} + + sampler := NewMemorySamplerOnInterval(DefaultSampleIntervalSeconds, &lp) + t.Cleanup(sampler.Close) + + am := New(DefaultConfig(), sampler, "test") + + // Test API endpoint + endpointType := am.recordMetric("/authzed.api.v1.PermissionsService/CheckPermission", true) + require.Equal(t, "api", endpointType) + + // Test dispatch endpoint + endpointType = am.recordMetric("/dispatch.v1.DispatchService/DispatchCheck", false) + require.Equal(t, "dispatch", endpointType) + + gaugeValue := testutil.ToFloat64(testRequestsProcessed.WithLabelValues("api", "true")) + require.Equal(t, float64(1), gaugeValue) + + gaugeValue = testutil.ToFloat64(testRequestsProcessed.WithLabelValues("dispatch", "false")) + require.Equal(t, float64(1), gaugeValue) +} + +type memoryProtectionTestServer struct { + testpb.UnimplementedTestServiceServer +} + +func (s *memoryProtectionTestServer) PingEmpty(_ context.Context, _ *testpb.PingEmptyRequest) (*testpb.PingEmptyResponse, error) { + return &testpb.PingEmptyResponse{}, nil +} + +func (s *memoryProtectionTestServer) PingStream(_ testpb.TestService_PingStreamServer) error { + return nil +} + +// unaryRequestBlockingTestSuite is a test suite for testing unary request blocking +type memoryProtectionMiddlewareTestSuite struct { + *testpb.InterceptorTestSuite + expectBlocked bool +} + +func TestMemoryProtectionMiddleware(t *testing.T) { + tests := []struct { + name string + memoryUsagePercent float64 + memoryLimit int64 + thresholdPercent float64 + expectBlocked bool + }{ + { + name: "below threshold - allowed", + memoryUsagePercent: 0.89, + thresholdPercent: 0.9, + expectBlocked: false, + }, + { + name: "above threshold - blocked", + memoryUsagePercent: 0.91, + thresholdPercent: 0.90, + expectBlocked: true, + }, + { + name: "at threshold - blocked", + memoryUsagePercent: 0.90, + thresholdPercent: 0.90, + expectBlocked: true, + }, + { + name: "disabled - allows all", + memoryUsagePercent: 0.90, + thresholdPercent: 0, + expectBlocked: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + config := Config{ + ThresholdPercent: tt.thresholdPercent, + } + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mocksampler := mocks.NewMockMemorySampler(ctrl) + mocksampler.EXPECT().GetMemoryUsagePercent().Return(tt.memoryUsagePercent).AnyTimes() + + am := New(config, mocksampler, "name") + + testSrv := &memoryProtectionTestServer{} + s := &memoryProtectionMiddlewareTestSuite{ + InterceptorTestSuite: &testpb.InterceptorTestSuite{ + TestService: testSrv, + ServerOpts: []grpc.ServerOption{ + grpc.ChainUnaryInterceptor( + am.UnaryServerInterceptor(), + ), + grpc.ChainStreamInterceptor( + am.StreamServerInterceptor(), + ), + }, + }, + expectBlocked: tt.expectBlocked, + } + suite.Run(t, s) + }) + } +} + +func (s *memoryProtectionMiddlewareTestSuite) TestUnaryInterceptor_EnforcesMemoryProtection() { + resp, err := s.Client.PingEmpty(context.Background(), &testpb.PingEmptyRequest{}) + + if s.expectBlocked { + // Request should be blocked + s.Require().Error(err, "Request should be blocked due to memory pressure") + s.Require().Nil(resp, "Response should be nil when blocked") + grpcErr, ok := status.FromError(err) + s.Require().True(ok, "Error should be a gRPC status error") + s.Require().Equal(codes.ResourceExhausted, grpcErr.Code(), "Should return ResourceExhausted error") + } else { + // Request should be allowed + s.Require().NoError(err, "Request should be allowed") + s.Require().NotNil(resp, "Response should not be nil when allowed") + } +} + +func (s *memoryProtectionMiddlewareTestSuite) TestStreamingInterceptor_EnforcesMemoryProtection() { + resp, err := s.Client.PingStream(context.Background()) + require.NoError(s.T(), err, "Request should be allowed") + + res, err := resp.Recv() + if errors.Is(err, io.EOF) { + return + } + if s.expectBlocked { + // Request should be blocked + s.Require().Error(err, "Request should be blocked due to memory pressure") + s.Require().Nil(res, "Response should be nil when blocked") + grpcErr, ok := status.FromError(err) + s.Require().True(ok, "Error should be a gRPC status error") + s.Require().Equal(codes.ResourceExhausted, grpcErr.Code(), "Should return ResourceExhausted error") + } else { + // Request should be allowed + s.Require().NoError(err, "Request should be allowed") + s.Require().NotNil(res, "Response should not be nil when allowed") + } +} diff --git a/internal/middleware/memoryprotection/memory_sampler.go b/internal/middleware/memoryprotection/memory_sampler.go new file mode 100644 index 000000000..cb62a1072 --- /dev/null +++ b/internal/middleware/memoryprotection/memory_sampler.go @@ -0,0 +1,168 @@ +//go:generate go run go.uber.org/mock/mockgen -source memory_sampler.go -destination ./mocks/mock_memory_sampler.go -package mocks MemorySampler +package memoryprotection + +import ( + "runtime/metrics" + "sync" + "time" + + "github.com/dustin/go-humanize" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + + log "github.com/authzed/spicedb/internal/logging" +) + +const DefaultSampleIntervalSeconds = 1 + +var MemoryUsageGauge = promauto.NewGauge(prometheus.GaugeOpts{ + Namespace: "spicedb", + Subsystem: "memory_middleware", + Name: "memory_usage_percent", + Help: "Current memory usage as percentage of GOMEMLIMIT (1-100)", +}) + +// MemorySampler provides memory usage sampling capabilities. +type MemorySampler interface { + // GetMemoryUsagePercent returns a number in the range [0,1] + GetMemoryUsagePercent() float64 + // GetTimestampLastMemorySample returns the timestamp of the most recent memory sample + GetTimestampLastMemorySample() time.Time + // GetIntervalSeconds returns how often the sampler runs + GetIntervalSeconds() int + // Close clears all resources. + Close() +} + +type NoOpSampler struct{} + +func (n NoOpSampler) GetIntervalSeconds() int { + return 0 +} + +func (n NoOpSampler) GetMemoryUsagePercent() float64 { + return 0 +} + +func (n NoOpSampler) GetTimestampLastMemorySample() time.Time { + return time.Now() +} + +func (n NoOpSampler) Close() { +} + +var ( + _ MemorySampler = (*NoOpSampler)(nil) + _ MemorySampler = (*MemorySamplerOnInterval)(nil) +) + +// MemorySamplerOnInterval provides shared memory sampling for multiple middleware instances. +// It runs a single background goroutine to sample memory usage periodically, avoiding +// duplicate sampling when multiple middleware instances are created. +type MemorySamplerOnInterval struct { + sampleIntervalSeconds int + memoryLimitInBytes uint64 + metricsSamples []metrics.Sample + stop chan struct{} + + usageLock *sync.RWMutex + lastMemorySampleInBytes uint64 // GUARDED_BY(usageLock) + timestampLastMemorySample time.Time // GUARDED_BY(usageLock) +} + +func (s *MemorySamplerOnInterval) GetIntervalSeconds() int { + return s.sampleIntervalSeconds +} + +// NewMemorySamplerOnInterval sets the current usage immediately and starts a goroutine to fetch memory usage on an interval. +// The goroutine stops when Close is called. +func NewMemorySamplerOnInterval(sampleIntervalSeconds int, limitProvider MemoryLimitProvider) *MemorySamplerOnInterval { + if sampleIntervalSeconds <= 0 { + log.Warn().Msgf("memory sampler sample interval cannot be zero or negative; using default value of %d seconds", DefaultSampleIntervalSeconds) + sampleIntervalSeconds = DefaultSampleIntervalSeconds + } + + sampler := &MemorySamplerOnInterval{ + stop: make(chan struct{}), + sampleIntervalSeconds: sampleIntervalSeconds, + memoryLimitInBytes: uint64(limitProvider.GetInBytes()), //nolint:gosec + usageLock: &sync.RWMutex{}, + timestampLastMemorySample: time.Now(), + lastMemorySampleInBytes: 0, + metricsSamples: []metrics.Sample{ + {Name: "/memory/classes/heap/objects:bytes"}, + }, + } + + // Initialize with current memory usage + sampler.sampleMemory() + + sampler.startBackgroundSampling() + + log.Info(). + Int("sample_interval_seconds", sampleIntervalSeconds). + Str("limit", humanize.Bytes(sampler.memoryLimitInBytes)). + Msg("shared memory sampler initialized with background sampling") + + return sampler +} + +// GetMemoryUsagePercent returns a number in the range [0,1] +func (s *MemorySamplerOnInterval) GetMemoryUsagePercent() float64 { + s.usageLock.RLock() + defer s.usageLock.RUnlock() + if s.memoryLimitInBytes <= 0 { + return 0 + } + return min(float64(1), float64(s.lastMemorySampleInBytes)/float64(s.memoryLimitInBytes)) +} + +// GetTimestampLastMemorySample returns the timestamp of the most recent memory sample +func (s *MemorySamplerOnInterval) GetTimestampLastMemorySample() time.Time { + s.usageLock.RLock() + defer s.usageLock.RUnlock() + return s.timestampLastMemorySample +} + +// startBackgroundSampling starts a background goroutine that samples memory usage on an interval +func (s *MemorySamplerOnInterval) startBackgroundSampling() { + interval := time.Duration(s.sampleIntervalSeconds) * time.Second + ticker := time.NewTicker(interval) + + go func() { + defer ticker.Stop() + // NOTE: this code might start running before the logger is setup, therefore we cannot log anything here + // or we will trigger a data race + + for { + select { + case <-ticker.C: + s.sampleMemory() + case <-s.stop: + return + } + } + }() +} + +// sampleMemory samples the current memory usage and timestamp and updates the gauge +func (s *MemorySamplerOnInterval) sampleMemory() { + metrics.Read(s.metricsSamples) + if len(s.metricsSamples) == 0 { + log.Warn().Msg("could not get current memory usage, no metrics available") + return + } + newUsage := s.metricsSamples[0].Value.Uint64() + + s.usageLock.Lock() + s.lastMemorySampleInBytes = newUsage + s.timestampLastMemorySample = time.Now() + s.usageLock.Unlock() + + MemoryUsageGauge.Set(s.GetMemoryUsagePercent() * 100) +} + +func (s *MemorySamplerOnInterval) Close() { + s.stop <- struct{}{} + log.Info().Msg("stopped memory sampler") +} diff --git a/internal/middleware/memoryprotection/memory_sampler_test.go b/internal/middleware/memoryprotection/memory_sampler_test.go new file mode 100644 index 000000000..5d8e26b65 --- /dev/null +++ b/internal/middleware/memoryprotection/memory_sampler_test.go @@ -0,0 +1,68 @@ +package memoryprotection + +import ( + "testing" + "testing/synctest" + "time" + + "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/require" +) + +func TestMemorySamplerOnInterval(t *testing.T) { + intervalSeconds := 1000 + interval := time.Duration(intervalSeconds) * time.Second + + testcases := map[string]struct { + limit int64 + expectGaugeAndUsageUpdate bool + }{ + `positive_limit`: { + limit: 100 * 1024 * 1024, + expectGaugeAndUsageUpdate: true, + }, + `very_low_limit`: { + limit: 1, + expectGaugeAndUsageUpdate: true, + }, + `negative_limit`: { + limit: -1, + expectGaugeAndUsageUpdate: false, + }, + `zero`: { + limit: 0, + expectGaugeAndUsageUpdate: false, + }, + } + + for name, tc := range testcases { + t.Run(name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + sampler := NewMemorySamplerOnInterval(intervalSeconds, &HardCodedMemoryLimitProvider{Hardcodedlimit: tc.limit}) + t.Cleanup(sampler.Close) + + now := time.Now() + time.Sleep(interval - 1*time.Millisecond) + synctest.Wait() + t.Log("When we get here, the sampling is guaranteed to NOT have run yet") + + require.False(t, sampler.GetTimestampLastMemorySample().After(now)) + + time.Sleep(1 * time.Millisecond) + synctest.Wait() + t.Log("When we get here, the sampling is guaranteed to have run") + + require.True(t, sampler.GetTimestampLastMemorySample().After(now)) + + gaugeValue := testutil.ToFloat64(MemoryUsageGauge) + if tc.expectGaugeAndUsageUpdate { + require.Greater(t, gaugeValue, float64(0)) + require.Greater(t, sampler.GetMemoryUsagePercent(), float64(0)) + require.LessOrEqual(t, sampler.GetMemoryUsagePercent(), float64(1), "percentage should be between 0 and 1") + } else { + require.InDelta(t, 0, gaugeValue, 0.001) // near zero + } + }) + }) + } +} diff --git a/internal/middleware/memoryprotection/mocks/mock_memory_sampler.go b/internal/middleware/memoryprotection/mocks/mock_memory_sampler.go new file mode 100644 index 000000000..8c9c427b3 --- /dev/null +++ b/internal/middleware/memoryprotection/mocks/mock_memory_sampler.go @@ -0,0 +1,95 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: memory_sampler.go +// +// Generated by this command: +// +// mockgen -source memory_sampler.go -destination ./mocks/mock_memory_sampler.go -package mocks MemorySampler +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + reflect "reflect" + time "time" + + gomock "go.uber.org/mock/gomock" +) + +// MockMemorySampler is a mock of MemorySampler interface. +type MockMemorySampler struct { + ctrl *gomock.Controller + recorder *MockMemorySamplerMockRecorder + isgomock struct{} +} + +// MockMemorySamplerMockRecorder is the mock recorder for MockMemorySampler. +type MockMemorySamplerMockRecorder struct { + mock *MockMemorySampler +} + +// NewMockMemorySampler creates a new mock instance. +func NewMockMemorySampler(ctrl *gomock.Controller) *MockMemorySampler { + mock := &MockMemorySampler{ctrl: ctrl} + mock.recorder = &MockMemorySamplerMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockMemorySampler) EXPECT() *MockMemorySamplerMockRecorder { + return m.recorder +} + +// Close mocks base method. +func (m *MockMemorySampler) Close() { + m.ctrl.T.Helper() + m.ctrl.Call(m, "Close") +} + +// Close indicates an expected call of Close. +func (mr *MockMemorySamplerMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockMemorySampler)(nil).Close)) +} + +// GetIntervalSeconds mocks base method. +func (m *MockMemorySampler) GetIntervalSeconds() int { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetIntervalSeconds") + ret0, _ := ret[0].(int) + return ret0 +} + +// GetIntervalSeconds indicates an expected call of GetIntervalSeconds. +func (mr *MockMemorySamplerMockRecorder) GetIntervalSeconds() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetIntervalSeconds", reflect.TypeOf((*MockMemorySampler)(nil).GetIntervalSeconds)) +} + +// GetMemoryUsagePercent mocks base method. +func (m *MockMemorySampler) GetMemoryUsagePercent() float64 { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetMemoryUsagePercent") + ret0, _ := ret[0].(float64) + return ret0 +} + +// GetMemoryUsagePercent indicates an expected call of GetMemoryUsagePercent. +func (mr *MockMemorySamplerMockRecorder) GetMemoryUsagePercent() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMemoryUsagePercent", reflect.TypeOf((*MockMemorySampler)(nil).GetMemoryUsagePercent)) +} + +// GetTimestampLastMemorySample mocks base method. +func (m *MockMemorySampler) GetTimestampLastMemorySample() time.Time { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetTimestampLastMemorySample") + ret0, _ := ret[0].(time.Time) + return ret0 +} + +// GetTimestampLastMemorySample indicates an expected call of GetTimestampLastMemorySample. +func (mr *MockMemorySamplerMockRecorder) GetTimestampLastMemorySample() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTimestampLastMemorySample", reflect.TypeOf((*MockMemorySampler)(nil).GetTimestampLastMemorySample)) +} diff --git a/internal/mocks/mock_datastore.go b/internal/mocks/mock_datastore.go new file mode 100644 index 000000000..3b271a8c7 --- /dev/null +++ b/internal/mocks/mock_datastore.go @@ -0,0 +1,2341 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: datastore.go +// +// Generated by this command: +// +// mockgen -source datastore.go -destination ../../internal/mocks/mock_datastore.go -package mocks Datastore +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + reflect "reflect" + + v1 "github.com/authzed/authzed-go/proto/authzed/api/v1" + datastore "github.com/authzed/spicedb/pkg/datastore" + options "github.com/authzed/spicedb/pkg/datastore/options" + corev1 "github.com/authzed/spicedb/pkg/proto/core/v1" + tuple "github.com/authzed/spicedb/pkg/tuple" + gomock "go.uber.org/mock/gomock" +) + +// MockSchemaDefinition is a mock of SchemaDefinition interface. +type MockSchemaDefinition struct { + ctrl *gomock.Controller + recorder *MockSchemaDefinitionMockRecorder + isgomock struct{} +} + +// MockSchemaDefinitionMockRecorder is the mock recorder for MockSchemaDefinition. +type MockSchemaDefinitionMockRecorder struct { + mock *MockSchemaDefinition +} + +// NewMockSchemaDefinition creates a new mock instance. +func NewMockSchemaDefinition(ctrl *gomock.Controller) *MockSchemaDefinition { + mock := &MockSchemaDefinition{ctrl: ctrl} + mock.recorder = &MockSchemaDefinitionMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSchemaDefinition) EXPECT() *MockSchemaDefinitionMockRecorder { + return m.recorder +} + +// GetName mocks base method. +func (m *MockSchemaDefinition) GetName() string { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetName") + ret0, _ := ret[0].(string) + return ret0 +} + +// GetName indicates an expected call of GetName. +func (mr *MockSchemaDefinitionMockRecorder) GetName() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetName", reflect.TypeOf((*MockSchemaDefinition)(nil).GetName)) +} + +// SizeVT mocks base method. +func (m *MockSchemaDefinition) SizeVT() int { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SizeVT") + ret0, _ := ret[0].(int) + return ret0 +} + +// SizeVT indicates an expected call of SizeVT. +func (mr *MockSchemaDefinitionMockRecorder) SizeVT() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SizeVT", reflect.TypeOf((*MockSchemaDefinition)(nil).SizeVT)) +} + +// MockReader is a mock of Reader interface. +type MockReader struct { + ctrl *gomock.Controller + recorder *MockReaderMockRecorder + isgomock struct{} +} + +// MockReaderMockRecorder is the mock recorder for MockReader. +type MockReaderMockRecorder struct { + mock *MockReader +} + +// NewMockReader creates a new mock instance. +func NewMockReader(ctrl *gomock.Controller) *MockReader { + mock := &MockReader{ctrl: ctrl} + mock.recorder = &MockReaderMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockReader) EXPECT() *MockReaderMockRecorder { + return m.recorder +} + +// CountRelationships mocks base method. +func (m *MockReader) CountRelationships(ctx context.Context, name string) (int, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CountRelationships", ctx, name) + ret0, _ := ret[0].(int) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// CountRelationships indicates an expected call of CountRelationships. +func (mr *MockReaderMockRecorder) CountRelationships(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountRelationships", reflect.TypeOf((*MockReader)(nil).CountRelationships), ctx, name) +} + +// ListAllCaveats mocks base method. +func (m *MockReader) ListAllCaveats(ctx context.Context) ([]datastore.RevisionedCaveat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAllCaveats", ctx) + ret0, _ := ret[0].([]datastore.RevisionedCaveat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAllCaveats indicates an expected call of ListAllCaveats. +func (mr *MockReaderMockRecorder) ListAllCaveats(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAllCaveats", reflect.TypeOf((*MockReader)(nil).ListAllCaveats), ctx) +} + +// ListAllNamespaces mocks base method. +func (m *MockReader) ListAllNamespaces(ctx context.Context) ([]datastore.RevisionedNamespace, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAllNamespaces", ctx) + ret0, _ := ret[0].([]datastore.RevisionedNamespace) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAllNamespaces indicates an expected call of ListAllNamespaces. +func (mr *MockReaderMockRecorder) ListAllNamespaces(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAllNamespaces", reflect.TypeOf((*MockReader)(nil).ListAllNamespaces), ctx) +} + +// LookupCaveatsWithNames mocks base method. +func (m *MockReader) LookupCaveatsWithNames(ctx context.Context, names []string) ([]datastore.RevisionedCaveat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupCaveatsWithNames", ctx, names) + ret0, _ := ret[0].([]datastore.RevisionedCaveat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LookupCaveatsWithNames indicates an expected call of LookupCaveatsWithNames. +func (mr *MockReaderMockRecorder) LookupCaveatsWithNames(ctx, names any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupCaveatsWithNames", reflect.TypeOf((*MockReader)(nil).LookupCaveatsWithNames), ctx, names) +} + +// LookupCounters mocks base method. +func (m *MockReader) LookupCounters(ctx context.Context) ([]datastore.RelationshipCounter, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupCounters", ctx) + ret0, _ := ret[0].([]datastore.RelationshipCounter) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LookupCounters indicates an expected call of LookupCounters. +func (mr *MockReaderMockRecorder) LookupCounters(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupCounters", reflect.TypeOf((*MockReader)(nil).LookupCounters), ctx) +} + +// LookupNamespacesWithNames mocks base method. +func (m *MockReader) LookupNamespacesWithNames(ctx context.Context, nsNames []string) ([]datastore.RevisionedNamespace, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupNamespacesWithNames", ctx, nsNames) + ret0, _ := ret[0].([]datastore.RevisionedNamespace) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LookupNamespacesWithNames indicates an expected call of LookupNamespacesWithNames. +func (mr *MockReaderMockRecorder) LookupNamespacesWithNames(ctx, nsNames any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupNamespacesWithNames", reflect.TypeOf((*MockReader)(nil).LookupNamespacesWithNames), ctx, nsNames) +} + +// QueryRelationships mocks base method. +func (m *MockReader) QueryRelationships(ctx context.Context, filter datastore.RelationshipsFilter, arg2 ...options.QueryOptionsOption) (datastore.RelationshipIterator, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, filter} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "QueryRelationships", varargs...) + ret0, _ := ret[0].(datastore.RelationshipIterator) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// QueryRelationships indicates an expected call of QueryRelationships. +func (mr *MockReaderMockRecorder) QueryRelationships(ctx, filter any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, filter}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "QueryRelationships", reflect.TypeOf((*MockReader)(nil).QueryRelationships), varargs...) +} + +// ReadCaveatByName mocks base method. +func (m *MockReader) ReadCaveatByName(ctx context.Context, name string) (*corev1.CaveatDefinition, datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadCaveatByName", ctx, name) + ret0, _ := ret[0].(*corev1.CaveatDefinition) + ret1, _ := ret[1].(datastore.Revision) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// ReadCaveatByName indicates an expected call of ReadCaveatByName. +func (mr *MockReaderMockRecorder) ReadCaveatByName(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadCaveatByName", reflect.TypeOf((*MockReader)(nil).ReadCaveatByName), ctx, name) +} + +// ReadNamespaceByName mocks base method. +func (m *MockReader) ReadNamespaceByName(ctx context.Context, nsName string) (*corev1.NamespaceDefinition, datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadNamespaceByName", ctx, nsName) + ret0, _ := ret[0].(*corev1.NamespaceDefinition) + ret1, _ := ret[1].(datastore.Revision) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// ReadNamespaceByName indicates an expected call of ReadNamespaceByName. +func (mr *MockReaderMockRecorder) ReadNamespaceByName(ctx, nsName any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadNamespaceByName", reflect.TypeOf((*MockReader)(nil).ReadNamespaceByName), ctx, nsName) +} + +// ReverseQueryRelationships mocks base method. +func (m *MockReader) ReverseQueryRelationships(ctx context.Context, subjectsFilter datastore.SubjectsFilter, arg2 ...options.ReverseQueryOptionsOption) (datastore.RelationshipIterator, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, subjectsFilter} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "ReverseQueryRelationships", varargs...) + ret0, _ := ret[0].(datastore.RelationshipIterator) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReverseQueryRelationships indicates an expected call of ReverseQueryRelationships. +func (mr *MockReaderMockRecorder) ReverseQueryRelationships(ctx, subjectsFilter any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, subjectsFilter}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReverseQueryRelationships", reflect.TypeOf((*MockReader)(nil).ReverseQueryRelationships), varargs...) +} + +// MockReadWriteTransaction is a mock of ReadWriteTransaction interface. +type MockReadWriteTransaction struct { + ctrl *gomock.Controller + recorder *MockReadWriteTransactionMockRecorder + isgomock struct{} +} + +// MockReadWriteTransactionMockRecorder is the mock recorder for MockReadWriteTransaction. +type MockReadWriteTransactionMockRecorder struct { + mock *MockReadWriteTransaction +} + +// NewMockReadWriteTransaction creates a new mock instance. +func NewMockReadWriteTransaction(ctrl *gomock.Controller) *MockReadWriteTransaction { + mock := &MockReadWriteTransaction{ctrl: ctrl} + mock.recorder = &MockReadWriteTransactionMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockReadWriteTransaction) EXPECT() *MockReadWriteTransactionMockRecorder { + return m.recorder +} + +// BulkLoad mocks base method. +func (m *MockReadWriteTransaction) BulkLoad(ctx context.Context, iter datastore.BulkWriteRelationshipSource) (uint64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "BulkLoad", ctx, iter) + ret0, _ := ret[0].(uint64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// BulkLoad indicates an expected call of BulkLoad. +func (mr *MockReadWriteTransactionMockRecorder) BulkLoad(ctx, iter any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BulkLoad", reflect.TypeOf((*MockReadWriteTransaction)(nil).BulkLoad), ctx, iter) +} + +// CountRelationships mocks base method. +func (m *MockReadWriteTransaction) CountRelationships(ctx context.Context, name string) (int, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CountRelationships", ctx, name) + ret0, _ := ret[0].(int) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// CountRelationships indicates an expected call of CountRelationships. +func (mr *MockReadWriteTransactionMockRecorder) CountRelationships(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountRelationships", reflect.TypeOf((*MockReadWriteTransaction)(nil).CountRelationships), ctx, name) +} + +// DeleteCaveats mocks base method. +func (m *MockReadWriteTransaction) DeleteCaveats(ctx context.Context, names []string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteCaveats", ctx, names) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteCaveats indicates an expected call of DeleteCaveats. +func (mr *MockReadWriteTransactionMockRecorder) DeleteCaveats(ctx, names any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteCaveats", reflect.TypeOf((*MockReadWriteTransaction)(nil).DeleteCaveats), ctx, names) +} + +// DeleteNamespaces mocks base method. +func (m *MockReadWriteTransaction) DeleteNamespaces(ctx context.Context, nsNames ...string) error { + m.ctrl.T.Helper() + varargs := []any{ctx} + for _, a := range nsNames { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "DeleteNamespaces", varargs...) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteNamespaces indicates an expected call of DeleteNamespaces. +func (mr *MockReadWriteTransactionMockRecorder) DeleteNamespaces(ctx any, nsNames ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx}, nsNames...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteNamespaces", reflect.TypeOf((*MockReadWriteTransaction)(nil).DeleteNamespaces), varargs...) +} + +// DeleteRelationships mocks base method. +func (m *MockReadWriteTransaction) DeleteRelationships(ctx context.Context, filter *v1.RelationshipFilter, arg2 ...options.DeleteOptionsOption) (uint64, bool, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, filter} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "DeleteRelationships", varargs...) + ret0, _ := ret[0].(uint64) + ret1, _ := ret[1].(bool) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// DeleteRelationships indicates an expected call of DeleteRelationships. +func (mr *MockReadWriteTransactionMockRecorder) DeleteRelationships(ctx, filter any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, filter}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteRelationships", reflect.TypeOf((*MockReadWriteTransaction)(nil).DeleteRelationships), varargs...) +} + +// ListAllCaveats mocks base method. +func (m *MockReadWriteTransaction) ListAllCaveats(ctx context.Context) ([]datastore.RevisionedCaveat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAllCaveats", ctx) + ret0, _ := ret[0].([]datastore.RevisionedCaveat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAllCaveats indicates an expected call of ListAllCaveats. +func (mr *MockReadWriteTransactionMockRecorder) ListAllCaveats(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAllCaveats", reflect.TypeOf((*MockReadWriteTransaction)(nil).ListAllCaveats), ctx) +} + +// ListAllNamespaces mocks base method. +func (m *MockReadWriteTransaction) ListAllNamespaces(ctx context.Context) ([]datastore.RevisionedNamespace, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAllNamespaces", ctx) + ret0, _ := ret[0].([]datastore.RevisionedNamespace) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAllNamespaces indicates an expected call of ListAllNamespaces. +func (mr *MockReadWriteTransactionMockRecorder) ListAllNamespaces(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAllNamespaces", reflect.TypeOf((*MockReadWriteTransaction)(nil).ListAllNamespaces), ctx) +} + +// LookupCaveatsWithNames mocks base method. +func (m *MockReadWriteTransaction) LookupCaveatsWithNames(ctx context.Context, names []string) ([]datastore.RevisionedCaveat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupCaveatsWithNames", ctx, names) + ret0, _ := ret[0].([]datastore.RevisionedCaveat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LookupCaveatsWithNames indicates an expected call of LookupCaveatsWithNames. +func (mr *MockReadWriteTransactionMockRecorder) LookupCaveatsWithNames(ctx, names any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupCaveatsWithNames", reflect.TypeOf((*MockReadWriteTransaction)(nil).LookupCaveatsWithNames), ctx, names) +} + +// LookupCounters mocks base method. +func (m *MockReadWriteTransaction) LookupCounters(ctx context.Context) ([]datastore.RelationshipCounter, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupCounters", ctx) + ret0, _ := ret[0].([]datastore.RelationshipCounter) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LookupCounters indicates an expected call of LookupCounters. +func (mr *MockReadWriteTransactionMockRecorder) LookupCounters(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupCounters", reflect.TypeOf((*MockReadWriteTransaction)(nil).LookupCounters), ctx) +} + +// LookupNamespacesWithNames mocks base method. +func (m *MockReadWriteTransaction) LookupNamespacesWithNames(ctx context.Context, nsNames []string) ([]datastore.RevisionedNamespace, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LookupNamespacesWithNames", ctx, nsNames) + ret0, _ := ret[0].([]datastore.RevisionedNamespace) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LookupNamespacesWithNames indicates an expected call of LookupNamespacesWithNames. +func (mr *MockReadWriteTransactionMockRecorder) LookupNamespacesWithNames(ctx, nsNames any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LookupNamespacesWithNames", reflect.TypeOf((*MockReadWriteTransaction)(nil).LookupNamespacesWithNames), ctx, nsNames) +} + +// QueryRelationships mocks base method. +func (m *MockReadWriteTransaction) QueryRelationships(ctx context.Context, filter datastore.RelationshipsFilter, arg2 ...options.QueryOptionsOption) (datastore.RelationshipIterator, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, filter} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "QueryRelationships", varargs...) + ret0, _ := ret[0].(datastore.RelationshipIterator) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// QueryRelationships indicates an expected call of QueryRelationships. +func (mr *MockReadWriteTransactionMockRecorder) QueryRelationships(ctx, filter any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, filter}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "QueryRelationships", reflect.TypeOf((*MockReadWriteTransaction)(nil).QueryRelationships), varargs...) +} + +// ReadCaveatByName mocks base method. +func (m *MockReadWriteTransaction) ReadCaveatByName(ctx context.Context, name string) (*corev1.CaveatDefinition, datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadCaveatByName", ctx, name) + ret0, _ := ret[0].(*corev1.CaveatDefinition) + ret1, _ := ret[1].(datastore.Revision) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// ReadCaveatByName indicates an expected call of ReadCaveatByName. +func (mr *MockReadWriteTransactionMockRecorder) ReadCaveatByName(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadCaveatByName", reflect.TypeOf((*MockReadWriteTransaction)(nil).ReadCaveatByName), ctx, name) +} + +// ReadNamespaceByName mocks base method. +func (m *MockReadWriteTransaction) ReadNamespaceByName(ctx context.Context, nsName string) (*corev1.NamespaceDefinition, datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadNamespaceByName", ctx, nsName) + ret0, _ := ret[0].(*corev1.NamespaceDefinition) + ret1, _ := ret[1].(datastore.Revision) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// ReadNamespaceByName indicates an expected call of ReadNamespaceByName. +func (mr *MockReadWriteTransactionMockRecorder) ReadNamespaceByName(ctx, nsName any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadNamespaceByName", reflect.TypeOf((*MockReadWriteTransaction)(nil).ReadNamespaceByName), ctx, nsName) +} + +// RegisterCounter mocks base method. +func (m *MockReadWriteTransaction) RegisterCounter(ctx context.Context, name string, filter *corev1.RelationshipFilter) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RegisterCounter", ctx, name, filter) + ret0, _ := ret[0].(error) + return ret0 +} + +// RegisterCounter indicates an expected call of RegisterCounter. +func (mr *MockReadWriteTransactionMockRecorder) RegisterCounter(ctx, name, filter any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterCounter", reflect.TypeOf((*MockReadWriteTransaction)(nil).RegisterCounter), ctx, name, filter) +} + +// ReverseQueryRelationships mocks base method. +func (m *MockReadWriteTransaction) ReverseQueryRelationships(ctx context.Context, subjectsFilter datastore.SubjectsFilter, arg2 ...options.ReverseQueryOptionsOption) (datastore.RelationshipIterator, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, subjectsFilter} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "ReverseQueryRelationships", varargs...) + ret0, _ := ret[0].(datastore.RelationshipIterator) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReverseQueryRelationships indicates an expected call of ReverseQueryRelationships. +func (mr *MockReadWriteTransactionMockRecorder) ReverseQueryRelationships(ctx, subjectsFilter any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, subjectsFilter}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReverseQueryRelationships", reflect.TypeOf((*MockReadWriteTransaction)(nil).ReverseQueryRelationships), varargs...) +} + +// StoreCounterValue mocks base method. +func (m *MockReadWriteTransaction) StoreCounterValue(ctx context.Context, name string, value int, computedAtRevision datastore.Revision) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "StoreCounterValue", ctx, name, value, computedAtRevision) + ret0, _ := ret[0].(error) + return ret0 +} + +// StoreCounterValue indicates an expected call of StoreCounterValue. +func (mr *MockReadWriteTransactionMockRecorder) StoreCounterValue(ctx, name, value, computedAtRevision any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StoreCounterValue", reflect.TypeOf((*MockReadWriteTransaction)(nil).StoreCounterValue), ctx, name, value, computedAtRevision) +} + +// UnregisterCounter mocks base method. +func (m *MockReadWriteTransaction) UnregisterCounter(ctx context.Context, name string) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UnregisterCounter", ctx, name) + ret0, _ := ret[0].(error) + return ret0 +} + +// UnregisterCounter indicates an expected call of UnregisterCounter. +func (mr *MockReadWriteTransactionMockRecorder) UnregisterCounter(ctx, name any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UnregisterCounter", reflect.TypeOf((*MockReadWriteTransaction)(nil).UnregisterCounter), ctx, name) +} + +// WriteCaveats mocks base method. +func (m *MockReadWriteTransaction) WriteCaveats(arg0 context.Context, arg1 []*corev1.CaveatDefinition) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WriteCaveats", arg0, arg1) + ret0, _ := ret[0].(error) + return ret0 +} + +// WriteCaveats indicates an expected call of WriteCaveats. +func (mr *MockReadWriteTransactionMockRecorder) WriteCaveats(arg0, arg1 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteCaveats", reflect.TypeOf((*MockReadWriteTransaction)(nil).WriteCaveats), arg0, arg1) +} + +// WriteNamespaces mocks base method. +func (m *MockReadWriteTransaction) WriteNamespaces(ctx context.Context, newConfigs ...*corev1.NamespaceDefinition) error { + m.ctrl.T.Helper() + varargs := []any{ctx} + for _, a := range newConfigs { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "WriteNamespaces", varargs...) + ret0, _ := ret[0].(error) + return ret0 +} + +// WriteNamespaces indicates an expected call of WriteNamespaces. +func (mr *MockReadWriteTransactionMockRecorder) WriteNamespaces(ctx any, newConfigs ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx}, newConfigs...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteNamespaces", reflect.TypeOf((*MockReadWriteTransaction)(nil).WriteNamespaces), varargs...) +} + +// WriteRelationships mocks base method. +func (m *MockReadWriteTransaction) WriteRelationships(ctx context.Context, mutations []tuple.RelationshipUpdate) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "WriteRelationships", ctx, mutations) + ret0, _ := ret[0].(error) + return ret0 +} + +// WriteRelationships indicates an expected call of WriteRelationships. +func (mr *MockReadWriteTransactionMockRecorder) WriteRelationships(ctx, mutations any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "WriteRelationships", reflect.TypeOf((*MockReadWriteTransaction)(nil).WriteRelationships), ctx, mutations) +} + +// MockBulkWriteRelationshipSource is a mock of BulkWriteRelationshipSource interface. +type MockBulkWriteRelationshipSource struct { + ctrl *gomock.Controller + recorder *MockBulkWriteRelationshipSourceMockRecorder + isgomock struct{} +} + +// MockBulkWriteRelationshipSourceMockRecorder is the mock recorder for MockBulkWriteRelationshipSource. +type MockBulkWriteRelationshipSourceMockRecorder struct { + mock *MockBulkWriteRelationshipSource +} + +// NewMockBulkWriteRelationshipSource creates a new mock instance. +func NewMockBulkWriteRelationshipSource(ctrl *gomock.Controller) *MockBulkWriteRelationshipSource { + mock := &MockBulkWriteRelationshipSource{ctrl: ctrl} + mock.recorder = &MockBulkWriteRelationshipSourceMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockBulkWriteRelationshipSource) EXPECT() *MockBulkWriteRelationshipSourceMockRecorder { + return m.recorder +} + +// Next mocks base method. +func (m *MockBulkWriteRelationshipSource) Next(ctx context.Context) (*tuple.Relationship, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Next", ctx) + ret0, _ := ret[0].(*tuple.Relationship) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Next indicates an expected call of Next. +func (mr *MockBulkWriteRelationshipSourceMockRecorder) Next(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Next", reflect.TypeOf((*MockBulkWriteRelationshipSource)(nil).Next), ctx) +} + +// MockReadOnlyDatastore is a mock of ReadOnlyDatastore interface. +type MockReadOnlyDatastore struct { + ctrl *gomock.Controller + recorder *MockReadOnlyDatastoreMockRecorder + isgomock struct{} +} + +// MockReadOnlyDatastoreMockRecorder is the mock recorder for MockReadOnlyDatastore. +type MockReadOnlyDatastoreMockRecorder struct { + mock *MockReadOnlyDatastore +} + +// NewMockReadOnlyDatastore creates a new mock instance. +func NewMockReadOnlyDatastore(ctrl *gomock.Controller) *MockReadOnlyDatastore { + mock := &MockReadOnlyDatastore{ctrl: ctrl} + mock.recorder = &MockReadOnlyDatastoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockReadOnlyDatastore) EXPECT() *MockReadOnlyDatastoreMockRecorder { + return m.recorder +} + +// CheckRevision mocks base method. +func (m *MockReadOnlyDatastore) CheckRevision(ctx context.Context, revision datastore.Revision) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CheckRevision", ctx, revision) + ret0, _ := ret[0].(error) + return ret0 +} + +// CheckRevision indicates an expected call of CheckRevision. +func (mr *MockReadOnlyDatastoreMockRecorder) CheckRevision(ctx, revision any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRevision", reflect.TypeOf((*MockReadOnlyDatastore)(nil).CheckRevision), ctx, revision) +} + +// Close mocks base method. +func (m *MockReadOnlyDatastore) Close() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close") + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockReadOnlyDatastoreMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockReadOnlyDatastore)(nil).Close)) +} + +// Features mocks base method. +func (m *MockReadOnlyDatastore) Features(ctx context.Context) (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Features", ctx) + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Features indicates an expected call of Features. +func (mr *MockReadOnlyDatastoreMockRecorder) Features(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Features", reflect.TypeOf((*MockReadOnlyDatastore)(nil).Features), ctx) +} + +// HeadRevision mocks base method. +func (m *MockReadOnlyDatastore) HeadRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "HeadRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// HeadRevision indicates an expected call of HeadRevision. +func (mr *MockReadOnlyDatastoreMockRecorder) HeadRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeadRevision", reflect.TypeOf((*MockReadOnlyDatastore)(nil).HeadRevision), ctx) +} + +// MetricsID mocks base method. +func (m *MockReadOnlyDatastore) MetricsID() (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MetricsID") + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// MetricsID indicates an expected call of MetricsID. +func (mr *MockReadOnlyDatastoreMockRecorder) MetricsID() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MetricsID", reflect.TypeOf((*MockReadOnlyDatastore)(nil).MetricsID)) +} + +// OfflineFeatures mocks base method. +func (m *MockReadOnlyDatastore) OfflineFeatures() (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OfflineFeatures") + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OfflineFeatures indicates an expected call of OfflineFeatures. +func (mr *MockReadOnlyDatastoreMockRecorder) OfflineFeatures() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OfflineFeatures", reflect.TypeOf((*MockReadOnlyDatastore)(nil).OfflineFeatures)) +} + +// OptimizedRevision mocks base method. +func (m *MockReadOnlyDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OptimizedRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OptimizedRevision indicates an expected call of OptimizedRevision. +func (mr *MockReadOnlyDatastoreMockRecorder) OptimizedRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OptimizedRevision", reflect.TypeOf((*MockReadOnlyDatastore)(nil).OptimizedRevision), ctx) +} + +// ReadyState mocks base method. +func (m *MockReadOnlyDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadyState", ctx) + ret0, _ := ret[0].(datastore.ReadyState) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadyState indicates an expected call of ReadyState. +func (mr *MockReadOnlyDatastoreMockRecorder) ReadyState(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadyState", reflect.TypeOf((*MockReadOnlyDatastore)(nil).ReadyState), ctx) +} + +// RevisionFromString mocks base method. +func (m *MockReadOnlyDatastore) RevisionFromString(serialized string) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RevisionFromString", serialized) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RevisionFromString indicates an expected call of RevisionFromString. +func (mr *MockReadOnlyDatastoreMockRecorder) RevisionFromString(serialized any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RevisionFromString", reflect.TypeOf((*MockReadOnlyDatastore)(nil).RevisionFromString), serialized) +} + +// SnapshotReader mocks base method. +func (m *MockReadOnlyDatastore) SnapshotReader(arg0 datastore.Revision) datastore.Reader { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SnapshotReader", arg0) + ret0, _ := ret[0].(datastore.Reader) + return ret0 +} + +// SnapshotReader indicates an expected call of SnapshotReader. +func (mr *MockReadOnlyDatastoreMockRecorder) SnapshotReader(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SnapshotReader", reflect.TypeOf((*MockReadOnlyDatastore)(nil).SnapshotReader), arg0) +} + +// Statistics mocks base method. +func (m *MockReadOnlyDatastore) Statistics(ctx context.Context) (datastore.Stats, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Statistics", ctx) + ret0, _ := ret[0].(datastore.Stats) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Statistics indicates an expected call of Statistics. +func (mr *MockReadOnlyDatastoreMockRecorder) Statistics(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Statistics", reflect.TypeOf((*MockReadOnlyDatastore)(nil).Statistics), ctx) +} + +// UniqueID mocks base method. +func (m *MockReadOnlyDatastore) UniqueID(arg0 context.Context) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UniqueID", arg0) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UniqueID indicates an expected call of UniqueID. +func (mr *MockReadOnlyDatastoreMockRecorder) UniqueID(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UniqueID", reflect.TypeOf((*MockReadOnlyDatastore)(nil).UniqueID), arg0) +} + +// Watch mocks base method. +func (m *MockReadOnlyDatastore) Watch(ctx context.Context, afterRevision datastore.Revision, arg2 datastore.WatchOptions) (<-chan datastore.RevisionChanges, <-chan error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Watch", ctx, afterRevision, arg2) + ret0, _ := ret[0].(<-chan datastore.RevisionChanges) + ret1, _ := ret[1].(<-chan error) + return ret0, ret1 +} + +// Watch indicates an expected call of Watch. +func (mr *MockReadOnlyDatastoreMockRecorder) Watch(ctx, afterRevision, arg2 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockReadOnlyDatastore)(nil).Watch), ctx, afterRevision, arg2) +} + +// MockDatastore is a mock of Datastore interface. +type MockDatastore struct { + ctrl *gomock.Controller + recorder *MockDatastoreMockRecorder + isgomock struct{} +} + +// MockDatastoreMockRecorder is the mock recorder for MockDatastore. +type MockDatastoreMockRecorder struct { + mock *MockDatastore +} + +// NewMockDatastore creates a new mock instance. +func NewMockDatastore(ctrl *gomock.Controller) *MockDatastore { + mock := &MockDatastore{ctrl: ctrl} + mock.recorder = &MockDatastoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockDatastore) EXPECT() *MockDatastoreMockRecorder { + return m.recorder +} + +// CheckRevision mocks base method. +func (m *MockDatastore) CheckRevision(ctx context.Context, revision datastore.Revision) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CheckRevision", ctx, revision) + ret0, _ := ret[0].(error) + return ret0 +} + +// CheckRevision indicates an expected call of CheckRevision. +func (mr *MockDatastoreMockRecorder) CheckRevision(ctx, revision any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRevision", reflect.TypeOf((*MockDatastore)(nil).CheckRevision), ctx, revision) +} + +// Close mocks base method. +func (m *MockDatastore) Close() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close") + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockDatastoreMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockDatastore)(nil).Close)) +} + +// Features mocks base method. +func (m *MockDatastore) Features(ctx context.Context) (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Features", ctx) + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Features indicates an expected call of Features. +func (mr *MockDatastoreMockRecorder) Features(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Features", reflect.TypeOf((*MockDatastore)(nil).Features), ctx) +} + +// HeadRevision mocks base method. +func (m *MockDatastore) HeadRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "HeadRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// HeadRevision indicates an expected call of HeadRevision. +func (mr *MockDatastoreMockRecorder) HeadRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeadRevision", reflect.TypeOf((*MockDatastore)(nil).HeadRevision), ctx) +} + +// MetricsID mocks base method. +func (m *MockDatastore) MetricsID() (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MetricsID") + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// MetricsID indicates an expected call of MetricsID. +func (mr *MockDatastoreMockRecorder) MetricsID() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MetricsID", reflect.TypeOf((*MockDatastore)(nil).MetricsID)) +} + +// OfflineFeatures mocks base method. +func (m *MockDatastore) OfflineFeatures() (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OfflineFeatures") + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OfflineFeatures indicates an expected call of OfflineFeatures. +func (mr *MockDatastoreMockRecorder) OfflineFeatures() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OfflineFeatures", reflect.TypeOf((*MockDatastore)(nil).OfflineFeatures)) +} + +// OptimizedRevision mocks base method. +func (m *MockDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OptimizedRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OptimizedRevision indicates an expected call of OptimizedRevision. +func (mr *MockDatastoreMockRecorder) OptimizedRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OptimizedRevision", reflect.TypeOf((*MockDatastore)(nil).OptimizedRevision), ctx) +} + +// ReadWriteTx mocks base method. +func (m *MockDatastore) ReadWriteTx(arg0 context.Context, arg1 datastore.TxUserFunc, arg2 ...options.RWTOptionsOption) (datastore.Revision, error) { + m.ctrl.T.Helper() + varargs := []any{arg0, arg1} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "ReadWriteTx", varargs...) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadWriteTx indicates an expected call of ReadWriteTx. +func (mr *MockDatastoreMockRecorder) ReadWriteTx(arg0, arg1 any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{arg0, arg1}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadWriteTx", reflect.TypeOf((*MockDatastore)(nil).ReadWriteTx), varargs...) +} + +// ReadyState mocks base method. +func (m *MockDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadyState", ctx) + ret0, _ := ret[0].(datastore.ReadyState) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadyState indicates an expected call of ReadyState. +func (mr *MockDatastoreMockRecorder) ReadyState(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadyState", reflect.TypeOf((*MockDatastore)(nil).ReadyState), ctx) +} + +// RevisionFromString mocks base method. +func (m *MockDatastore) RevisionFromString(serialized string) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RevisionFromString", serialized) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RevisionFromString indicates an expected call of RevisionFromString. +func (mr *MockDatastoreMockRecorder) RevisionFromString(serialized any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RevisionFromString", reflect.TypeOf((*MockDatastore)(nil).RevisionFromString), serialized) +} + +// SnapshotReader mocks base method. +func (m *MockDatastore) SnapshotReader(arg0 datastore.Revision) datastore.Reader { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SnapshotReader", arg0) + ret0, _ := ret[0].(datastore.Reader) + return ret0 +} + +// SnapshotReader indicates an expected call of SnapshotReader. +func (mr *MockDatastoreMockRecorder) SnapshotReader(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SnapshotReader", reflect.TypeOf((*MockDatastore)(nil).SnapshotReader), arg0) +} + +// Statistics mocks base method. +func (m *MockDatastore) Statistics(ctx context.Context) (datastore.Stats, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Statistics", ctx) + ret0, _ := ret[0].(datastore.Stats) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Statistics indicates an expected call of Statistics. +func (mr *MockDatastoreMockRecorder) Statistics(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Statistics", reflect.TypeOf((*MockDatastore)(nil).Statistics), ctx) +} + +// UniqueID mocks base method. +func (m *MockDatastore) UniqueID(arg0 context.Context) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UniqueID", arg0) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UniqueID indicates an expected call of UniqueID. +func (mr *MockDatastoreMockRecorder) UniqueID(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UniqueID", reflect.TypeOf((*MockDatastore)(nil).UniqueID), arg0) +} + +// Watch mocks base method. +func (m *MockDatastore) Watch(ctx context.Context, afterRevision datastore.Revision, arg2 datastore.WatchOptions) (<-chan datastore.RevisionChanges, <-chan error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Watch", ctx, afterRevision, arg2) + ret0, _ := ret[0].(<-chan datastore.RevisionChanges) + ret1, _ := ret[1].(<-chan error) + return ret0, ret1 +} + +// Watch indicates an expected call of Watch. +func (mr *MockDatastoreMockRecorder) Watch(ctx, afterRevision, arg2 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockDatastore)(nil).Watch), ctx, afterRevision, arg2) +} + +// MockExplainable is a mock of Explainable interface. +type MockExplainable struct { + ctrl *gomock.Controller + recorder *MockExplainableMockRecorder + isgomock struct{} +} + +// MockExplainableMockRecorder is the mock recorder for MockExplainable. +type MockExplainableMockRecorder struct { + mock *MockExplainable +} + +// NewMockExplainable creates a new mock instance. +func NewMockExplainable(ctrl *gomock.Controller) *MockExplainable { + mock := &MockExplainable{ctrl: ctrl} + mock.recorder = &MockExplainableMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockExplainable) EXPECT() *MockExplainableMockRecorder { + return m.recorder +} + +// BuildExplainQuery mocks base method. +func (m *MockExplainable) BuildExplainQuery(sql string, args []any) (string, []any, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "BuildExplainQuery", sql, args) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].([]any) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// BuildExplainQuery indicates an expected call of BuildExplainQuery. +func (mr *MockExplainableMockRecorder) BuildExplainQuery(sql, args any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BuildExplainQuery", reflect.TypeOf((*MockExplainable)(nil).BuildExplainQuery), sql, args) +} + +// ParseExplain mocks base method. +func (m *MockExplainable) ParseExplain(explain string) (datastore.ParsedExplain, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ParseExplain", explain) + ret0, _ := ret[0].(datastore.ParsedExplain) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ParseExplain indicates an expected call of ParseExplain. +func (mr *MockExplainableMockRecorder) ParseExplain(explain any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ParseExplain", reflect.TypeOf((*MockExplainable)(nil).ParseExplain), explain) +} + +// PreExplainStatements mocks base method. +func (m *MockExplainable) PreExplainStatements() []string { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "PreExplainStatements") + ret0, _ := ret[0].([]string) + return ret0 +} + +// PreExplainStatements indicates an expected call of PreExplainStatements. +func (mr *MockExplainableMockRecorder) PreExplainStatements() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PreExplainStatements", reflect.TypeOf((*MockExplainable)(nil).PreExplainStatements)) +} + +// MockSQLDatastore is a mock of SQLDatastore interface. +type MockSQLDatastore struct { + ctrl *gomock.Controller + recorder *MockSQLDatastoreMockRecorder + isgomock struct{} +} + +// MockSQLDatastoreMockRecorder is the mock recorder for MockSQLDatastore. +type MockSQLDatastoreMockRecorder struct { + mock *MockSQLDatastore +} + +// NewMockSQLDatastore creates a new mock instance. +func NewMockSQLDatastore(ctrl *gomock.Controller) *MockSQLDatastore { + mock := &MockSQLDatastore{ctrl: ctrl} + mock.recorder = &MockSQLDatastoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSQLDatastore) EXPECT() *MockSQLDatastoreMockRecorder { + return m.recorder +} + +// BuildExplainQuery mocks base method. +func (m *MockSQLDatastore) BuildExplainQuery(sql string, args []any) (string, []any, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "BuildExplainQuery", sql, args) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].([]any) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// BuildExplainQuery indicates an expected call of BuildExplainQuery. +func (mr *MockSQLDatastoreMockRecorder) BuildExplainQuery(sql, args any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BuildExplainQuery", reflect.TypeOf((*MockSQLDatastore)(nil).BuildExplainQuery), sql, args) +} + +// CheckRevision mocks base method. +func (m *MockSQLDatastore) CheckRevision(ctx context.Context, revision datastore.Revision) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CheckRevision", ctx, revision) + ret0, _ := ret[0].(error) + return ret0 +} + +// CheckRevision indicates an expected call of CheckRevision. +func (mr *MockSQLDatastoreMockRecorder) CheckRevision(ctx, revision any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRevision", reflect.TypeOf((*MockSQLDatastore)(nil).CheckRevision), ctx, revision) +} + +// Close mocks base method. +func (m *MockSQLDatastore) Close() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close") + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockSQLDatastoreMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockSQLDatastore)(nil).Close)) +} + +// Features mocks base method. +func (m *MockSQLDatastore) Features(ctx context.Context) (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Features", ctx) + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Features indicates an expected call of Features. +func (mr *MockSQLDatastoreMockRecorder) Features(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Features", reflect.TypeOf((*MockSQLDatastore)(nil).Features), ctx) +} + +// HeadRevision mocks base method. +func (m *MockSQLDatastore) HeadRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "HeadRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// HeadRevision indicates an expected call of HeadRevision. +func (mr *MockSQLDatastoreMockRecorder) HeadRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeadRevision", reflect.TypeOf((*MockSQLDatastore)(nil).HeadRevision), ctx) +} + +// MetricsID mocks base method. +func (m *MockSQLDatastore) MetricsID() (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MetricsID") + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// MetricsID indicates an expected call of MetricsID. +func (mr *MockSQLDatastoreMockRecorder) MetricsID() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MetricsID", reflect.TypeOf((*MockSQLDatastore)(nil).MetricsID)) +} + +// OfflineFeatures mocks base method. +func (m *MockSQLDatastore) OfflineFeatures() (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OfflineFeatures") + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OfflineFeatures indicates an expected call of OfflineFeatures. +func (mr *MockSQLDatastoreMockRecorder) OfflineFeatures() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OfflineFeatures", reflect.TypeOf((*MockSQLDatastore)(nil).OfflineFeatures)) +} + +// OptimizedRevision mocks base method. +func (m *MockSQLDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OptimizedRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OptimizedRevision indicates an expected call of OptimizedRevision. +func (mr *MockSQLDatastoreMockRecorder) OptimizedRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OptimizedRevision", reflect.TypeOf((*MockSQLDatastore)(nil).OptimizedRevision), ctx) +} + +// ParseExplain mocks base method. +func (m *MockSQLDatastore) ParseExplain(explain string) (datastore.ParsedExplain, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ParseExplain", explain) + ret0, _ := ret[0].(datastore.ParsedExplain) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ParseExplain indicates an expected call of ParseExplain. +func (mr *MockSQLDatastoreMockRecorder) ParseExplain(explain any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ParseExplain", reflect.TypeOf((*MockSQLDatastore)(nil).ParseExplain), explain) +} + +// PreExplainStatements mocks base method. +func (m *MockSQLDatastore) PreExplainStatements() []string { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "PreExplainStatements") + ret0, _ := ret[0].([]string) + return ret0 +} + +// PreExplainStatements indicates an expected call of PreExplainStatements. +func (mr *MockSQLDatastoreMockRecorder) PreExplainStatements() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PreExplainStatements", reflect.TypeOf((*MockSQLDatastore)(nil).PreExplainStatements)) +} + +// ReadWriteTx mocks base method. +func (m *MockSQLDatastore) ReadWriteTx(arg0 context.Context, arg1 datastore.TxUserFunc, arg2 ...options.RWTOptionsOption) (datastore.Revision, error) { + m.ctrl.T.Helper() + varargs := []any{arg0, arg1} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "ReadWriteTx", varargs...) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadWriteTx indicates an expected call of ReadWriteTx. +func (mr *MockSQLDatastoreMockRecorder) ReadWriteTx(arg0, arg1 any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{arg0, arg1}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadWriteTx", reflect.TypeOf((*MockSQLDatastore)(nil).ReadWriteTx), varargs...) +} + +// ReadyState mocks base method. +func (m *MockSQLDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadyState", ctx) + ret0, _ := ret[0].(datastore.ReadyState) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadyState indicates an expected call of ReadyState. +func (mr *MockSQLDatastoreMockRecorder) ReadyState(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadyState", reflect.TypeOf((*MockSQLDatastore)(nil).ReadyState), ctx) +} + +// RevisionFromString mocks base method. +func (m *MockSQLDatastore) RevisionFromString(serialized string) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RevisionFromString", serialized) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RevisionFromString indicates an expected call of RevisionFromString. +func (mr *MockSQLDatastoreMockRecorder) RevisionFromString(serialized any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RevisionFromString", reflect.TypeOf((*MockSQLDatastore)(nil).RevisionFromString), serialized) +} + +// SnapshotReader mocks base method. +func (m *MockSQLDatastore) SnapshotReader(arg0 datastore.Revision) datastore.Reader { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SnapshotReader", arg0) + ret0, _ := ret[0].(datastore.Reader) + return ret0 +} + +// SnapshotReader indicates an expected call of SnapshotReader. +func (mr *MockSQLDatastoreMockRecorder) SnapshotReader(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SnapshotReader", reflect.TypeOf((*MockSQLDatastore)(nil).SnapshotReader), arg0) +} + +// Statistics mocks base method. +func (m *MockSQLDatastore) Statistics(ctx context.Context) (datastore.Stats, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Statistics", ctx) + ret0, _ := ret[0].(datastore.Stats) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Statistics indicates an expected call of Statistics. +func (mr *MockSQLDatastoreMockRecorder) Statistics(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Statistics", reflect.TypeOf((*MockSQLDatastore)(nil).Statistics), ctx) +} + +// UniqueID mocks base method. +func (m *MockSQLDatastore) UniqueID(arg0 context.Context) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UniqueID", arg0) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UniqueID indicates an expected call of UniqueID. +func (mr *MockSQLDatastoreMockRecorder) UniqueID(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UniqueID", reflect.TypeOf((*MockSQLDatastore)(nil).UniqueID), arg0) +} + +// Watch mocks base method. +func (m *MockSQLDatastore) Watch(ctx context.Context, afterRevision datastore.Revision, arg2 datastore.WatchOptions) (<-chan datastore.RevisionChanges, <-chan error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Watch", ctx, afterRevision, arg2) + ret0, _ := ret[0].(<-chan datastore.RevisionChanges) + ret1, _ := ret[1].(<-chan error) + return ret0, ret1 +} + +// Watch indicates an expected call of Watch. +func (mr *MockSQLDatastoreMockRecorder) Watch(ctx, afterRevision, arg2 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockSQLDatastore)(nil).Watch), ctx, afterRevision, arg2) +} + +// MockStrictReadDatastore is a mock of StrictReadDatastore interface. +type MockStrictReadDatastore struct { + ctrl *gomock.Controller + recorder *MockStrictReadDatastoreMockRecorder + isgomock struct{} +} + +// MockStrictReadDatastoreMockRecorder is the mock recorder for MockStrictReadDatastore. +type MockStrictReadDatastoreMockRecorder struct { + mock *MockStrictReadDatastore +} + +// NewMockStrictReadDatastore creates a new mock instance. +func NewMockStrictReadDatastore(ctrl *gomock.Controller) *MockStrictReadDatastore { + mock := &MockStrictReadDatastore{ctrl: ctrl} + mock.recorder = &MockStrictReadDatastoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockStrictReadDatastore) EXPECT() *MockStrictReadDatastoreMockRecorder { + return m.recorder +} + +// CheckRevision mocks base method. +func (m *MockStrictReadDatastore) CheckRevision(ctx context.Context, revision datastore.Revision) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CheckRevision", ctx, revision) + ret0, _ := ret[0].(error) + return ret0 +} + +// CheckRevision indicates an expected call of CheckRevision. +func (mr *MockStrictReadDatastoreMockRecorder) CheckRevision(ctx, revision any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRevision", reflect.TypeOf((*MockStrictReadDatastore)(nil).CheckRevision), ctx, revision) +} + +// Close mocks base method. +func (m *MockStrictReadDatastore) Close() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close") + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockStrictReadDatastoreMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockStrictReadDatastore)(nil).Close)) +} + +// Features mocks base method. +func (m *MockStrictReadDatastore) Features(ctx context.Context) (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Features", ctx) + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Features indicates an expected call of Features. +func (mr *MockStrictReadDatastoreMockRecorder) Features(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Features", reflect.TypeOf((*MockStrictReadDatastore)(nil).Features), ctx) +} + +// HeadRevision mocks base method. +func (m *MockStrictReadDatastore) HeadRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "HeadRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// HeadRevision indicates an expected call of HeadRevision. +func (mr *MockStrictReadDatastoreMockRecorder) HeadRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeadRevision", reflect.TypeOf((*MockStrictReadDatastore)(nil).HeadRevision), ctx) +} + +// IsStrictReadModeEnabled mocks base method. +func (m *MockStrictReadDatastore) IsStrictReadModeEnabled() bool { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IsStrictReadModeEnabled") + ret0, _ := ret[0].(bool) + return ret0 +} + +// IsStrictReadModeEnabled indicates an expected call of IsStrictReadModeEnabled. +func (mr *MockStrictReadDatastoreMockRecorder) IsStrictReadModeEnabled() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsStrictReadModeEnabled", reflect.TypeOf((*MockStrictReadDatastore)(nil).IsStrictReadModeEnabled)) +} + +// MetricsID mocks base method. +func (m *MockStrictReadDatastore) MetricsID() (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MetricsID") + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// MetricsID indicates an expected call of MetricsID. +func (mr *MockStrictReadDatastoreMockRecorder) MetricsID() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MetricsID", reflect.TypeOf((*MockStrictReadDatastore)(nil).MetricsID)) +} + +// OfflineFeatures mocks base method. +func (m *MockStrictReadDatastore) OfflineFeatures() (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OfflineFeatures") + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OfflineFeatures indicates an expected call of OfflineFeatures. +func (mr *MockStrictReadDatastoreMockRecorder) OfflineFeatures() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OfflineFeatures", reflect.TypeOf((*MockStrictReadDatastore)(nil).OfflineFeatures)) +} + +// OptimizedRevision mocks base method. +func (m *MockStrictReadDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OptimizedRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OptimizedRevision indicates an expected call of OptimizedRevision. +func (mr *MockStrictReadDatastoreMockRecorder) OptimizedRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OptimizedRevision", reflect.TypeOf((*MockStrictReadDatastore)(nil).OptimizedRevision), ctx) +} + +// ReadWriteTx mocks base method. +func (m *MockStrictReadDatastore) ReadWriteTx(arg0 context.Context, arg1 datastore.TxUserFunc, arg2 ...options.RWTOptionsOption) (datastore.Revision, error) { + m.ctrl.T.Helper() + varargs := []any{arg0, arg1} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "ReadWriteTx", varargs...) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadWriteTx indicates an expected call of ReadWriteTx. +func (mr *MockStrictReadDatastoreMockRecorder) ReadWriteTx(arg0, arg1 any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{arg0, arg1}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadWriteTx", reflect.TypeOf((*MockStrictReadDatastore)(nil).ReadWriteTx), varargs...) +} + +// ReadyState mocks base method. +func (m *MockStrictReadDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadyState", ctx) + ret0, _ := ret[0].(datastore.ReadyState) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadyState indicates an expected call of ReadyState. +func (mr *MockStrictReadDatastoreMockRecorder) ReadyState(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadyState", reflect.TypeOf((*MockStrictReadDatastore)(nil).ReadyState), ctx) +} + +// RevisionFromString mocks base method. +func (m *MockStrictReadDatastore) RevisionFromString(serialized string) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RevisionFromString", serialized) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RevisionFromString indicates an expected call of RevisionFromString. +func (mr *MockStrictReadDatastoreMockRecorder) RevisionFromString(serialized any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RevisionFromString", reflect.TypeOf((*MockStrictReadDatastore)(nil).RevisionFromString), serialized) +} + +// SnapshotReader mocks base method. +func (m *MockStrictReadDatastore) SnapshotReader(arg0 datastore.Revision) datastore.Reader { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SnapshotReader", arg0) + ret0, _ := ret[0].(datastore.Reader) + return ret0 +} + +// SnapshotReader indicates an expected call of SnapshotReader. +func (mr *MockStrictReadDatastoreMockRecorder) SnapshotReader(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SnapshotReader", reflect.TypeOf((*MockStrictReadDatastore)(nil).SnapshotReader), arg0) +} + +// Statistics mocks base method. +func (m *MockStrictReadDatastore) Statistics(ctx context.Context) (datastore.Stats, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Statistics", ctx) + ret0, _ := ret[0].(datastore.Stats) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Statistics indicates an expected call of Statistics. +func (mr *MockStrictReadDatastoreMockRecorder) Statistics(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Statistics", reflect.TypeOf((*MockStrictReadDatastore)(nil).Statistics), ctx) +} + +// UniqueID mocks base method. +func (m *MockStrictReadDatastore) UniqueID(arg0 context.Context) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UniqueID", arg0) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UniqueID indicates an expected call of UniqueID. +func (mr *MockStrictReadDatastoreMockRecorder) UniqueID(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UniqueID", reflect.TypeOf((*MockStrictReadDatastore)(nil).UniqueID), arg0) +} + +// Watch mocks base method. +func (m *MockStrictReadDatastore) Watch(ctx context.Context, afterRevision datastore.Revision, arg2 datastore.WatchOptions) (<-chan datastore.RevisionChanges, <-chan error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Watch", ctx, afterRevision, arg2) + ret0, _ := ret[0].(<-chan datastore.RevisionChanges) + ret1, _ := ret[1].(<-chan error) + return ret0, ret1 +} + +// Watch indicates an expected call of Watch. +func (mr *MockStrictReadDatastoreMockRecorder) Watch(ctx, afterRevision, arg2 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockStrictReadDatastore)(nil).Watch), ctx, afterRevision, arg2) +} + +// MockStartableDatastore is a mock of StartableDatastore interface. +type MockStartableDatastore struct { + ctrl *gomock.Controller + recorder *MockStartableDatastoreMockRecorder + isgomock struct{} +} + +// MockStartableDatastoreMockRecorder is the mock recorder for MockStartableDatastore. +type MockStartableDatastoreMockRecorder struct { + mock *MockStartableDatastore +} + +// NewMockStartableDatastore creates a new mock instance. +func NewMockStartableDatastore(ctrl *gomock.Controller) *MockStartableDatastore { + mock := &MockStartableDatastore{ctrl: ctrl} + mock.recorder = &MockStartableDatastoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockStartableDatastore) EXPECT() *MockStartableDatastoreMockRecorder { + return m.recorder +} + +// CheckRevision mocks base method. +func (m *MockStartableDatastore) CheckRevision(ctx context.Context, revision datastore.Revision) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CheckRevision", ctx, revision) + ret0, _ := ret[0].(error) + return ret0 +} + +// CheckRevision indicates an expected call of CheckRevision. +func (mr *MockStartableDatastoreMockRecorder) CheckRevision(ctx, revision any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRevision", reflect.TypeOf((*MockStartableDatastore)(nil).CheckRevision), ctx, revision) +} + +// Close mocks base method. +func (m *MockStartableDatastore) Close() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close") + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockStartableDatastoreMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockStartableDatastore)(nil).Close)) +} + +// Features mocks base method. +func (m *MockStartableDatastore) Features(ctx context.Context) (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Features", ctx) + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Features indicates an expected call of Features. +func (mr *MockStartableDatastoreMockRecorder) Features(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Features", reflect.TypeOf((*MockStartableDatastore)(nil).Features), ctx) +} + +// HeadRevision mocks base method. +func (m *MockStartableDatastore) HeadRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "HeadRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// HeadRevision indicates an expected call of HeadRevision. +func (mr *MockStartableDatastoreMockRecorder) HeadRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeadRevision", reflect.TypeOf((*MockStartableDatastore)(nil).HeadRevision), ctx) +} + +// MetricsID mocks base method. +func (m *MockStartableDatastore) MetricsID() (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MetricsID") + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// MetricsID indicates an expected call of MetricsID. +func (mr *MockStartableDatastoreMockRecorder) MetricsID() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MetricsID", reflect.TypeOf((*MockStartableDatastore)(nil).MetricsID)) +} + +// OfflineFeatures mocks base method. +func (m *MockStartableDatastore) OfflineFeatures() (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OfflineFeatures") + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OfflineFeatures indicates an expected call of OfflineFeatures. +func (mr *MockStartableDatastoreMockRecorder) OfflineFeatures() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OfflineFeatures", reflect.TypeOf((*MockStartableDatastore)(nil).OfflineFeatures)) +} + +// OptimizedRevision mocks base method. +func (m *MockStartableDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OptimizedRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OptimizedRevision indicates an expected call of OptimizedRevision. +func (mr *MockStartableDatastoreMockRecorder) OptimizedRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OptimizedRevision", reflect.TypeOf((*MockStartableDatastore)(nil).OptimizedRevision), ctx) +} + +// ReadWriteTx mocks base method. +func (m *MockStartableDatastore) ReadWriteTx(arg0 context.Context, arg1 datastore.TxUserFunc, arg2 ...options.RWTOptionsOption) (datastore.Revision, error) { + m.ctrl.T.Helper() + varargs := []any{arg0, arg1} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "ReadWriteTx", varargs...) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadWriteTx indicates an expected call of ReadWriteTx. +func (mr *MockStartableDatastoreMockRecorder) ReadWriteTx(arg0, arg1 any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{arg0, arg1}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadWriteTx", reflect.TypeOf((*MockStartableDatastore)(nil).ReadWriteTx), varargs...) +} + +// ReadyState mocks base method. +func (m *MockStartableDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadyState", ctx) + ret0, _ := ret[0].(datastore.ReadyState) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadyState indicates an expected call of ReadyState. +func (mr *MockStartableDatastoreMockRecorder) ReadyState(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadyState", reflect.TypeOf((*MockStartableDatastore)(nil).ReadyState), ctx) +} + +// RevisionFromString mocks base method. +func (m *MockStartableDatastore) RevisionFromString(serialized string) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RevisionFromString", serialized) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RevisionFromString indicates an expected call of RevisionFromString. +func (mr *MockStartableDatastoreMockRecorder) RevisionFromString(serialized any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RevisionFromString", reflect.TypeOf((*MockStartableDatastore)(nil).RevisionFromString), serialized) +} + +// SnapshotReader mocks base method. +func (m *MockStartableDatastore) SnapshotReader(arg0 datastore.Revision) datastore.Reader { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SnapshotReader", arg0) + ret0, _ := ret[0].(datastore.Reader) + return ret0 +} + +// SnapshotReader indicates an expected call of SnapshotReader. +func (mr *MockStartableDatastoreMockRecorder) SnapshotReader(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SnapshotReader", reflect.TypeOf((*MockStartableDatastore)(nil).SnapshotReader), arg0) +} + +// Start mocks base method. +func (m *MockStartableDatastore) Start(ctx context.Context) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Start", ctx) + ret0, _ := ret[0].(error) + return ret0 +} + +// Start indicates an expected call of Start. +func (mr *MockStartableDatastoreMockRecorder) Start(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Start", reflect.TypeOf((*MockStartableDatastore)(nil).Start), ctx) +} + +// Statistics mocks base method. +func (m *MockStartableDatastore) Statistics(ctx context.Context) (datastore.Stats, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Statistics", ctx) + ret0, _ := ret[0].(datastore.Stats) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Statistics indicates an expected call of Statistics. +func (mr *MockStartableDatastoreMockRecorder) Statistics(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Statistics", reflect.TypeOf((*MockStartableDatastore)(nil).Statistics), ctx) +} + +// UniqueID mocks base method. +func (m *MockStartableDatastore) UniqueID(arg0 context.Context) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UniqueID", arg0) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UniqueID indicates an expected call of UniqueID. +func (mr *MockStartableDatastoreMockRecorder) UniqueID(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UniqueID", reflect.TypeOf((*MockStartableDatastore)(nil).UniqueID), arg0) +} + +// Watch mocks base method. +func (m *MockStartableDatastore) Watch(ctx context.Context, afterRevision datastore.Revision, arg2 datastore.WatchOptions) (<-chan datastore.RevisionChanges, <-chan error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Watch", ctx, afterRevision, arg2) + ret0, _ := ret[0].(<-chan datastore.RevisionChanges) + ret1, _ := ret[1].(<-chan error) + return ret0, ret1 +} + +// Watch indicates an expected call of Watch. +func (mr *MockStartableDatastoreMockRecorder) Watch(ctx, afterRevision, arg2 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockStartableDatastore)(nil).Watch), ctx, afterRevision, arg2) +} + +// MockRepairableDatastore is a mock of RepairableDatastore interface. +type MockRepairableDatastore struct { + ctrl *gomock.Controller + recorder *MockRepairableDatastoreMockRecorder + isgomock struct{} +} + +// MockRepairableDatastoreMockRecorder is the mock recorder for MockRepairableDatastore. +type MockRepairableDatastoreMockRecorder struct { + mock *MockRepairableDatastore +} + +// NewMockRepairableDatastore creates a new mock instance. +func NewMockRepairableDatastore(ctrl *gomock.Controller) *MockRepairableDatastore { + mock := &MockRepairableDatastore{ctrl: ctrl} + mock.recorder = &MockRepairableDatastoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockRepairableDatastore) EXPECT() *MockRepairableDatastoreMockRecorder { + return m.recorder +} + +// CheckRevision mocks base method. +func (m *MockRepairableDatastore) CheckRevision(ctx context.Context, revision datastore.Revision) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CheckRevision", ctx, revision) + ret0, _ := ret[0].(error) + return ret0 +} + +// CheckRevision indicates an expected call of CheckRevision. +func (mr *MockRepairableDatastoreMockRecorder) CheckRevision(ctx, revision any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRevision", reflect.TypeOf((*MockRepairableDatastore)(nil).CheckRevision), ctx, revision) +} + +// Close mocks base method. +func (m *MockRepairableDatastore) Close() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close") + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockRepairableDatastoreMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockRepairableDatastore)(nil).Close)) +} + +// Features mocks base method. +func (m *MockRepairableDatastore) Features(ctx context.Context) (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Features", ctx) + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Features indicates an expected call of Features. +func (mr *MockRepairableDatastoreMockRecorder) Features(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Features", reflect.TypeOf((*MockRepairableDatastore)(nil).Features), ctx) +} + +// HeadRevision mocks base method. +func (m *MockRepairableDatastore) HeadRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "HeadRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// HeadRevision indicates an expected call of HeadRevision. +func (mr *MockRepairableDatastoreMockRecorder) HeadRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HeadRevision", reflect.TypeOf((*MockRepairableDatastore)(nil).HeadRevision), ctx) +} + +// MetricsID mocks base method. +func (m *MockRepairableDatastore) MetricsID() (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MetricsID") + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// MetricsID indicates an expected call of MetricsID. +func (mr *MockRepairableDatastoreMockRecorder) MetricsID() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MetricsID", reflect.TypeOf((*MockRepairableDatastore)(nil).MetricsID)) +} + +// OfflineFeatures mocks base method. +func (m *MockRepairableDatastore) OfflineFeatures() (*datastore.Features, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OfflineFeatures") + ret0, _ := ret[0].(*datastore.Features) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OfflineFeatures indicates an expected call of OfflineFeatures. +func (mr *MockRepairableDatastoreMockRecorder) OfflineFeatures() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OfflineFeatures", reflect.TypeOf((*MockRepairableDatastore)(nil).OfflineFeatures)) +} + +// OptimizedRevision mocks base method. +func (m *MockRepairableDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "OptimizedRevision", ctx) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// OptimizedRevision indicates an expected call of OptimizedRevision. +func (mr *MockRepairableDatastoreMockRecorder) OptimizedRevision(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OptimizedRevision", reflect.TypeOf((*MockRepairableDatastore)(nil).OptimizedRevision), ctx) +} + +// ReadWriteTx mocks base method. +func (m *MockRepairableDatastore) ReadWriteTx(arg0 context.Context, arg1 datastore.TxUserFunc, arg2 ...options.RWTOptionsOption) (datastore.Revision, error) { + m.ctrl.T.Helper() + varargs := []any{arg0, arg1} + for _, a := range arg2 { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "ReadWriteTx", varargs...) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadWriteTx indicates an expected call of ReadWriteTx. +func (mr *MockRepairableDatastoreMockRecorder) ReadWriteTx(arg0, arg1 any, arg2 ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{arg0, arg1}, arg2...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadWriteTx", reflect.TypeOf((*MockRepairableDatastore)(nil).ReadWriteTx), varargs...) +} + +// ReadyState mocks base method. +func (m *MockRepairableDatastore) ReadyState(ctx context.Context) (datastore.ReadyState, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadyState", ctx) + ret0, _ := ret[0].(datastore.ReadyState) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReadyState indicates an expected call of ReadyState. +func (mr *MockRepairableDatastoreMockRecorder) ReadyState(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadyState", reflect.TypeOf((*MockRepairableDatastore)(nil).ReadyState), ctx) +} + +// Repair mocks base method. +func (m *MockRepairableDatastore) Repair(ctx context.Context, operationName string, outputProgress bool) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Repair", ctx, operationName, outputProgress) + ret0, _ := ret[0].(error) + return ret0 +} + +// Repair indicates an expected call of Repair. +func (mr *MockRepairableDatastoreMockRecorder) Repair(ctx, operationName, outputProgress any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Repair", reflect.TypeOf((*MockRepairableDatastore)(nil).Repair), ctx, operationName, outputProgress) +} + +// RepairOperations mocks base method. +func (m *MockRepairableDatastore) RepairOperations() []datastore.RepairOperation { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RepairOperations") + ret0, _ := ret[0].([]datastore.RepairOperation) + return ret0 +} + +// RepairOperations indicates an expected call of RepairOperations. +func (mr *MockRepairableDatastoreMockRecorder) RepairOperations() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RepairOperations", reflect.TypeOf((*MockRepairableDatastore)(nil).RepairOperations)) +} + +// RevisionFromString mocks base method. +func (m *MockRepairableDatastore) RevisionFromString(serialized string) (datastore.Revision, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RevisionFromString", serialized) + ret0, _ := ret[0].(datastore.Revision) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RevisionFromString indicates an expected call of RevisionFromString. +func (mr *MockRepairableDatastoreMockRecorder) RevisionFromString(serialized any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RevisionFromString", reflect.TypeOf((*MockRepairableDatastore)(nil).RevisionFromString), serialized) +} + +// SnapshotReader mocks base method. +func (m *MockRepairableDatastore) SnapshotReader(arg0 datastore.Revision) datastore.Reader { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "SnapshotReader", arg0) + ret0, _ := ret[0].(datastore.Reader) + return ret0 +} + +// SnapshotReader indicates an expected call of SnapshotReader. +func (mr *MockRepairableDatastoreMockRecorder) SnapshotReader(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SnapshotReader", reflect.TypeOf((*MockRepairableDatastore)(nil).SnapshotReader), arg0) +} + +// Statistics mocks base method. +func (m *MockRepairableDatastore) Statistics(ctx context.Context) (datastore.Stats, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Statistics", ctx) + ret0, _ := ret[0].(datastore.Stats) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Statistics indicates an expected call of Statistics. +func (mr *MockRepairableDatastoreMockRecorder) Statistics(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Statistics", reflect.TypeOf((*MockRepairableDatastore)(nil).Statistics), ctx) +} + +// UniqueID mocks base method. +func (m *MockRepairableDatastore) UniqueID(arg0 context.Context) (string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UniqueID", arg0) + ret0, _ := ret[0].(string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UniqueID indicates an expected call of UniqueID. +func (mr *MockRepairableDatastoreMockRecorder) UniqueID(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UniqueID", reflect.TypeOf((*MockRepairableDatastore)(nil).UniqueID), arg0) +} + +// Watch mocks base method. +func (m *MockRepairableDatastore) Watch(ctx context.Context, afterRevision datastore.Revision, arg2 datastore.WatchOptions) (<-chan datastore.RevisionChanges, <-chan error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Watch", ctx, afterRevision, arg2) + ret0, _ := ret[0].(<-chan datastore.RevisionChanges) + ret1, _ := ret[1].(<-chan error) + return ret0, ret1 +} + +// Watch indicates an expected call of Watch. +func (mr *MockRepairableDatastoreMockRecorder) Watch(ctx, afterRevision, arg2 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Watch", reflect.TypeOf((*MockRepairableDatastore)(nil).Watch), ctx, afterRevision, arg2) +} + +// MockUnwrappableDatastore is a mock of UnwrappableDatastore interface. +type MockUnwrappableDatastore struct { + ctrl *gomock.Controller + recorder *MockUnwrappableDatastoreMockRecorder + isgomock struct{} +} + +// MockUnwrappableDatastoreMockRecorder is the mock recorder for MockUnwrappableDatastore. +type MockUnwrappableDatastoreMockRecorder struct { + mock *MockUnwrappableDatastore +} + +// NewMockUnwrappableDatastore creates a new mock instance. +func NewMockUnwrappableDatastore(ctrl *gomock.Controller) *MockUnwrappableDatastore { + mock := &MockUnwrappableDatastore{ctrl: ctrl} + mock.recorder = &MockUnwrappableDatastoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockUnwrappableDatastore) EXPECT() *MockUnwrappableDatastoreMockRecorder { + return m.recorder +} + +// Unwrap mocks base method. +func (m *MockUnwrappableDatastore) Unwrap() datastore.Datastore { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Unwrap") + ret0, _ := ret[0].(datastore.Datastore) + return ret0 +} + +// Unwrap indicates an expected call of Unwrap. +func (mr *MockUnwrappableDatastoreMockRecorder) Unwrap() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Unwrap", reflect.TypeOf((*MockUnwrappableDatastore)(nil).Unwrap)) +} + +// MockRevision is a mock of Revision interface. +type MockRevision struct { + ctrl *gomock.Controller + recorder *MockRevisionMockRecorder + isgomock struct{} +} + +// MockRevisionMockRecorder is the mock recorder for MockRevision. +type MockRevisionMockRecorder struct { + mock *MockRevision +} + +// NewMockRevision creates a new mock instance. +func NewMockRevision(ctrl *gomock.Controller) *MockRevision { + mock := &MockRevision{ctrl: ctrl} + mock.recorder = &MockRevisionMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockRevision) EXPECT() *MockRevisionMockRecorder { + return m.recorder +} + +// ByteSortable mocks base method. +func (m *MockRevision) ByteSortable() bool { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ByteSortable") + ret0, _ := ret[0].(bool) + return ret0 +} + +// ByteSortable indicates an expected call of ByteSortable. +func (mr *MockRevisionMockRecorder) ByteSortable() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ByteSortable", reflect.TypeOf((*MockRevision)(nil).ByteSortable)) +} + +// Equal mocks base method. +func (m *MockRevision) Equal(arg0 datastore.Revision) bool { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Equal", arg0) + ret0, _ := ret[0].(bool) + return ret0 +} + +// Equal indicates an expected call of Equal. +func (mr *MockRevisionMockRecorder) Equal(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Equal", reflect.TypeOf((*MockRevision)(nil).Equal), arg0) +} + +// GreaterThan mocks base method. +func (m *MockRevision) GreaterThan(arg0 datastore.Revision) bool { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GreaterThan", arg0) + ret0, _ := ret[0].(bool) + return ret0 +} + +// GreaterThan indicates an expected call of GreaterThan. +func (mr *MockRevisionMockRecorder) GreaterThan(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GreaterThan", reflect.TypeOf((*MockRevision)(nil).GreaterThan), arg0) +} + +// LessThan mocks base method. +func (m *MockRevision) LessThan(arg0 datastore.Revision) bool { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LessThan", arg0) + ret0, _ := ret[0].(bool) + return ret0 +} + +// LessThan indicates an expected call of LessThan. +func (mr *MockRevisionMockRecorder) LessThan(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LessThan", reflect.TypeOf((*MockRevision)(nil).LessThan), arg0) +} + +// String mocks base method. +func (m *MockRevision) String() string { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "String") + ret0, _ := ret[0].(string) + return ret0 +} + +// String indicates an expected call of String. +func (mr *MockRevisionMockRecorder) String() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "String", reflect.TypeOf((*MockRevision)(nil).String)) +} diff --git a/internal/mocks/mock_dispatcher.go b/internal/mocks/mock_dispatcher.go new file mode 100644 index 000000000..01725702f --- /dev/null +++ b/internal/mocks/mock_dispatcher.go @@ -0,0 +1,386 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: dispatch.go +// +// Generated by this command: +// +// mockgen -source dispatch.go -destination ../mocks/mock_dispatcher.go -package mocks Dispatcher +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + reflect "reflect" + + dispatch "github.com/authzed/spicedb/internal/dispatch" + dispatchv1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1" + zerolog "github.com/rs/zerolog" + gomock "go.uber.org/mock/gomock" +) + +// MockDispatcher is a mock of Dispatcher interface. +type MockDispatcher struct { + ctrl *gomock.Controller + recorder *MockDispatcherMockRecorder + isgomock struct{} +} + +// MockDispatcherMockRecorder is the mock recorder for MockDispatcher. +type MockDispatcherMockRecorder struct { + mock *MockDispatcher +} + +// NewMockDispatcher creates a new mock instance. +func NewMockDispatcher(ctrl *gomock.Controller) *MockDispatcher { + mock := &MockDispatcher{ctrl: ctrl} + mock.recorder = &MockDispatcherMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockDispatcher) EXPECT() *MockDispatcherMockRecorder { + return m.recorder +} + +// Close mocks base method. +func (m *MockDispatcher) Close() error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Close") + ret0, _ := ret[0].(error) + return ret0 +} + +// Close indicates an expected call of Close. +func (mr *MockDispatcherMockRecorder) Close() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Close", reflect.TypeOf((*MockDispatcher)(nil).Close)) +} + +// DispatchCheck mocks base method. +func (m *MockDispatcher) DispatchCheck(ctx context.Context, req *dispatchv1.DispatchCheckRequest) (*dispatchv1.DispatchCheckResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchCheck", ctx, req) + ret0, _ := ret[0].(*dispatchv1.DispatchCheckResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DispatchCheck indicates an expected call of DispatchCheck. +func (mr *MockDispatcherMockRecorder) DispatchCheck(ctx, req any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchCheck", reflect.TypeOf((*MockDispatcher)(nil).DispatchCheck), ctx, req) +} + +// DispatchExpand mocks base method. +func (m *MockDispatcher) DispatchExpand(ctx context.Context, req *dispatchv1.DispatchExpandRequest) (*dispatchv1.DispatchExpandResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchExpand", ctx, req) + ret0, _ := ret[0].(*dispatchv1.DispatchExpandResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DispatchExpand indicates an expected call of DispatchExpand. +func (mr *MockDispatcherMockRecorder) DispatchExpand(ctx, req any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchExpand", reflect.TypeOf((*MockDispatcher)(nil).DispatchExpand), ctx, req) +} + +// DispatchLookupResources2 mocks base method. +func (m *MockDispatcher) DispatchLookupResources2(req *dispatchv1.DispatchLookupResources2Request, stream dispatch.LookupResources2Stream) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchLookupResources2", req, stream) + ret0, _ := ret[0].(error) + return ret0 +} + +// DispatchLookupResources2 indicates an expected call of DispatchLookupResources2. +func (mr *MockDispatcherMockRecorder) DispatchLookupResources2(req, stream any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchLookupResources2", reflect.TypeOf((*MockDispatcher)(nil).DispatchLookupResources2), req, stream) +} + +// DispatchLookupResources3 mocks base method. +func (m *MockDispatcher) DispatchLookupResources3(req *dispatchv1.DispatchLookupResources3Request, stream dispatch.LookupResources3Stream) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchLookupResources3", req, stream) + ret0, _ := ret[0].(error) + return ret0 +} + +// DispatchLookupResources3 indicates an expected call of DispatchLookupResources3. +func (mr *MockDispatcherMockRecorder) DispatchLookupResources3(req, stream any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchLookupResources3", reflect.TypeOf((*MockDispatcher)(nil).DispatchLookupResources3), req, stream) +} + +// DispatchLookupSubjects mocks base method. +func (m *MockDispatcher) DispatchLookupSubjects(req *dispatchv1.DispatchLookupSubjectsRequest, stream dispatch.LookupSubjectsStream) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchLookupSubjects", req, stream) + ret0, _ := ret[0].(error) + return ret0 +} + +// DispatchLookupSubjects indicates an expected call of DispatchLookupSubjects. +func (mr *MockDispatcherMockRecorder) DispatchLookupSubjects(req, stream any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchLookupSubjects", reflect.TypeOf((*MockDispatcher)(nil).DispatchLookupSubjects), req, stream) +} + +// ReadyState mocks base method. +func (m *MockDispatcher) ReadyState() dispatch.ReadyState { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReadyState") + ret0, _ := ret[0].(dispatch.ReadyState) + return ret0 +} + +// ReadyState indicates an expected call of ReadyState. +func (mr *MockDispatcherMockRecorder) ReadyState() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReadyState", reflect.TypeOf((*MockDispatcher)(nil).ReadyState)) +} + +// MockCheck is a mock of Check interface. +type MockCheck struct { + ctrl *gomock.Controller + recorder *MockCheckMockRecorder + isgomock struct{} +} + +// MockCheckMockRecorder is the mock recorder for MockCheck. +type MockCheckMockRecorder struct { + mock *MockCheck +} + +// NewMockCheck creates a new mock instance. +func NewMockCheck(ctrl *gomock.Controller) *MockCheck { + mock := &MockCheck{ctrl: ctrl} + mock.recorder = &MockCheckMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockCheck) EXPECT() *MockCheckMockRecorder { + return m.recorder +} + +// DispatchCheck mocks base method. +func (m *MockCheck) DispatchCheck(ctx context.Context, req *dispatchv1.DispatchCheckRequest) (*dispatchv1.DispatchCheckResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchCheck", ctx, req) + ret0, _ := ret[0].(*dispatchv1.DispatchCheckResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DispatchCheck indicates an expected call of DispatchCheck. +func (mr *MockCheckMockRecorder) DispatchCheck(ctx, req any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchCheck", reflect.TypeOf((*MockCheck)(nil).DispatchCheck), ctx, req) +} + +// MockExpand is a mock of Expand interface. +type MockExpand struct { + ctrl *gomock.Controller + recorder *MockExpandMockRecorder + isgomock struct{} +} + +// MockExpandMockRecorder is the mock recorder for MockExpand. +type MockExpandMockRecorder struct { + mock *MockExpand +} + +// NewMockExpand creates a new mock instance. +func NewMockExpand(ctrl *gomock.Controller) *MockExpand { + mock := &MockExpand{ctrl: ctrl} + mock.recorder = &MockExpandMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockExpand) EXPECT() *MockExpandMockRecorder { + return m.recorder +} + +// DispatchExpand mocks base method. +func (m *MockExpand) DispatchExpand(ctx context.Context, req *dispatchv1.DispatchExpandRequest) (*dispatchv1.DispatchExpandResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchExpand", ctx, req) + ret0, _ := ret[0].(*dispatchv1.DispatchExpandResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DispatchExpand indicates an expected call of DispatchExpand. +func (mr *MockExpandMockRecorder) DispatchExpand(ctx, req any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchExpand", reflect.TypeOf((*MockExpand)(nil).DispatchExpand), ctx, req) +} + +// MockLookupResources2 is a mock of LookupResources2 interface. +type MockLookupResources2 struct { + ctrl *gomock.Controller + recorder *MockLookupResources2MockRecorder + isgomock struct{} +} + +// MockLookupResources2MockRecorder is the mock recorder for MockLookupResources2. +type MockLookupResources2MockRecorder struct { + mock *MockLookupResources2 +} + +// NewMockLookupResources2 creates a new mock instance. +func NewMockLookupResources2(ctrl *gomock.Controller) *MockLookupResources2 { + mock := &MockLookupResources2{ctrl: ctrl} + mock.recorder = &MockLookupResources2MockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockLookupResources2) EXPECT() *MockLookupResources2MockRecorder { + return m.recorder +} + +// DispatchLookupResources2 mocks base method. +func (m *MockLookupResources2) DispatchLookupResources2(req *dispatchv1.DispatchLookupResources2Request, stream dispatch.LookupResources2Stream) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchLookupResources2", req, stream) + ret0, _ := ret[0].(error) + return ret0 +} + +// DispatchLookupResources2 indicates an expected call of DispatchLookupResources2. +func (mr *MockLookupResources2MockRecorder) DispatchLookupResources2(req, stream any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchLookupResources2", reflect.TypeOf((*MockLookupResources2)(nil).DispatchLookupResources2), req, stream) +} + +// MockLookupResources3 is a mock of LookupResources3 interface. +type MockLookupResources3 struct { + ctrl *gomock.Controller + recorder *MockLookupResources3MockRecorder + isgomock struct{} +} + +// MockLookupResources3MockRecorder is the mock recorder for MockLookupResources3. +type MockLookupResources3MockRecorder struct { + mock *MockLookupResources3 +} + +// NewMockLookupResources3 creates a new mock instance. +func NewMockLookupResources3(ctrl *gomock.Controller) *MockLookupResources3 { + mock := &MockLookupResources3{ctrl: ctrl} + mock.recorder = &MockLookupResources3MockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockLookupResources3) EXPECT() *MockLookupResources3MockRecorder { + return m.recorder +} + +// DispatchLookupResources3 mocks base method. +func (m *MockLookupResources3) DispatchLookupResources3(req *dispatchv1.DispatchLookupResources3Request, stream dispatch.LookupResources3Stream) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchLookupResources3", req, stream) + ret0, _ := ret[0].(error) + return ret0 +} + +// DispatchLookupResources3 indicates an expected call of DispatchLookupResources3. +func (mr *MockLookupResources3MockRecorder) DispatchLookupResources3(req, stream any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchLookupResources3", reflect.TypeOf((*MockLookupResources3)(nil).DispatchLookupResources3), req, stream) +} + +// MockLookupSubjects is a mock of LookupSubjects interface. +type MockLookupSubjects struct { + ctrl *gomock.Controller + recorder *MockLookupSubjectsMockRecorder + isgomock struct{} +} + +// MockLookupSubjectsMockRecorder is the mock recorder for MockLookupSubjects. +type MockLookupSubjectsMockRecorder struct { + mock *MockLookupSubjects +} + +// NewMockLookupSubjects creates a new mock instance. +func NewMockLookupSubjects(ctrl *gomock.Controller) *MockLookupSubjects { + mock := &MockLookupSubjects{ctrl: ctrl} + mock.recorder = &MockLookupSubjectsMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockLookupSubjects) EXPECT() *MockLookupSubjectsMockRecorder { + return m.recorder +} + +// DispatchLookupSubjects mocks base method. +func (m *MockLookupSubjects) DispatchLookupSubjects(req *dispatchv1.DispatchLookupSubjectsRequest, stream dispatch.LookupSubjectsStream) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DispatchLookupSubjects", req, stream) + ret0, _ := ret[0].(error) + return ret0 +} + +// DispatchLookupSubjects indicates an expected call of DispatchLookupSubjects. +func (mr *MockLookupSubjectsMockRecorder) DispatchLookupSubjects(req, stream any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DispatchLookupSubjects", reflect.TypeOf((*MockLookupSubjects)(nil).DispatchLookupSubjects), req, stream) +} + +// MockDispatchableRequest is a mock of DispatchableRequest interface. +type MockDispatchableRequest struct { + ctrl *gomock.Controller + recorder *MockDispatchableRequestMockRecorder + isgomock struct{} +} + +// MockDispatchableRequestMockRecorder is the mock recorder for MockDispatchableRequest. +type MockDispatchableRequestMockRecorder struct { + mock *MockDispatchableRequest +} + +// NewMockDispatchableRequest creates a new mock instance. +func NewMockDispatchableRequest(ctrl *gomock.Controller) *MockDispatchableRequest { + mock := &MockDispatchableRequest{ctrl: ctrl} + mock.recorder = &MockDispatchableRequestMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockDispatchableRequest) EXPECT() *MockDispatchableRequestMockRecorder { + return m.recorder +} + +// GetMetadata mocks base method. +func (m *MockDispatchableRequest) GetMetadata() *dispatchv1.ResolverMeta { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetMetadata") + ret0, _ := ret[0].(*dispatchv1.ResolverMeta) + return ret0 +} + +// GetMetadata indicates an expected call of GetMetadata. +func (mr *MockDispatchableRequestMockRecorder) GetMetadata() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMetadata", reflect.TypeOf((*MockDispatchableRequest)(nil).GetMetadata)) +} + +// MarshalZerologObject mocks base method. +func (m *MockDispatchableRequest) MarshalZerologObject(e *zerolog.Event) { + m.ctrl.T.Helper() + m.ctrl.Call(m, "MarshalZerologObject", e) +} + +// MarshalZerologObject indicates an expected call of MarshalZerologObject. +func (mr *MockDispatchableRequestMockRecorder) MarshalZerologObject(e any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarshalZerologObject", reflect.TypeOf((*MockDispatchableRequest)(nil).MarshalZerologObject), e) +} diff --git a/magefiles/tools.go b/magefiles/tools.go index 0f88baab7..2a62d35d9 100644 --- a/magefiles/tools.go +++ b/magefiles/tools.go @@ -10,6 +10,7 @@ import ( _ "github.com/envoyproxy/protoc-gen-validate" _ "github.com/planetscale/vtprotobuf/cmd/protoc-gen-go-vtproto" _ "github.com/planetscale/vtprotobuf/protohelpers" + _ "go.uber.org/mock/mockgen" _ "golang.org/x/tools/cmd/stringer" _ "google.golang.org/grpc/cmd/protoc-gen-go-grpc" _ "google.golang.org/protobuf/cmd/protoc-gen-go" diff --git a/pkg/cmd/serve.go b/pkg/cmd/serve.go index de8b6975b..284c1f425 100644 --- a/pkg/cmd/serve.go +++ b/pkg/cmd/serve.go @@ -121,6 +121,12 @@ func RegisterServeFlags(cmd *cobra.Command, config *server.Config) error { apiFlags.BoolVar(&config.EnablePerformanceInsightMetrics, "enable-performance-insight-metrics", false, "enables performance insight metrics, which are used to track the latency of API calls by shape") apiFlags.StringVar(&config.MismatchZedTokenBehavior, "mismatch-zed-token-behavior", "full-consistency", "behavior to enforce when an API call receives a zedtoken that was originally intended for a different kind of datastore. One of: full-consistency (treat as a full-consistency call, ignoring the zedtoken), min-latency (treat as a min-latency call, ignoring the zedtoken), error (return an error). defaults to full-consistency for safety.") + // Memory Protection flags + apiFlags.BoolVar(&config.MemoryProtectionEnabled, "memory-protection-enabled", true, "enables a memory-based middleware that rejects requests with error code ResourceExhausted when memory usage is too high. Use in conjunction with GOMEMLIMIT.") + apiFlags.Float64Var(&config.MemoryProtectionNormalAPIThresholdPercent, "memory-protection-normal-api-threshold", 0.90, "float percentage (where 1 = 100%) of memory usage such that when the server receives an external API request and the server exceeeds this usage, it will reject it") + apiFlags.Float64Var(&config.MemoryProtectionDispatchAPIThresholdPercent, "memory-protection-dispatch-api-threshold", 0.95, "float percentage (where 1 = 100%) of memory usage such that when the server receives an internal (dispatch) API request and the server exceeeds this usage, it will reject it") + apiFlags.IntVar(&config.MemoryProtectionSampleIntervalSeconds, "memory-protection-sample-interval", 1, "how often to sample current memory usage in seconds") + datastoreFlags := nfs.FlagSet(BoldBlue("Datastore")) // Flags for the datastore if err := datastore.RegisterDatastoreFlags(datastoreFlags, &config.DatastoreConfig); err != nil { diff --git a/pkg/cmd/server/defaults.go b/pkg/cmd/server/defaults.go index 2cf514eda..0e1772513 100644 --- a/pkg/cmd/server/defaults.go +++ b/pkg/cmd/server/defaults.go @@ -35,6 +35,7 @@ import ( "github.com/authzed/spicedb/internal/logging" datastoremw "github.com/authzed/spicedb/internal/middleware/datastore" dispatchmw "github.com/authzed/spicedb/internal/middleware/dispatcher" + "github.com/authzed/spicedb/internal/middleware/memoryprotection" "github.com/authzed/spicedb/internal/middleware/servicespecific" "github.com/authzed/spicedb/pkg/datastore" consistencymw "github.com/authzed/spicedb/pkg/middleware/consistency" @@ -165,12 +166,13 @@ var alwaysDebugOption = grpclog.WithLevels(func(code codes.Code) grpclog.Level { }) const ( - DefaultMiddlewareRequestID = "requestid" - DefaultMiddlewareLog = "log" - DefaultMiddlewareGRPCLog = "grpclog" - DefaultMiddlewareGRPCAuth = "grpcauth" - DefaultMiddlewareGRPCProm = "grpcprom" - DefaultMiddlewareServerVersion = "serverversion" + DefaultMiddlewareRequestID = "requestid" + DefaultMiddlewareLog = "log" + DefaultMiddlewareGRPCLog = "grpclog" + DefaultMiddlewareGRPCAuth = "grpcauth" + DefaultMiddlewareGRPCProm = "grpcprom" + DefaultMiddlewareServerVersion = "serverversion" + DefaultMiddlewareMemoryProtection = "memoryprotection" DefaultInternalMiddlewareDispatch = "dispatch" DefaultInternalMiddlewareDatastore = "datastore" @@ -190,6 +192,10 @@ type MiddlewareOption struct { MiddlewareServiceLabel string `debugmap:"visible"` MismatchingZedTokenOption consistencymw.MismatchingTokenOption `debugmap:"visible"` + // Memory Protection + MemoryProtectionConfig memoryprotection.Config `debugmap:"visible"` + MemorySampler memoryprotection.MemorySampler `debugmap:"hidden"` + unaryDatastoreMiddleware *ReferenceableMiddleware[grpc.UnaryServerInterceptor] `debugmap:"hidden"` streamDatastoreMiddleware *ReferenceableMiddleware[grpc.StreamServerInterceptor] `debugmap:"hidden"` } @@ -212,19 +218,10 @@ func (m MiddlewareOption) WithDatastoreMiddleware(middleware Middleware) Middlew WithInterceptor(middleware.StreamServerInterceptor()). Done() - return MiddlewareOption{ - Logger: m.Logger, - AuthFunc: m.AuthFunc, - EnableVersionResponse: m.EnableVersionResponse, - DispatcherForMiddleware: m.DispatcherForMiddleware, - EnableRequestLog: m.EnableRequestLog, - EnableResponseLog: m.EnableResponseLog, - DisableGRPCHistogram: m.DisableGRPCHistogram, - MiddlewareServiceLabel: m.MiddlewareServiceLabel, - MismatchingZedTokenOption: m.MismatchingZedTokenOption, - unaryDatastoreMiddleware: &unary, - streamDatastoreMiddleware: &stream, - } + m.unaryDatastoreMiddleware = &unary + m.streamDatastoreMiddleware = &stream + + return m } func (m MiddlewareOption) WithDatastore(ds datastore.Datastore) MiddlewareOption { @@ -240,19 +237,10 @@ func (m MiddlewareOption) WithDatastore(ds datastore.Datastore) MiddlewareOption WithInterceptor(datastoremw.StreamServerInterceptor(ds)). Done() - return MiddlewareOption{ - Logger: m.Logger, - AuthFunc: m.AuthFunc, - EnableVersionResponse: m.EnableVersionResponse, - DispatcherForMiddleware: m.DispatcherForMiddleware, - EnableRequestLog: m.EnableRequestLog, - EnableResponseLog: m.EnableResponseLog, - DisableGRPCHistogram: m.DisableGRPCHistogram, - MiddlewareServiceLabel: m.MiddlewareServiceLabel, - MismatchingZedTokenOption: m.MismatchingZedTokenOption, - unaryDatastoreMiddleware: &unary, - streamDatastoreMiddleware: &stream, - } + m.unaryDatastoreMiddleware = &unary + m.streamDatastoreMiddleware = &stream + + return m } // gRPCMetricsUnaryInterceptor creates the default prometheus metrics interceptor for unary gRPCs @@ -289,6 +277,7 @@ func doesNotMatchRoute(route string) func(_ context.Context, c interceptors.Call // DefaultUnaryMiddleware generates the default middleware chain used for the public SpiceDB Unary gRPC methods func DefaultUnaryMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.UnaryServerInterceptor], error) { grpcMetricsUnaryInterceptor, _ := GRPCMetrics(opts.DisableGRPCHistogram) + memoryProtectionUnaryInterceptor := memoryprotection.New(opts.MemoryProtectionConfig, opts.MemorySampler, "unary-middleware") chain, err := NewMiddlewareChain([]ReferenceableMiddleware[grpc.UnaryServerInterceptor]{ NewUnaryMiddleware(). WithName(DefaultMiddlewareRequestID). @@ -319,10 +308,16 @@ func DefaultUnaryMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.UnaryS WithInterceptor(grpcMetricsUnaryInterceptor). Done(), + NewUnaryMiddleware(). + WithName(DefaultMiddlewareMemoryProtection). + WithInterceptor(memoryProtectionUnaryInterceptor.UnaryServerInterceptor()). + Done(), + NewUnaryMiddleware(). WithName(DefaultMiddlewareGRPCAuth). WithInterceptor(grpcauth.UnaryServerInterceptor(opts.AuthFunc)). - EnsureAlreadyExecuted(DefaultMiddlewareGRPCProm). // so that prom middleware reports auth failures + EnsureAlreadyExecuted(DefaultMiddlewareGRPCProm). // so that prom middleware reports auth failures + EnsureAlreadyExecuted(DefaultMiddlewareMemoryProtection). // so that memory protection happens before auth Done(), NewUnaryMiddleware(). @@ -355,6 +350,7 @@ func DefaultUnaryMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.UnaryS // DefaultStreamingMiddleware generates the default middleware chain used for the public SpiceDB Streaming gRPC methods func DefaultStreamingMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.StreamServerInterceptor], error) { _, grpcMetricsStreamingInterceptor := GRPCMetrics(opts.DisableGRPCHistogram) + memoryProtectionStreamInterceptor := memoryprotection.New(opts.MemoryProtectionConfig, opts.MemorySampler, "stream-middleware") chain, err := NewMiddlewareChain([]ReferenceableMiddleware[grpc.StreamServerInterceptor]{ NewStreamMiddleware(). WithName(DefaultMiddlewareRequestID). @@ -385,10 +381,16 @@ func DefaultStreamingMiddleware(opts MiddlewareOption) (*MiddlewareChain[grpc.St WithInterceptor(grpcMetricsStreamingInterceptor). Done(), + NewStreamMiddleware(). + WithName(DefaultMiddlewareMemoryProtection). + WithInterceptor(memoryProtectionStreamInterceptor.StreamServerInterceptor()). + Done(), + NewStreamMiddleware(). WithName(DefaultMiddlewareGRPCAuth). WithInterceptor(grpcauth.StreamServerInterceptor(opts.AuthFunc)). - EnsureInterceptorAlreadyExecuted(DefaultMiddlewareGRPCProm). // so that prom middleware reports auth failures + EnsureInterceptorAlreadyExecuted(DefaultMiddlewareGRPCProm). // so that prom middleware reports auth failures + EnsureInterceptorAlreadyExecuted(DefaultMiddlewareMemoryProtection). // so that memory protection happens before auth Done(), NewStreamMiddleware(). @@ -432,15 +434,16 @@ func determineEventsToLog(opts MiddlewareOption) grpclog.Option { } // DefaultDispatchMiddleware generates the default middleware chain used for the internal dispatch SpiceDB gRPC API -func DefaultDispatchMiddleware(logger zerolog.Logger, authFunc grpcauth.AuthFunc, ds datastore.Datastore, - disableGRPCLatencyHistogram bool, -) ([]grpc.UnaryServerInterceptor, []grpc.StreamServerInterceptor) { +func DefaultDispatchMiddleware(logger zerolog.Logger, authFunc grpcauth.AuthFunc, ds datastore.Datastore, disableGRPCLatencyHistogram bool, memoryConfig memoryprotection.Config, sampler memoryprotection.MemorySampler) ([]grpc.UnaryServerInterceptor, []grpc.StreamServerInterceptor) { grpcMetricsUnaryInterceptor, grpcMetricsStreamingInterceptor := GRPCMetrics(disableGRPCLatencyHistogram) + dispatchMemoryProtection := memoryprotection.New(memoryConfig, sampler, "dispatch-middleware") + return []grpc.UnaryServerInterceptor{ requestid.UnaryServerInterceptor(requestid.GenerateIfMissing(true)), logmw.UnaryServerInterceptor(logmw.ExtractMetadataField(string(requestmeta.RequestIDKey), "requestID")), grpclog.UnaryServerInterceptor(InterceptorLogger(logger), dispatchDefaultCodeToLevel, durationFieldOption, traceIDFieldOption), grpcMetricsUnaryInterceptor, + dispatchMemoryProtection.UnaryServerInterceptor(), grpcauth.UnaryServerInterceptor(authFunc), datastoremw.UnaryServerInterceptor(ds), servicespecific.UnaryServerInterceptor, @@ -449,6 +452,7 @@ func DefaultDispatchMiddleware(logger zerolog.Logger, authFunc grpcauth.AuthFunc // when returning streaming messages. requestid.StreamServerInterceptor(requestid.GenerateIfMissing(true)), grpcMetricsStreamingInterceptor, + dispatchMemoryProtection.StreamServerInterceptor(), grpcauth.StreamServerInterceptor(authFunc), datastoremw.StreamServerInterceptor(ds), servicespecific.StreamServerInterceptor, diff --git a/pkg/cmd/server/defaults_test.go b/pkg/cmd/server/defaults_test.go index c57116b2f..b6349f4b3 100644 --- a/pkg/cmd/server/defaults_test.go +++ b/pkg/cmd/server/defaults_test.go @@ -11,6 +11,7 @@ import ( "github.com/authzed/spicedb/internal/datastore/memdb" "github.com/authzed/spicedb/internal/dispatch" + "github.com/authzed/spicedb/internal/middleware/memoryprotection" "github.com/authzed/spicedb/internal/middleware/pertoken" caveattypes "github.com/authzed/spicedb/pkg/caveats/types" "github.com/authzed/spicedb/pkg/middleware/consistency" @@ -33,6 +34,8 @@ func TestWithDatastore(t *testing.T) { false, "service", consistency.TreatMismatchingTokensAsError, + memoryprotection.Config{ThresholdPercent: 0}, + memoryprotection.NewMemorySamplerOnInterval(memoryprotection.DefaultSampleIntervalSeconds, &memoryprotection.DefaultMemoryLimitProvider{}), nil, nil, } @@ -73,8 +76,10 @@ func TestWithDatastoreMiddleware(t *testing.T) { true, true, false, - "anotherservice", + "service", consistency.TreatMismatchingTokensAsError, + memoryprotection.Config{ThresholdPercent: 0}, + memoryprotection.NewMemorySamplerOnInterval(memoryprotection.DefaultSampleIntervalSeconds, &memoryprotection.DefaultMemoryLimitProvider{}), nil, nil, } diff --git a/pkg/cmd/server/server.go b/pkg/cmd/server/server.go index ab9d31230..8c05ddb3b 100644 --- a/pkg/cmd/server/server.go +++ b/pkg/cmd/server/server.go @@ -37,6 +37,7 @@ import ( "github.com/authzed/spicedb/internal/dispatch/keys" "github.com/authzed/spicedb/internal/gateway" log "github.com/authzed/spicedb/internal/logging" + "github.com/authzed/spicedb/internal/middleware/memoryprotection" "github.com/authzed/spicedb/internal/services" dispatchSvc "github.com/authzed/spicedb/internal/services/dispatch" "github.com/authzed/spicedb/internal/services/health" @@ -141,6 +142,12 @@ type Config struct { UnaryMiddlewareModification []MiddlewareModification[grpc.UnaryServerInterceptor] `debugmap:"hidden"` StreamingMiddlewareModification []MiddlewareModification[grpc.StreamServerInterceptor] `debugmap:"hidden"` + // Memory Protection + MemoryProtectionEnabled bool `debugmap:"visible" default:"true"` + MemoryProtectionNormalAPIThresholdPercent float64 `debugmap:"visible" default:"0.90"` + MemoryProtectionDispatchAPIThresholdPercent float64 `debugmap:"visible" default:"0.95"` + MemoryProtectionSampleIntervalSeconds int `debugmap:"visible" default:"1"` + // Middleware for internal dispatch API DispatchUnaryMiddleware []grpc.UnaryServerInterceptor `debugmap:"hidden"` DispatchStreamingMiddleware []grpc.StreamServerInterceptor `debugmap:"hidden"` @@ -335,14 +342,6 @@ func (c *Config) Complete(ctx context.Context) (RunnableServer, error) { } closeables.AddWithError(dispatcher.Close) - if len(c.DispatchUnaryMiddleware) == 0 && len(c.DispatchStreamingMiddleware) == 0 { - if c.GRPCAuthFunc == nil { - c.DispatchUnaryMiddleware, c.DispatchStreamingMiddleware = DefaultDispatchMiddleware(log.Logger, auth.MustRequirePresharedKey(c.PresharedSecureKey), ds, c.DisableGRPCLatencyHistogram) - } else { - c.DispatchUnaryMiddleware, c.DispatchStreamingMiddleware = DefaultDispatchMiddleware(log.Logger, c.GRPCAuthFunc, ds, c.DisableGRPCLatencyHistogram) - } - } - var cachingClusterDispatch dispatch.Dispatcher if c.DispatchServer.Enabled { cdcc, err := CompleteCache[keys.DispatchCacheKey, any](c.ClusterDispatchCacheConfig.WithRevisionParameters( @@ -372,25 +371,6 @@ func (c *Config) Complete(ctx context.Context) (RunnableServer, error) { closeables.AddWithError(cachingClusterDispatch.Close) } - // Build OTel stats handler options (shared by both gRPC servers) - // Always disable health check tracing to reduce trace volume - statsHandlerOpts := []otelgrpc.Option{ - otelgrpc.WithFilter(filters.Not(filters.HealthCheck())), - } - - dispatchGrpcServer, err := c.DispatchServer.Complete(zerolog.InfoLevel, - func(server *grpc.Server) { - dispatchSvc.RegisterGrpcServices(server, cachingClusterDispatch) - }, - grpc.ChainUnaryInterceptor(c.DispatchUnaryMiddleware...), - grpc.ChainStreamInterceptor(c.DispatchStreamingMiddleware...), - grpc.StatsHandler(otelgrpc.NewServerHandler(statsHandlerOpts...)), - ) - if err != nil { - return nil, fmt.Errorf("failed to create dispatch gRPC server: %w", err) - } - closeables.AddWithoutError(dispatchGrpcServer.GracefulStop) - datastoreFeatures, err := ds.Features(ctx) if err != nil { return nil, fmt.Errorf("error determining datastore features: %w", err) @@ -432,21 +412,37 @@ func (c *Config) Complete(ctx context.Context) (RunnableServer, error) { return nil, fmt.Errorf("unknown mismatched zedtoken behavior: %s", c.MismatchZedTokenBehavior) } + apiMemoryProtectionConfig, dispatchMemoryProtectionConfig, sharedMemorySampler, err := c.buildMemoryProtectionConfigs(&closeables) + if err != nil { + return nil, err + } + opts := MiddlewareOption{ - log.Logger, - c.GRPCAuthFunc, - !c.DisableVersionResponse, - dispatcher, - c.EnableRequestLogs, - c.EnableResponseLogs, - c.DisableGRPCLatencyHistogram, - serverName, - mismatchZedTokenOption, - nil, - nil, + Logger: log.Logger, + AuthFunc: c.GRPCAuthFunc, + EnableVersionResponse: !c.DisableVersionResponse, + DispatcherForMiddleware: dispatcher, + EnableRequestLog: c.EnableRequestLogs, + EnableResponseLog: c.EnableResponseLogs, + DisableGRPCHistogram: c.DisableGRPCLatencyHistogram, + MiddlewareServiceLabel: serverName, + MismatchingZedTokenOption: mismatchZedTokenOption, + MemoryProtectionConfig: apiMemoryProtectionConfig, + MemorySampler: sharedMemorySampler, } opts = opts.WithDatastore(ds) + // Build OTel stats handler options (shared by both gRPC servers) + // Always disable health check tracing to reduce trace volume + statsHandlerOpts := []otelgrpc.Option{ + otelgrpc.WithFilter(filters.Not(filters.HealthCheck())), + } + + dispatchGrpcServer, err := c.buildDispatchServer(dispatchMemoryProtectionConfig, sharedMemorySampler, ds, cachingClusterDispatch, &closeables, statsHandlerOpts) + if err != nil { + return nil, err + } + defaultUnaryMiddlewareChain, err := DefaultUnaryMiddleware(opts) if err != nil { return nil, fmt.Errorf("error building default middlewares: %w", err) @@ -587,6 +583,51 @@ func (c *Config) Complete(ctx context.Context) (RunnableServer, error) { }, nil } +// buildMemoryProtectionConfigs returns the API config and the Dispatch API config. +func (c *Config) buildMemoryProtectionConfigs(closeables *closeableStack) (memoryprotection.Config, memoryprotection.Config, memoryprotection.MemorySampler, error) { + var apiConfig, dispatchConfig memoryprotection.Config // default is zero value + var sampler memoryprotection.MemorySampler + sampler = memoryprotection.NoOpSampler{} + if c.MemoryProtectionNormalAPIThresholdPercent > 1 || c.MemoryProtectionDispatchAPIThresholdPercent > 1 { + return apiConfig, dispatchConfig, sampler, errors.New("invalid memory protection configuration") + } + if c.MemoryProtectionEnabled { + apiConfig = memoryprotection.Config{ + ThresholdPercent: c.MemoryProtectionNormalAPIThresholdPercent, + } + dispatchConfig = memoryprotection.Config{ + ThresholdPercent: c.MemoryProtectionDispatchAPIThresholdPercent, + } + sampler = memoryprotection.NewMemorySamplerOnInterval(c.MemoryProtectionSampleIntervalSeconds, &memoryprotection.DefaultMemoryLimitProvider{}) + closeables.AddWithoutError(sampler.Close) + } + return apiConfig, dispatchConfig, sampler, nil +} + +func (c *Config) buildDispatchServer(m memoryprotection.Config, sampler memoryprotection.MemorySampler, ds datastore.Datastore, cachingClusterDispatch dispatch.Dispatcher, closeables *closeableStack, otelOpts []otelgrpc.Option) (util.RunnableGRPCServer, error) { + if len(c.DispatchUnaryMiddleware) == 0 && len(c.DispatchStreamingMiddleware) == 0 { + if c.GRPCAuthFunc == nil { + c.DispatchUnaryMiddleware, c.DispatchStreamingMiddleware = DefaultDispatchMiddleware(log.Logger, auth.MustRequirePresharedKey(c.PresharedSecureKey), ds, c.DisableGRPCLatencyHistogram, m, sampler) + } else { + c.DispatchUnaryMiddleware, c.DispatchStreamingMiddleware = DefaultDispatchMiddleware(log.Logger, c.GRPCAuthFunc, ds, c.DisableGRPCLatencyHistogram, m, sampler) + } + } + + dispatchGrpcServer, err := c.DispatchServer.Complete(zerolog.InfoLevel, + func(server *grpc.Server) { + dispatchSvc.RegisterGrpcServices(server, cachingClusterDispatch) + }, + grpc.ChainUnaryInterceptor(c.DispatchUnaryMiddleware...), + grpc.ChainStreamInterceptor(c.DispatchStreamingMiddleware...), + grpc.StatsHandler(otelgrpc.NewServerHandler(otelOpts...)), + ) + if err != nil { + return nil, fmt.Errorf("failed to create dispatch gRPC server: %w", err) + } + closeables.AddWithoutError(dispatchGrpcServer.GracefulStop) + return dispatchGrpcServer, nil +} + func (c *Config) supportOldAndNewReadReplicaConnectionPoolFlags() { defaultReadConnPoolCfg := *datastorecfg.DefaultReadConnPool() if c.DatastoreConfig.ReadReplicaConnPool.MaxOpenConns == defaultReadConnPoolCfg.MaxOpenConns && c.DatastoreConfig. diff --git a/pkg/cmd/server/server_test.go b/pkg/cmd/server/server_test.go index c20652898..6914e2aed 100644 --- a/pkg/cmd/server/server_test.go +++ b/pkg/cmd/server/server_test.go @@ -13,6 +13,7 @@ import ( "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" "go.uber.org/goleak" + "go.uber.org/mock/gomock" "google.golang.org/grpc" "google.golang.org/grpc/codes" healthpb "google.golang.org/grpc/health/grpc_health_v1" @@ -24,6 +25,8 @@ import ( "github.com/authzed/spicedb/internal/datastore/dsfortesting" "github.com/authzed/spicedb/internal/logging" + "github.com/authzed/spicedb/internal/middleware/memoryprotection" + "github.com/authzed/spicedb/internal/mocks" "github.com/authzed/spicedb/pkg/cmd/datastore" "github.com/authzed/spicedb/pkg/cmd/util" "github.com/authzed/spicedb/pkg/middleware/consistency" @@ -51,6 +54,7 @@ func TestServerGracefulTermination(t *testing.T) { WithClusterDispatchCacheConfig(CacheConfig{Enabled: true}), WithHTTPGateway(util.HTTPServerConfig{HTTPEnabled: true, HTTPAddress: ":"}), WithMetricsAPI(util.HTTPServerConfig{HTTPEnabled: true, HTTPAddress: ":"}), + WithMemoryProtectionEnabled(false), ) rs, err := c.Complete(ctx) require.NoError(t, err) @@ -102,6 +106,7 @@ func TestOTelReporting(t *testing.T) { WithNamespaceCacheConfig(CacheConfig{Enabled: false, Metrics: false}), WithClusterDispatchCacheConfig(CacheConfig{Enabled: false, Metrics: false}), WithDatastore(ds), + WithMemoryProtectionEnabled(false), } srv, err := NewConfigWithOptionsAndDefaults(configOpts...).Complete(ctx) @@ -294,6 +299,7 @@ func TestRetryPolicy(t *testing.T) { WithNamespaceCacheConfig(CacheConfig{Enabled: false, Metrics: false}), WithClusterDispatchCacheConfig(CacheConfig{Enabled: false, Metrics: false}), WithDatastore(ds), + WithMemoryProtectionEnabled(false), SetUnaryMiddlewareModification([]MiddlewareModification[grpc.UnaryServerInterceptor]{ { Operation: OperationAppend, @@ -380,7 +386,7 @@ func TestServerGracefulTerminationOnError(t *testing.T) { GRPCServer: util.GRPCServerConfig{ Network: util.BufferedNetwork, }, - }, WithPresharedSecureKey("psk"), WithDatastore(ds)) + }, WithPresharedSecureKey("psk"), WithDatastore(ds), WithMemoryProtectionEnabled(false)) cancel() _, err = c.Complete(ctx) require.NoError(t, err) @@ -440,7 +446,7 @@ func TestModifyUnaryMiddleware(t *testing.T) { }, }} - opt := MiddlewareOption{logging.Logger, nil, false, nil, false, false, false, "testing", consistency.TreatMismatchingTokensAsFullConsistency, nil, nil} + opt := MiddlewareOption{logging.Logger, nil, false, nil, false, false, false, "testing", consistency.TreatMismatchingTokensAsFullConsistency, memoryprotection.Config{ThresholdPercent: 0}, memoryprotection.NewMemorySamplerOnInterval(memoryprotection.DefaultSampleIntervalSeconds, &memoryprotection.DefaultMemoryLimitProvider{}), nil, nil} opt = opt.WithDatastore(nil) defaultMw, err := DefaultUnaryMiddleware(opt) @@ -468,7 +474,7 @@ func TestModifyStreamingMiddleware(t *testing.T) { }, }} - opt := MiddlewareOption{logging.Logger, nil, false, nil, false, false, false, "testing", consistency.TreatMismatchingTokensAsFullConsistency, nil, nil} + opt := MiddlewareOption{logging.Logger, nil, false, nil, false, false, false, "testing", consistency.TreatMismatchingTokensAsFullConsistency, memoryprotection.Config{ThresholdPercent: 0}, memoryprotection.NewMemorySamplerOnInterval(memoryprotection.DefaultSampleIntervalSeconds, &memoryprotection.DefaultMemoryLimitProvider{}), nil, nil} opt = opt.WithDatastore(nil) defaultMw, err := DefaultStreamingMiddleware(opt) @@ -578,3 +584,139 @@ func TestSupportOldAndNewReadReplicaConnectionPoolFlags(t *testing.T) { }) } } + +func TestBuildMemoryProtectionConfig(t *testing.T) { + testcases := map[string]struct { + config *Config + expectedAPIConfig memoryprotection.Config + expectedDispatchConfig memoryprotection.Config + expectedSamplerInterval int + expectAddToCloseables bool + expectedErr string + }{ + `disabled`: { + config: &Config{ + MemoryProtectionEnabled: false, + }, + expectedAPIConfig: memoryprotection.Config{ + ThresholdPercent: 0, + }, + expectedDispatchConfig: memoryprotection.Config{ + ThresholdPercent: 0, + }, + expectedSamplerInterval: 0, + expectAddToCloseables: false, + }, + `enabled`: { + config: &Config{ + MemoryProtectionEnabled: true, + MemoryProtectionSampleIntervalSeconds: 10, + MemoryProtectionNormalAPIThresholdPercent: 0.80, + MemoryProtectionDispatchAPIThresholdPercent: 0.70, + }, + expectedAPIConfig: memoryprotection.Config{ + ThresholdPercent: 0.80, + }, + expectedDispatchConfig: memoryprotection.Config{ + ThresholdPercent: 0.70, + }, + expectedSamplerInterval: 10, + expectAddToCloseables: true, + }, + `err_invalid_above_one_1`: { + config: &Config{ + MemoryProtectionEnabled: true, + MemoryProtectionSampleIntervalSeconds: 10, + MemoryProtectionNormalAPIThresholdPercent: 1.1, + MemoryProtectionDispatchAPIThresholdPercent: 0.70, + }, + expectedSamplerInterval: 0, + expectedErr: "invalid memory protection configuration", + expectAddToCloseables: false, + }, + `err_invalid_above_one_2`: { + config: &Config{ + MemoryProtectionEnabled: true, + MemoryProtectionSampleIntervalSeconds: 10, + MemoryProtectionNormalAPIThresholdPercent: 0.70, + MemoryProtectionDispatchAPIThresholdPercent: 1.1, + }, + expectedSamplerInterval: 0, + expectedErr: "invalid memory protection configuration", + expectAddToCloseables: false, + }, + } + + for name, tc := range testcases { + t.Run(name, func(t *testing.T) { + closeables := closeableStack{} + actualAPIConfig, actualDispatchConfig, sampler, err := tc.config.buildMemoryProtectionConfigs(&closeables) + if tc.expectedErr != "" { + require.ErrorContains(t, err, tc.expectedErr) + return + } + require.NoError(t, err) + require.InDelta(t, tc.expectedAPIConfig.ThresholdPercent, actualAPIConfig.ThresholdPercent, 0.01) + require.InDelta(t, tc.expectedDispatchConfig.ThresholdPercent, actualDispatchConfig.ThresholdPercent, 0.01) + require.NotNil(t, sampler) + require.Equal(t, tc.expectedSamplerInterval, sampler.GetIntervalSeconds()) + if tc.expectAddToCloseables { + require.Len(t, closeables.closers, 1) + } else { + require.Len(t, closeables.closers, 0) + } + }) + } +} + +func TestBuildDispatchServer(t *testing.T) { + testcases := map[string]struct { + config *Config + expectedDispatchUnnaryMiddleware int + expectedDispatchStreamingMiddleware int + }{ + `auth:preshared key`: { + config: &Config{ + PresharedSecureKey: []string{"securekey"}, + }, + expectedDispatchUnnaryMiddleware: 8, + expectedDispatchStreamingMiddleware: 6, + }, + `auth:custom`: { + config: &Config{ + GRPCAuthFunc: func(ctx context.Context) (context.Context, error) { + return ctx, nil + }, + }, + expectedDispatchUnnaryMiddleware: 8, + expectedDispatchStreamingMiddleware: 6, + }, + } + + for name, tc := range testcases { + t.Run(name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockDatastore := mocks.NewMockDatastore(ctrl) + mockDispatcher := mocks.NewMockDispatcher(ctrl) + + mp := memoryprotection.DefaultDispatchConfig() + + closeables := closeableStack{} + t.Cleanup(func() { + _ = closeables.Close() + }) + + sampler := memoryprotection.NewMemorySamplerOnInterval(memoryprotection.DefaultSampleIntervalSeconds, &memoryprotection.DefaultMemoryLimitProvider{}) + t.Cleanup(sampler.Close) + + srv, err := tc.config.buildDispatchServer(mp, sampler, mockDatastore, mockDispatcher, &closeables, nil) + require.NoError(t, err) + require.NotNil(t, srv) + require.Len(t, closeables.closers, 1) + require.Len(t, tc.config.DispatchUnaryMiddleware, tc.expectedDispatchUnnaryMiddleware) + require.Len(t, tc.config.DispatchStreamingMiddleware, tc.expectedDispatchStreamingMiddleware) + }) + } +} diff --git a/pkg/cmd/server/zz_generated.middlewareoption.go b/pkg/cmd/server/zz_generated.middlewareoption.go index ed1dd3430..2ea23cbe4 100644 --- a/pkg/cmd/server/zz_generated.middlewareoption.go +++ b/pkg/cmd/server/zz_generated.middlewareoption.go @@ -3,6 +3,7 @@ package server import ( dispatch "github.com/authzed/spicedb/internal/dispatch" + memoryprotection "github.com/authzed/spicedb/internal/middleware/memoryprotection" consistency "github.com/authzed/spicedb/pkg/middleware/consistency" defaults "github.com/creasty/defaults" helpers "github.com/ecordell/optgen/helpers" @@ -43,6 +44,8 @@ func (m *MiddlewareOption) ToOption() MiddlewareOptionOption { to.DisableGRPCHistogram = m.DisableGRPCHistogram to.MiddlewareServiceLabel = m.MiddlewareServiceLabel to.MismatchingZedTokenOption = m.MismatchingZedTokenOption + to.MemoryProtectionConfig = m.MemoryProtectionConfig + to.MemorySampler = m.MemorySampler to.unaryDatastoreMiddleware = m.unaryDatastoreMiddleware to.streamDatastoreMiddleware = m.streamDatastoreMiddleware } @@ -57,6 +60,7 @@ func (m MiddlewareOption) DebugMap() map[string]any { debugMap["DisableGRPCHistogram"] = helpers.DebugValue(m.DisableGRPCHistogram, false) debugMap["MiddlewareServiceLabel"] = helpers.DebugValue(m.MiddlewareServiceLabel, false) debugMap["MismatchingZedTokenOption"] = helpers.DebugValue(m.MismatchingZedTokenOption, false) + debugMap["MemoryProtectionConfig"] = helpers.DebugValue(m.MemoryProtectionConfig, false) return debugMap } @@ -138,3 +142,17 @@ func WithMismatchingZedTokenOption(mismatchingZedTokenOption consistency.Mismatc m.MismatchingZedTokenOption = mismatchingZedTokenOption } } + +// WithMemoryProtectionConfig returns an option that can set MemoryProtectionConfig on a MiddlewareOption +func WithMemoryProtectionConfig(memoryProtectionConfig memoryprotection.Config) MiddlewareOptionOption { + return func(m *MiddlewareOption) { + m.MemoryProtectionConfig = memoryProtectionConfig + } +} + +// WithMemorySampler returns an option that can set MemorySampler on a MiddlewareOption +func WithMemorySampler(memorySampler memoryprotection.MemorySampler) MiddlewareOptionOption { + return func(m *MiddlewareOption) { + m.MemorySampler = memorySampler + } +} diff --git a/pkg/cmd/server/zz_generated.options.go b/pkg/cmd/server/zz_generated.options.go index fd2c777a1..e8abce9c9 100644 --- a/pkg/cmd/server/zz_generated.options.go +++ b/pkg/cmd/server/zz_generated.options.go @@ -100,6 +100,10 @@ func (c *Config) ToOption() ConfigOption { to.MetricsAPI = c.MetricsAPI to.UnaryMiddlewareModification = c.UnaryMiddlewareModification to.StreamingMiddlewareModification = c.StreamingMiddlewareModification + to.MemoryProtectionEnabled = c.MemoryProtectionEnabled + to.MemoryProtectionNormalAPIThresholdPercent = c.MemoryProtectionNormalAPIThresholdPercent + to.MemoryProtectionDispatchAPIThresholdPercent = c.MemoryProtectionDispatchAPIThresholdPercent + to.MemoryProtectionSampleIntervalSeconds = c.MemoryProtectionSampleIntervalSeconds to.DispatchUnaryMiddleware = c.DispatchUnaryMiddleware to.DispatchStreamingMiddleware = c.DispatchStreamingMiddleware to.SilentlyDisableTelemetry = c.SilentlyDisableTelemetry @@ -174,6 +178,10 @@ func (c Config) DebugMap() map[string]any { debugMap["EnablePerformanceInsightMetrics"] = helpers.DebugValue(c.EnablePerformanceInsightMetrics, false) debugMap["MismatchZedTokenBehavior"] = helpers.DebugValue(c.MismatchZedTokenBehavior, false) debugMap["MetricsAPI"] = helpers.DebugValue(c.MetricsAPI, false) + debugMap["MemoryProtectionEnabled"] = helpers.DebugValue(c.MemoryProtectionEnabled, false) + debugMap["MemoryProtectionNormalAPIThresholdPercent"] = helpers.DebugValue(c.MemoryProtectionNormalAPIThresholdPercent, false) + debugMap["MemoryProtectionDispatchAPIThresholdPercent"] = helpers.DebugValue(c.MemoryProtectionDispatchAPIThresholdPercent, false) + debugMap["MemoryProtectionSampleIntervalSeconds"] = helpers.DebugValue(c.MemoryProtectionSampleIntervalSeconds, false) debugMap["SilentlyDisableTelemetry"] = helpers.DebugValue(c.SilentlyDisableTelemetry, false) debugMap["TelemetryCAOverridePath"] = helpers.DebugValue(c.TelemetryCAOverridePath, false) debugMap["TelemetryEndpoint"] = helpers.DebugValue(c.TelemetryEndpoint, false) @@ -683,6 +691,34 @@ func SetStreamingMiddlewareModification(streamingMiddlewareModification []Middle } } +// WithMemoryProtectionEnabled returns an option that can set MemoryProtectionEnabled on a Config +func WithMemoryProtectionEnabled(memoryProtectionEnabled bool) ConfigOption { + return func(c *Config) { + c.MemoryProtectionEnabled = memoryProtectionEnabled + } +} + +// WithMemoryProtectionNormalAPIThresholdPercent returns an option that can set MemoryProtectionNormalAPIThresholdPercent on a Config +func WithMemoryProtectionNormalAPIThresholdPercent(memoryProtectionNormalAPIThresholdPercent float64) ConfigOption { + return func(c *Config) { + c.MemoryProtectionNormalAPIThresholdPercent = memoryProtectionNormalAPIThresholdPercent + } +} + +// WithMemoryProtectionDispatchAPIThresholdPercent returns an option that can set MemoryProtectionDispatchAPIThresholdPercent on a Config +func WithMemoryProtectionDispatchAPIThresholdPercent(memoryProtectionDispatchAPIThresholdPercent float64) ConfigOption { + return func(c *Config) { + c.MemoryProtectionDispatchAPIThresholdPercent = memoryProtectionDispatchAPIThresholdPercent + } +} + +// WithMemoryProtectionSampleIntervalSeconds returns an option that can set MemoryProtectionSampleIntervalSeconds on a Config +func WithMemoryProtectionSampleIntervalSeconds(memoryProtectionSampleIntervalSeconds int) ConfigOption { + return func(c *Config) { + c.MemoryProtectionSampleIntervalSeconds = memoryProtectionSampleIntervalSeconds + } +} + // WithDispatchUnaryMiddleware returns an option that can append DispatchUnaryMiddlewares to Config.DispatchUnaryMiddleware func WithDispatchUnaryMiddleware(dispatchUnaryMiddleware grpc.UnaryServerInterceptor) ConfigOption { return func(c *Config) { diff --git a/pkg/datastore/datastore.go b/pkg/datastore/datastore.go index 5f5a1baf9..cceb55955 100644 --- a/pkg/datastore/datastore.go +++ b/pkg/datastore/datastore.go @@ -1,3 +1,5 @@ +//go:generate go run go.uber.org/mock/mockgen -source datastore.go -destination ../../internal/mocks/mock_datastore.go -package mocks Datastore + package datastore import (