feat: add Memory Protection Middleware #2691

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

Pre-reqs

chore: support ldflags in go-build actions#105

Description

This PR 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 APIs, and it uses this lib: https://github.com/odigos-io/go-rtml

⚠️ It does NOT reject health check requests based on memory usage,
⚠️ It needs a special build flag.

References

Alternative implementation: #2646.

Upon testing both implementations, it was found that:

this PR rejects less requests.
this PR didn't OOM the docker container when mem_limit was hit.

codecov · 2025-11-11T02:32:50Z

Codecov Report

❌ Patch coverage is 93.06931% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.01%. Comparing base (ec7032a) to head (967b1b9).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
pkg/cmd/server/server.go	88.24%	2 Missing and 2 partials ⚠️
pkg/cmd/testserver/testserver.go	0.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2691      +/-   ##
==========================================
- Coverage   77.01%   77.01%   -0.00%     
==========================================
  Files         462      464       +2     
  Lines       49134    49184      +50     
==========================================
+ Hits        37837    37872      +35     
- Misses       8509     8520      +11     
- Partials     2788     2792       +4

☔ 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.

internal/middleware/memoryprotection/memory_protection_test.go

ecordell · 2025-11-12T15:14:23Z

pkg/cmd/server/server.go

+	if c.MemoryProtectionEnabled {
+		return memoryprotection.NewRealTimeMemoryUsageProvider()
+	}
+	return &memoryprotection.HarcodedMemoryLimitProvider{AcceptAllRequests: true}


I would've expected that if the feature was disabled we would just not install the middleware at all, instead of installing it with a no-op limiter?

I'm wondering now why do we have the MemoryProtectionEnabled flag at all 🤔

