feat: add Memory Protection Middleware #2646

ecordell · 2025-10-23T17:19:05Z

Description

This introduces a new memory protection middleware to help prevent out-of-memory conditions in SpiceDB by implementing admission control based on current memory usage.

This is not a perfect solution (doesn't prevent non-traffic-related sources of OOM) and is meant to support other future improvements to resource sharing in a single SpiceDB node.

The middleware is installed both in the main api and in dispatch, but at different thresholds. Memory usage is polled in the background, and if in-flight memory rises above the threshold, backpressure is placed on incoming requests.

The dispatch threshold is higher than the API threshold to preserve already admitted traffic as much as possible.

Testing

Unit tests included
Manual E2E test:
- Modify docker-compose.yaml to set mem_limit: "300mb" to the spicedb containers
- Run docker-compose up --build
- Run this

zed context set example localhost:50051 foobar --insecure
zed import development/schema.yaml
{
    echo '{"items":['
    for i in $(seq 1 200); do
      d=$(( (RANDOM % 9999) + 1 ))
      echo -n "{\"resource\":{\"objectTy  pe\":\"document\",\"objectId\": \"${d}\"}, \"permission\":\"view\",\"subject\":{ \"object\": {\"objectType\": \"user\", \"objectId\": \"1\"}}}"
      [ $i -lt 200 ] && echo -n ","
    done
    echo "], \"with_tracing\": true}"
} > payload.json
ab -n 100000 -c 200 -T 'application/json' -H 'Authorization: Bearer foobar' -p payload.json http://localhost:8443/v1/permissions/checkbulk

you should see logs such as:

{
  "level": "warn",
  "traceID": "125b7b37c5775af1f2d9ebf253dcf3d1",
  "protocol": "grpc",
  "grpc.component": "server",
  "grpc.service": "authzed.api.v1.PermissionsService",
  "grpc.method": "CheckBulkPermissions",
  "grpc.method_type": "unary",
  "requestID": "d3vurmuoqu8s73bsdom0",
  "peer.address": "127.0.0.1:60376",
  "grpc.start_time": "2025-10-27T22:10:35Z",
  "grpc.code": "ResourceExhausted",
  "grpc.error": "rpc error: code = ResourceExhausted desc = server rejected the request because memory usage is above configured threshold",
  "grpc.time_ms": 0,
  "time": "2025-10-27T22:10:35Z",
  "message": "finished call"
}

and this graph in Grafana:

codecov · 2025-10-23T17:24:51Z

Codecov Report

❌ Patch coverage is 11.97605% with 1470 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.08%. Comparing base (e55404b) to head (8b47b79).
⚠️ Report is 11 commits behind head on main.

Files with missing lines	Patch %	Lines
internal/mocks/mock_datastore.go	0.32%	1277 Missing ⚠️
internal/mocks/mock_dispatcher.go	2.57%	152 Missing ⚠️
...ware/memoryprotection/mocks/mock_memory_sampler.go	38.89%	22 Missing ⚠️
...rnal/middleware/memoryprotection/memory_sampler.go	83.08%	9 Missing and 2 partials ⚠️
pkg/cmd/server/server.go	87.24%	3 Missing and 3 partials ⚠️
...l/middleware/memoryprotection/memory_protection.go	96.50%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2646      +/-   ##
==========================================
- Coverage   79.46%   77.08%   -2.37%     
==========================================
  Files         455      461       +6     
  Lines       47161    48940    +1779     
==========================================
+ Hits        37470    37722     +252     
- Misses       6945     8449    +1504     
- Partials     2746     2769      +23

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

miparnisari · 2025-10-27T23:51:57Z

development/prometheus.yaml

  - job_name: "spicedb"
    static_configs:
-      - targets: ["spicedb:9090"]
+      - targets: ["spicedb-1:9090"]


FYI so that we can verify the new metrics in Grafana

miparnisari · 2025-10-28T16:52:35Z

internal/mocks/mock_datastore.go

FYI i can move the mock generation to a different PR

This is fine.

miparnisari · 2025-10-28T16:53:45Z

.gitattributes

@@ -0,0 +1,4 @@
+internal/mocks/*.go linguist-generated=true


FYI so that when changing these, the "View PR" view on Github says that the files are automatically generated and people can just mark "viewed" on them

internal/middleware/memoryprotection/memoryprotection.go

miparnisari · 2025-10-28T18:30:18Z

pkg/cmd/serve.go


+	// Memory Protection flags
+	apiFlags.BoolVar(&config.MemoryProtectionEnabled, "memory-protection-enabled", true, "enables a memory-based middleware that rejects requests when memory usage is too high")
+	apiFlags.IntVar(&config.MemoryProtectionAPIThresholdPercent, "memory-protection-api-threshold", 90, "memory usage threshold percentage for regular API requests (0-100)")


i don't like that we have the default 90 here and also in the struct itself 😭 i'd love a unified approach in the future

pkg/cmd/serve.go

pkg/cmd/server/defaults.go

pkg/cmd/server/server_test.go

internal/middleware/memoryprotection/memoryprotection.go

vroldanbet · 2025-10-28T22:22:43Z

pkg/cmd/serve.go


+	// Memory Protection flags
+	apiFlags.BoolVar(&config.MemoryProtectionEnabled, "memory-protection-enabled", true, "enables a memory-based middleware that rejects requests when memory usage is too high")
+	apiFlags.IntVar(&config.MemoryProtectionAPIThresholdPercent, "memory-protection-api-threshold", 90, "memory usage threshold percentage for regular API requests (0-100)")


We have some percent-based flags that are defined as [0...1] floats. Worth having a look and deciding which approach to commit to.

Made it similar to this

--datastore-revision-quantization-max-staleness-percent float float percentage (where 1 = 100%)

i don't love it but I prefer being consistent

Funny, I just found that you can pass a large number to that flag and the server accepts it 😆

- "SPICEDB_DATASTORE_REVISION_QUANTIZATION_MAX_STALENESS_PERCENT=10000"

seems like it's missing some validation :D

internal/middleware/memoryprotection/memoryprotection.go

The commit introduces a new memory protection middleware to help prevent out-of-memory conditions in SpiceDB by implementing admission control based on current memory usage. This is not a perfect solution (doesn't prevent non-traffic-related sources of OOM) and is meant to support other future improvements to resource sharing in a single SpiceDB node. The middleware is installed both in the main api and in dispatch, but at different thresholds. Memory usage is polled in the background, and if in-flight memory rises above the threshold, backpressure is placed on incoming requests. The API threshold is higher than the dispatch threshold to preserve already admitted traffic as much as possible.

tstirrat15

See comments

tstirrat15 · 2025-10-28T16:52:03Z

internal/middleware/memoryprotection/memoryprotection.go

+	_ MemoryLimitProvider = (*DefaultMemoryLimitProvider)(nil)
+	_ MemoryLimitProvider = (*HardCodedMemoryLimitProvider)(nil)


Why this syntax as opposed to constructing one?

WDYM? This is standard go code to assert that a struct satisfies an interface

internal/middleware/memoryprotection/memoryprotection.go

tstirrat15 · 2025-10-28T20:50:09Z

internal/mocks/mock_datastore.go

This is fine.

pkg/cmd/server/server.go

tstirrat15 · 2025-10-31T17:26:47Z

internal/middleware/memoryprotection/memory_sampler_test.go

+				synctest.Wait()
+				t.Log("When we get here, the sampling is guaranteed to have run")
+
+				require.True(t, sampler.GetTimestampLastMemorySample().After(now))


This is in a goroutine :P

ecordell · 2025-11-10T14:46:57Z

We had a discussion on slack where we looked at some implementation details of runtime/metrics. Although it's not stop-the-world like ReadMemStats is, it does acquire a metrics lock (metricsSema) to read the values. Joey found https://github.com/odigos-io/go-rtml/ which looked promising as a way to get these values faster, potentially even on every request (i.e. without the extra goroutines polling in the background).

vroldanbet

LGTM, I defer to y'all looking into an alternative implementation to retrieve runtime stats with lower overhead

miparnisari · 2025-11-11T02:40:29Z

@ecordell @vroldanbet please see #2691

github-actions bot added area/cli Affects the command line area/dependencies Affects dependencies area/tooling Affects the dev or user toolchain (e.g. tests, ci, build tools) labels Oct 23, 2025

ecordell force-pushed the oomprotect branch from 59a8715 to 31d1947 Compare October 23, 2025 18:15

miparnisari force-pushed the oomprotect branch 6 times, most recently from d65a7c0 to 7afdbc9 Compare October 27, 2025 23:51

miparnisari reviewed Oct 27, 2025

View reviewed changes

github-actions bot added the area/dispatch Affects dispatching of requests label Oct 27, 2025

miparnisari force-pushed the oomprotect branch 5 times, most recently from 7d7fb92 to dd062d7 Compare October 28, 2025 01:06

miparnisari marked this pull request as ready for review October 28, 2025 05:01

miparnisari requested a review from a team as a code owner October 28, 2025 05:01

miparnisari reviewed Oct 28, 2025

View reviewed changes

miparnisari force-pushed the oomprotect branch from dd062d7 to 9ff3096 Compare October 28, 2025 18:26

miparnisari reviewed Oct 28, 2025

View reviewed changes

internal/middleware/memoryprotection/memoryprotection.go Outdated Show resolved Hide resolved

miparnisari reviewed Oct 28, 2025

View reviewed changes

vroldanbet reviewed Oct 28, 2025

View reviewed changes

miparnisari force-pushed the oomprotect branch 4 times, most recently from 77c4458 to e73f76c Compare October 31, 2025 03:56

ecordell and others added 2 commits October 31, 2025 09:22

chore: improvements and fix building of dispatch middleware

4a317d3

miparnisari force-pushed the oomprotect branch from e73f76c to 4a317d3 Compare October 31, 2025 16:31

tstirrat15 reviewed Oct 31, 2025

View reviewed changes

vroldanbet previously approved these changes Nov 10, 2025

View reviewed changes

miparnisari added this pull request to the merge queue Nov 10, 2025

miparnisari removed this pull request from the merge queue due to a manual request Nov 10, 2025

miparnisari dismissed vroldanbet’s stale review via 44c3557 November 10, 2025 22:48

miparnisari force-pushed the oomprotect branch 2 times, most recently from 772f37b to 24b9025 Compare November 10, 2025 23:35

fix percentage (0-100) should be (0-1). also amend grafana dashboards

8b47b79

miparnisari force-pushed the oomprotect branch from 24b9025 to 8b47b79 Compare November 11, 2025 00:32

miparnisari mentioned this pull request Nov 11, 2025

feat: add Memory Protection Middleware #2691

Merged

miparnisari marked this pull request as draft November 11, 2025 02:40

miparnisari closed this Nov 11, 2025

github-actions bot locked and limited conversation to collaborators Nov 11, 2025

		_ MemoryLimitProvider = (*DefaultMemoryLimitProvider)(nil)
		_ MemoryLimitProvider = (*HardCodedMemoryLimitProvider)(nil)

feat: add Memory Protection Middleware #2646

feat: add Memory Protection Middleware #2646

Conversation

ecordell commented Oct 23, 2025 • edited by miparnisari Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Testing

Uh oh!

codecov bot commented Oct 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

miparnisari Oct 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tstirrat15 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ecordell commented Nov 10, 2025

Uh oh!

vroldanbet left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

miparnisari commented Nov 11, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ecordell commented Oct 23, 2025 •

edited by miparnisari

Loading

codecov bot commented Oct 23, 2025 •

edited

Loading

miparnisari Oct 28, 2025 •

edited

Loading