Rax size tracking

[knggk]

Introduce a size_t field into the rax struct to track allocation size.
Update the allocation size on rax insert and deletes.
Return the allocation size when raxAllocSize is called.

This size tracking is now used in MEMORY USAGE and MEMORY STATS in place of the previous method based on sampling.

The module API allows to create sorted dictionaries, which are backed by rax. Users now also get precise memory allocation for them (through ValkeyModule_MallocSizeDict).

Fixes #677.

For the release notes:

• MEMORY USAGE and MEMORY STATS are now exact for streams, rather than based on sampling.

[kyle-yh-kim]

To resolve DCO complaints, you can sign-off the previous commits.

// Sign-off the last 3 commits.
git rebase --signoff HEAD~3

// Force push to update your PR.
git push -f

[kyle-yh-kim]

Discussed with @knggk offline. A few parallelizable work for this PR;

• Import rax-test.c and stabilize it [source].
• Devise new test cases under rax-test.c, for size tracking.
• MEMORY STATS integration.

[codecov]

Codecov Report

Attention: Patch coverage is 75.00000% with 7 lines in your changes missing coverage. Please review.

Project coverage is 70.60%. Comparing base (9827eef) to head (fee849d).
Report is 1 commits behind head on unstable.

Files with missing lines	Patch %	Lines
src/object.c	0.00%	3 Missing ⚠️
src/module.c	0.00%	2 Missing ⚠️
src/rax.c	91.30%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           unstable     #688      +/-   ##
============================================
+ Coverage     70.47%   70.60%   +0.13%     
============================================
  Files           114      114              
  Lines         61695    61710      +15     
============================================
+ Hits          43479    43572      +93     
+ Misses        18216    18138      -78     

Files with missing lines	Coverage Δ
src/module.c	9.64% <0.00%> (+<0.01%)	⬆️
src/rax.c	82.92% <91.30%> (+0.22%)	⬆️
src/object.c	79.18% <0.00%> (+0.53%)	⬆️

... and 13 files with indirect coverage changes

[zuiderkwast]

In general LGTM. The test needs to be ported to the new kind of unit tests.

Is is possible to test the calculation of the alloc size by comparing it to the stats we can get from the allocator? If we run a fuzz test and it does no other allocations, then it could work I think.

[knggk]

In general LGTM. The test needs to be ported to the new kind of unit tests.

Is is possible to test the calculation of the alloc size by comparing it to the stats we can get from the allocator? If we run a fuzz test and it does no other allocations, then it could work I think.

Great Idea. I wanted also to check that after every element from the rax was removed, memory goes back to 0. Turns out there is an issue right now that needs further investigation, ie src/valkey-unit-tests --large-memory says:

...
[ok] - test_rax.c:test_hugeKey
...
[END] - test_rax.c: 12 tests, 12 passed, 0 failed
...
[START] - test_zmalloc.c
[test_zmallocInitialUsedMemory - unit/test_zmalloc.c:9] Failed assertion: zmalloc_used_memory() == 0

[zuiderkwast]

Just an idea: If it's a problem with zmalloc, you can create another rax_test_malloc.h to instead of rax_malloc.h just for the rax test. If the test defines RAX_MALLOC_INCLUDE before it includes rax.c, then rax.c will use it. It's already like this in rax.c:

#ifndef RAX_MALLOC_INCLUDE
#define RAX_MALLOC_INCLUDE "rax_malloc.h"
#endif

#include RAX_MALLOC_INCLUDE

[knggk]

RAX_MALLOC_INCLUDE

Thanks for the tip @zuiderkwast . I think it wasn't a problem with zmalloc, see the latest two commits. There was a zfree missing which was the cause of valkey-unit-tests --large-memory failing.

However you gave me another idea: what if instead of manually tracking rax->alloc_size in the main code body, we made rax_alloc_size and friends do it, as in rax_alloc_size(rax, size) { void *x = zmalloc(size); rax->alloc_size += zmalloc_usable_memory(x); x } ? It sounds like it'd be harder to make mistakes like the one fixed in "Tentative fix for address sanitizer", which tried to use an address that's already been freed. What do you think?