I'd be in favor of having it be always-on, and then we just don't install the middleware in tests (if that's possible).

I looked at our existing middlewares to see if we have an example of one that is conditionally installed, and it doesn't exist... if a middleware isn't enabled, it still runs, but it just no-ops. Example:

spicedb/internal/middleware/perfinsights/perfinsights.go

Lines 157 to 160 in d4ff660

func (r *reporter) ServerReporter(ctx context.Context, callMeta interceptors.CallMeta) (interceptors.Reporter, context.Context) {

if !r.isEnabled {

return &interceptors.NoopReporter{}, ctx

}

Captured in #2723

And I think I'm going to split the difference and not expose it as a user-facing flag, but leave the ability to configure it. My thought is that if (e.g.) our steelthread tests start hitting OOM conditions that trigger this middleware, it'll be good to disable the middleware so that we get failures instead of ResourceExhausteds

pkg/cmd/server/server_test.go

ecordell · 2025-11-12T15:21:31Z

internal/middleware/memoryprotection/memory_usage_provider.go

+type GoRealTimeMemoryLimiter struct{}
+
+func (g GoRealTimeMemoryLimiter) IsMemLimitReached() bool {
+	return rtml.IsMemLimitReached()


I wish we had a way to set a threshold here. Since this will get called both for dispatch and primary api calls, it will start rejecting them both equally, when really we want to reject primary api before we start rejecting dispatch.

I don't think this is worth blocking on, but I do wonder if we should contribute a threshold % to the upstream library.

Good idea. Opened odigos-io/go-rtml#3

vroldanbet · 2025-11-12T10:06:06Z

magefiles/util.go

 		"-failfast",
 		"-count=1",
 		"-race",
+		"-timeout=20m",


why increasing from 10m to 20m?

Ah sorry, that was a leftover from a different PR. The answer is https://github.com/authzed/spicedb/pull/2674/files#r2505759338

Tl;dr instead of putting timeouts in the test code itself, i think it's better to put it in go test -timeout

vroldanbet · 2025-11-12T10:15:00Z

pkg/cmd/server/defaults.go

 			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


Why does it have to run before auth?

The rationale for EnsureAlreadyExecuted is to declare a dependency between middlewares explicitly: "I need this to run before me because I depend on its side effects". So why does the auth middleware depend on memory protection?

If not, you can remove it.

If anything, I'd argue that memory protection depends on all the previous middlewares, because we want observability to take place before we kick memory protection, otherwise we would be running blind when ResourceExhausted is returned.

Oh this was a leftover from @ecordell's original PR. I don't have a strong opinion around this, I didn't know that EnsureAlreadyExecuted meant an explicit dependency.

@ecordell do you have any thoughts?

I don't have a strong opinion here, I just thought that the earlier the memory protection could run the better.

pkg/cmd/server/defaults.go

tstirrat15

See comments, otherwise looks good.

internal/middleware/memoryprotection/memory_protection.go

tstirrat15 · 2025-11-12T20:00:03Z

internal/middleware/memoryprotection/memory_protection.go

+	acceptedStr := strconv.FormatBool(accepted)
+
+	RequestsProcessed.WithLabelValues(endpointType, acceptedStr).Inc()


Ah, we're doing it with a label? Hmm...

tstirrat15

I'd be keen to see what @ecordell has to say, and I'm down with the idea of making this memory middleware non-optional.

tstirrat15 · 2025-11-14T22:04:47Z

pkg/cmd/server/server.go

+	if c.MemoryProtectionEnabled {
+		return memoryprotection.NewRealTimeMemoryUsageProvider()
+	}
+	return &memoryprotection.HarcodedMemoryLimitProvider{AcceptAllRequests: true}


I'd be in favor of having it be always-on, and then we just don't install the middleware in tests (if that's possible).

tstirrat15

This looks like it should be good.

tstirrat15 · 2025-11-21T21:29:41Z

magefiles/build.go

-	return sh.RunWithV(
-		map[string]string{},


No vars, so don't need With

tstirrat15

LGTM

miparnisari changed the title ~~feat: add Memory Protection Middleware~~ feat: add Memory Protection Middleware (alternative implementation) Nov 11, 2025

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) area/dispatch Affects dispatching of requests labels Nov 11, 2025

miparnisari mentioned this pull request Nov 11, 2025

feat: add Memory Protection Middleware #2646

Closed

miparnisari force-pushed the oomprotect-with-rtml branch from 5ec48e4 to 4ddac59 Compare November 11, 2025 02:48

miparnisari changed the title ~~feat: add Memory Protection Middleware (alternative implementation)~~ feat: add Memory Protection Middleware (preferred implementation) Nov 11, 2025

This was referenced Nov 11, 2025

chore: support ldflags in go-build authzed/actions#105

Merged

chore: update grafana dashboard, .gitattributes, github actions #2693

Merged

ecordell reviewed Nov 12, 2025

View reviewed changes

vroldanbet reviewed Nov 12, 2025

View reviewed changes

miparnisari force-pushed the oomprotect-with-rtml branch from 4ddac59 to 269a1c1 Compare November 12, 2025 17:27

tstirrat15 reviewed Nov 12, 2025

View reviewed changes

miparnisari force-pushed the oomprotect-with-rtml branch 13 times, most recently from 112d59b to 491398d Compare November 14, 2025 01:33

miparnisari mentioned this pull request Nov 14, 2025

chore: add docs command to be used by mage #2703

Merged

miparnisari force-pushed the oomprotect-with-rtml branch from 491398d to 6844c39 Compare November 14, 2025 20:24

miparnisari marked this pull request as ready for review November 14, 2025 20:56

miparnisari requested a review from a team as a code owner November 14, 2025 20:56

tstirrat15 reviewed Nov 14, 2025

View reviewed changes

miparnisari changed the title ~~feat: add Memory Protection Middleware (preferred implementation)~~ feat: add Memory Protection Middleware Nov 14, 2025

miparnisari mentioned this pull request Nov 14, 2025

chore: introduce mockgen #2705

Merged

miparnisari force-pushed the oomprotect-with-rtml branch from 10b5284 to 8f5d0b9 Compare November 14, 2025 23:24

github-actions bot removed the area/dispatch Affects dispatching of requests label Nov 14, 2025

miparnisari force-pushed the oomprotect-with-rtml branch from 8f5d0b9 to 3d26690 Compare November 14, 2025 23:25

tstirrat15 force-pushed the oomprotect-with-rtml branch 3 times, most recently from 173b7ce to 27e46d0 Compare November 21, 2025 21:28

tstirrat15 previously approved these changes Nov 21, 2025

View reviewed changes

magefiles/build.go

Comment on lines -16 to -17

return sh.RunWithV(

map[string]string{},

Copy link

Contributor

tstirrat15 Nov 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No vars, so don't need With

tstirrat15 dismissed their stale review via 05c7766 November 21, 2025 21:30

tstirrat15 force-pushed the oomprotect-with-rtml branch from 27e46d0 to 05c7766 Compare November 21, 2025 21:30

tstirrat15 approved these changes Nov 21, 2025

View reviewed changes

ecordell and others added 7 commits November 24, 2025 13:58

feat: add Memory Protection Middleware via RTML

3dc0bab

chore: Fix lint issues

66f17ae

chore: Reorder middlewares

ac52b71

chore: remove user-facing flag

12fbc9a

chore: remove unnecessary env vars

2e895c9

chore: regen docs to remove flag

bb00197

chore: format

967b1b9

tstirrat15 force-pushed the oomprotect-with-rtml branch from 83189c4 to 967b1b9 Compare November 24, 2025 20:59

tstirrat15 added this pull request to the merge queue Nov 25, 2025

github-merge-queue bot removed this pull request from the merge queue due to failed status checks Nov 25, 2025

tstirrat15 added this pull request to the merge queue Nov 25, 2025

Merged via the queue into main with commit e7390fe Nov 25, 2025
45 checks passed

tstirrat15 deleted the oomprotect-with-rtml branch November 25, 2025 17:41

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

	func (r *reporter) ServerReporter(ctx context.Context, callMeta interceptors.CallMeta) (interceptors.Reporter, context.Context) {
	if !r.isEnabled {
	return &interceptors.NoopReporter{}, ctx
	}

		acceptedStr := strconv.FormatBool(accepted)

		RequestsProcessed.WithLabelValues(endpointType, acceptedStr).Inc()

feat: add Memory Protection Middleware #2691

feat: add Memory Protection Middleware #2691

Uh oh!

Conversation

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

Uh oh!

Pre-reqs

Description

References

Uh oh!

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

Uh oh!

Codecov Report

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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!

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!

tstirrat15 left a comment • 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!

Choose a reason for hiding this comment

Uh oh!

tstirrat15 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tstirrat15 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tstirrat15 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

miparnisari commented Nov 11, 2025 •

edited

Loading

codecov bot commented Nov 11, 2025 •

edited

Loading

tstirrat15 left a comment •

edited

Loading