Implement VM_AllocateExternalMemory - #4128
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds external memory accounting in zmalloc, exposes it through new module APIs, reports it in the INFO Debug section, and adds unit and module integration tests. ChangesExternal memory accounting
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ModuleTest
participant ModuleAPI
participant Zmalloc
participant INFO
ModuleTest->>ModuleAPI: set external memory
ModuleAPI->>Zmalloc: update used_memory_external
Zmalloc-->>ModuleAPI: return update status
ModuleAPI-->>ModuleTest: return command result
INFO->>Zmalloc: read external memory usage
Zmalloc-->>INFO: return external usage
INFO-->>ModuleTest: emit used_memory_module_external in Debug
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/unit/test_zmalloc.cpp (1)
70-81: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd overflow/underflow error path tests.
The test covers the happy path well but doesn't validate the overflow/underflow guards in
zmalloc_increase_used_memory_externalandzmalloc_decrease_used_memory_external. Consider adding assertions for the-1return on invalid updates.♻️ Suggested additional test cases
TEST_F(ZmallocTest, TestZmallocExternalUsedMemory) { size_t used_memory_before = zmalloc_used_memory(); size_t external_memory_before = zmalloc_used_external_memory(); ASSERT_EQ(zmalloc_increase_used_memory_external(123), 0); ASSERT_EQ(zmalloc_used_external_memory(), external_memory_before + 123); ASSERT_EQ(zmalloc_used_memory(), used_memory_before + 123); ASSERT_EQ(zmalloc_decrease_used_memory_external(123), 0); ASSERT_EQ(zmalloc_used_external_memory(), external_memory_before); ASSERT_EQ(zmalloc_used_memory(), used_memory_before); + + /* Underflow: decreasing more than allocated should fail. */ + ASSERT_EQ(zmalloc_decrease_used_memory_external(1), -1); + + /* Overflow: increasing beyond SIZE_MAX should fail. */ + ASSERT_EQ(zmalloc_increase_used_memory_external(SIZE_MAX), 0); + ASSERT_EQ(zmalloc_increase_used_memory_external(1), -1); + ASSERT_EQ(zmalloc_decrease_used_memory_external(SIZE_MAX), 0); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/unit/test_zmalloc.cpp` around lines 70 - 81, The Zmalloc external memory test only covers successful updates and misses the guard paths in zmalloc_increase_used_memory_external and zmalloc_decrease_used_memory_external. Extend TestZmallocExternalUsedMemory in test_zmalloc.cpp with negative cases that force overflow/underflow or invalid updates and assert the functions return -1, while also verifying zmalloc_used_memory and zmalloc_used_external_memory remain unchanged after those failures.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/zmalloc.c`:
- Around line 117-118: `used_memory_external` is a shared counter with
concurrent reads and writes, so it needs the same atomic handling used by
`used_memory_for_additional_threads`. Update the declaration in `src/zmalloc.c`
to use an atomic size type and modify the read/write paths in
`zmalloc_used_memory()` and the module callback update logic to use relaxed
atomic load/store or fetch-add operations, keeping `moduleGIL` as the writer
serialization but removing the C data race.
---
Nitpick comments:
In `@src/unit/test_zmalloc.cpp`:
- Around line 70-81: The Zmalloc external memory test only covers successful
updates and misses the guard paths in zmalloc_increase_used_memory_external and
zmalloc_decrease_used_memory_external. Extend TestZmallocExternalUsedMemory in
test_zmalloc.cpp with negative cases that force overflow/underflow or invalid
updates and assert the functions return -1, while also verifying
zmalloc_used_memory and zmalloc_used_external_memory remain unchanged after
those failures.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 094e030e-1ad3-444e-83e2-470588de8fe3
📒 Files selected for processing (6)
src/module.csrc/server.csrc/unit/test_zmalloc.cppsrc/valkeymodule.hsrc/zmalloc.csrc/zmalloc.h
Signed-off-by: Su Ko <rhtn1128@gmail.com>
cf84948 to
eab7d7a
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## unstable #4128 +/- ##
============================================
+ Coverage 76.76% 76.80% +0.03%
============================================
Files 162 162
Lines 81169 81525 +356
============================================
+ Hits 62313 62618 +305
- Misses 18856 18907 +51
🚀 New features to boost your workflow:
|
![valkey-review-bot[bot]](https://avatars.githubusercontent.com/in/3633497?s=60&v=4){kind=small}
| info, | ||
| "# Memory\r\n" FMTARGS( | ||
| "used_memory:%zu\r\n", zmalloc_used, | ||
| "used_memory_module_external:%zu\r\n", module_external_memory, |
There was a problem hiding this comment.
Adding a dedicated used_memory_module_external field here makes the new category visible, but getMemoryOverheadData() still derives mh->dataset from zmalloc_used_memory() and subtracts only the built-in overhead buckets (src/object.c:1403, src/object.c:1499-1508). Because zmalloc_used_memory() now includes used_memory_external, the same bytes also flow into used_memory_dataset, MEMORY STATS dataset.bytes, and keys.bytes-per-key (src/object.c:1895-1902).
That reclassifies module-only external memory as dataset memory: a module that accounts background-queue or mmap-backed state here will make the dataset counters grow even when no additional keys were stored. Keep the counter in the total/maxmemory path, but thread it through serverMemOverhead as its own bucket (or subtract zmalloc_used_external_memory() before computing mh->dataset) so the dataset metrics keep describing keyspace memory.
| /* Track memory that a module allocated outside the server allocator. | ||
| * The caller must invoke this API from a command callback or while holding | ||
| * a thread safe context lock. */ | ||
| int VM_AllocateExternalMemory(size_t bytes) { | ||
| if (zmalloc_increase_used_memory_external(bytes) != 0) { | ||
| errno = ERANGE; | ||
| return VALKEYMODULE_ERR; | ||
| } | ||
| return VALKEYMODULE_OK; | ||
| } | ||
|
|
There was a problem hiding this comment.
Why having the locking constraint? Ultimately this devolves to an atomic increment, so why have the restriction?
If we decide to retain the restriction, then we should check & abort when it's violated.
| int zmalloc_increase_used_memory_external(size_t size) { | ||
| size_t current = atomic_load_explicit(&used_memory_external, memory_order_relaxed); | ||
| if (SIZE_MAX - current < size) return -1; | ||
| atomic_store_explicit(&used_memory_external, current + size, memory_order_relaxed); | ||
| return 0; | ||
| } | ||
|
|
||
| int zmalloc_decrease_used_memory_external(size_t size) { | ||
| size_t current = atomic_load_explicit(&used_memory_external, memory_order_relaxed); | ||
| if (current < size) return -1; | ||
| atomic_store_explicit(&used_memory_external, current - size, memory_order_relaxed); | ||
| return 0; |
There was a problem hiding this comment.
I really like the "failure" when the sizes go out of bounds. But it looks like this is the only place that actually requires external mutual exclusion -- so it could be recoded with a compare & swap to provide equivalent functionality but remove the single-thread/external locking constraint.
Signed-off-by: Su Ko <rhtn1128@gmail.com>
| info, | ||
| "# Memory\r\n" FMTARGS( | ||
| "used_memory:%zu\r\n", zmalloc_used, | ||
| "used_memory_module_external:%zu\r\n", module_external_memory, |
There was a problem hiding this comment.
Do we anticipate users needing this? Modules themselves can report their own memory usage
There was a problem hiding this comment.
For developers, having visibility into this field could be useful. I don't see it being particularly helpful to end-users.
There was a problem hiding this comment.
I think reporting at both the core and module levels would be useful for developer
The core could report the total amount of external memory, while each module could report how much external memory it uses.
There was a problem hiding this comment.
Perhaps it is a good case for a debug INFO field then? https://github.com/valkey-io/valkey/blob/50a0397f4f225603c6c7b68e5c172aa3d3102dad/src/server.c#L6829C33-L6840
There was a problem hiding this comment.
I agree. This seems like a good use case for it 👍
There was a problem hiding this comment.
So the INFO field issue is resolved now, right? Is this now in a state for voting?
There was a problem hiding this comment.
The code still has it in memory stats, right? @bandalgomsu are we aligned to move it to DEBUG?
There was a problem hiding this comment.
Yes, I moved it to DEBUG in the latest commit.
|
Once we decide on the INFO field, we can do a vote on the major decision |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/server.c`:
- Around line 6830-6835: Add an INFO integration regression test in the existing
tests/INFO coverage that verifies used_memory_module_external appears in INFO
DEBUG with the expected accounted value and is absent from the general INFO
MEMORY output. Reuse the existing test setup and INFO parsing/assertion helpers
rather than adding new infrastructure.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
Signed-off-by: Su Ko <rhtn1128@gmail.com>
7960bdf to
9f50527
Compare
3 participants
Summary
This PR adds module APIs for accounting memory that is not tracked by
zmalloc, and includes that memory inused_memory/ maxmemory bookkeeping.New APIs:
ValkeyModule_AllocateExternalMemory(size_t bytes)ValkeyModule_FreeExternalMemory(size_t bytes)These APIs are intended to be used by modules from a command callback or from a locked thread-safe context.
Problem
Some module architectures, especially asynchronous indexing pipelines, can create memory pressure that the core does not currently see.
A concrete example is the Search module ingestion flow:
That gap degrades OOM detection and allows the process to run significantly past the configured
maxmemorythreshold before the core reacts.This also applies to memory that modules allocate outside
zmalloc, for example OS-backed allocations such asmmapregions used for huge-page-backed data structures.Solution
This change introduces an explicit external-memory accounting path for modules.
Internally:
zmalloc_used_memory()adds that counter to the allocator-tracked totalused_memory, andused_memory_peakall include module external memoryCloses : #3339