[zuiderkwast]

Please check clang-format and spellcheck. You can change "ba" to "by" (etc.) in the rax_test, or if this actually affects the test, it's OK to add this file to .config/typos.toml to exclude the file from spell check.

[knggk]

I see. I had a quick exchange with @touitou-dan who was suggesting a 2% degradation in pipeline only can be ignored. After all pipeline is not the standard mode of benchmark.

In light of that, do we know the degradation of the memory tracking for the other pieces? If similar, maybe it's worth not making it configurable. Thoughts @zuiderkwast?

[zuiderkwast]

This will have to wait until 8.2 8.1, so we should think about it after 8.0 rc1...

[ranshid]

@knggk Did we consider ways to avoid calling zmalloc_size on each rax operation? IIRC these can be expensive and we should aim to avoid them. For example It might be messy but we could try and use the jemalloc *_usable API when allocating memory and retrieve the allocation size.
Also as we now track the used_memory_thread it might be easier to just take the diff during RAX operations instead of going through the "expensive" zmalloc_size calls?

[knggk]

Did we consider ways to avoid calling zmalloc_size on each rax operation?

But it only happens on rax mutations right?

try and use the jemalloc *_usable API when allocating memory and retrieve the allocation size.

What are those? Are you saying they are less expensive than zmalloc_size? What's the context behind zmalloc_size being expensive?

Also as we now track the used_memory_thread it might be easier to just take the diff during RAX operations instead of going through the "expensive" zmalloc_size calls?

Agree, good idea if this new used_memory_thread is updated during allocation calls and would capture a diff when called right before allocating eg a raxNode, and again right after. Is this how used_memory_thread behaves?

[ranshid]

Did we consider ways to avoid calling zmalloc_size on each rax operation?

But it only happens on rax mutations right?

yes

try and use the jemalloc *_usable API when allocating memory and retrieve the allocation size.

What are those? Are you saying they are less expensive than zmalloc_size? What's the context behind zmalloc_size being expensive?

I mean for example there is the 'oid *zmalloc_usable(size_t size, size_t *usable)' which both allocate memory and return the allocation size, so it can be used in a single call instead of later calling 'zmalloc_size'

Also as we now track the used_memory_thread it might be easier to just take the diff during RAX operations instead of going through the "expensive" zmalloc_size calls?

Agree, good idea if this new used_memory_thread is updated during allocation calls and would capture a diff when called right before allocating eg a raxNode, and again right after. Is this how used_memory_thread behaves?

Basically since we are taking the diff without being preempted and doing "something else" I think we can count of this diff to reflect the exact amount of memory we allocated/freed during the rax operation.

[zuiderkwast]

I mean for example there is the void *zmalloc_usable(size_t size, size_t *usable) which both allocate memory and return the allocation size, so it can be used in a single call instead of later calling 'zmalloc_size'

Good idea. If it's not too complicated, it would be better to use it, but I'm not sure zmalloc_size is very slow either. Maybe it's not a significant part of the work in the rax update.

Basically since we are taking the diff without being preempted and doing "something else" I think we can count of this diff to reflect the exact amount of memory we allocated/freed during the rax operation.

I think this can work, but note that used_memory_thread per thread is a size_t which can wrap around zero. For example, if one thread T1 allocates some memory and another thread T2 frees the memory, then T2 can end up freeing more memory than it allocates, i.e. negative number, wrapping from 0 to SIZE_MAX. It can probably be handled. Just be careful.

IMO, we can also accept this PR with zmalloc_size now and allow it to be optimized in the future.

[ranshid]

IMO, we can also accept this PR with zmalloc_size now and allow it to be optimized in the future.

Can we open an issue for that then?

[ranshid]

I think this can work, but note that used_memory_thread per thread is a size_t which can wrap around zero. For example, if one thread T1 allocates some memory and another thread T2 frees the memory, then T2 can end up freeing more memory than it allocates, i.e. negative number, wrapping from 0 to SIZE_MAX. It can probably be handled. Just be careful.

Not sure there is a problem. we only take a diff of the thread local variable, so even if it was negative the diff calculation will still be fine.

[zuiderkwast]

I think this can work, but note that used_memory_thread per thread is a size_t which can wrap around zero. For example, if one thread T1 allocates some memory and another thread T2 frees the memory, then T2 can end up freeing more memory than it allocates, i.e. negative number, wrapping from 0 to SIZE_MAX. It can probably be handled. Just be careful.

Not sure there is a problem. we only take a diff of the thread local variable, so even if it was negative the diff calculation will still be fine.

@ranshid 🤔 You're probably right.

Another problem is that we only track a certain number of threads. If a module creates a lot of threads and they use the rax (e.g. using ValkeyModule_DictXxxx functions) these threads will not have their own variable.

I don't know how to solve that. Do you have an idea? The per-thread size tracking was not designed for what we're discussing here. It was only designed to speed up the total size tracking. Letting the rax track its own allocations seems more robust to me somehow.

Can we open an issue for that then?

Sure.

@knggk, will you merge latest unstable to this branch?

[knggk]

@knggk, will you merge latest unstable to this branch?

I tried just now but am getting:

% make valkey-unit-tests
...
    LINK valkey-unit-tests
duplicate symbol '__serverAssert' in:
    /Users/knggk/valkey/src/unit/test_main.o
    libvalkey.a[36](debug.o)
duplicate symbol '_main' in:
    /Users/knggk/valkey/src/unit/test_main.o
    libvalkey.a[9](server.o)
ld: 2 duplicate symbols
clang: error: linker command failed with exit code 1 (use -v to see invocation)

Edit: Nvm, this also happens on clean unstable. Might be a problem with my local toolchain?

[zuiderkwast]

Edit: Nvm, this also happens on clean unstable. Might be a problem with my local toolchain?

Seems to be, because the build works in the CI. Did you try make distclean?

[knggk]

Seems to be, because the build works in the CI. Did you try make distclean?

I tried with distclean but still getting the same issue on link. I am on Mac Sonoma if that's any useful info. Could is be because of missing pkg-config? I get:

% make valkey-unit-tests
make valkey-unit-tests
cd src && /Library/Developer/CommandLineTools/usr/bin/make valkey-unit-tests
/bin/sh: pkg-config: command not found
/bin/sh: pkg-config: command not found
/bin/sh: pkg-config: command not found
...

PS pushed commit to fix formatting

[zuiderkwast]

Great, thanks.

Could is be because of missing pkg-config? I get:

Possibly. I'm running Fedora on my macbook. It seemed better for programmers. :D

pkg-config is used in the Makefile and maybe it just exits if it's not available. Why don't you have it? And why don't you have make in your path? Maybe you just need some basic "build util" package from brew or something...?

[zuiderkwast]

@knggk There was a build failure with undefined-santitizer in Daily:

...
Cluster Fuzz test [keys:99963652 keylen:0]: ok with 71243660 final keys
Cluster Fuzz test [keys:37437816 keylen:0]: ok with 34387942 final keys
Cluster Fuzz test [keys:22357214 keylen:0]: ok with 5758528 final keys
Fuzz test in mode 0 [9067]: 8396 elements inserted
unit/test_rax.c:168:15: runtime error: left shift of 36625 by 16 places cannot be represented in type 'int'
SUMMARY: UndefinedBehaviorSanitizer: undefined-behavior unit/test_rax.c:168:15 in 
Fuzz test in mode 1 [7504]: 

Do you want to look into it? You need to build with SANITIZER=undefined and run the tests with --accurate.

I have a fix for it in #1122.

[knggk]

Thanks Viktor for the turn around, I hadn't seen the issue.