1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
use std::alloc::{GlobalAlloc, Layout};
use crate::raw;
/// Panics with a message without using an allocator.
/// Useful when using the allocator should be avoided or it is
/// inaccessible. The default [std::panic] performs allocations and so
/// will cause a double panic without a meaningful message if the
/// allocator can't be used. This function makes sure we can panic with
/// a reasonable message even without the allocator working.
fn allocation_free_panic(message: &'static str) -> ! {
use std::os::unix::io::AsRawFd;
let _ = nix::unistd::write(std::io::stderr().as_raw_fd(), message.as_bytes());
std::process::abort();
}
const VALKEY_ALLOCATOR_NOT_AVAILABLE_MESSAGE: &str =
"Critical error: the Valkey Allocator isn't available.\n";
/// Defines the Valkey allocator. This allocator delegates the allocation
/// and deallocation tasks to the Valkey server when available, otherwise
/// it panics.
#[derive(Copy, Clone)]
pub struct ValkeyAlloc;
unsafe impl GlobalAlloc for ValkeyAlloc {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
/*
* To make sure the memory allocation by Valkey is aligned to the according to the layout,
* we need to align the size of the allocation to the layout.
*
* "Memory is conceptually broken into equal-sized chunks,
* where the chunk size is a power of two that is greater than the page size.
* Chunks are always aligned to multiples of the chunk size.
* This alignment makes it possible to find metadata for user objects very quickly."
*
* From: https://linux.die.net/man/3/jemalloc
*/
if cfg!(feature = "enable-system-alloc") {
return std::alloc::System.alloc(layout);
}
let size = (layout.size() + layout.align() - 1) & (!(layout.align() - 1));
match raw::RedisModule_Alloc {
Some(alloc) => alloc(size).cast(),
None => allocation_free_panic(VALKEY_ALLOCATOR_NOT_AVAILABLE_MESSAGE),
}
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
/*
* To make sure the memory allocation by Valkey is aligned to the according to the layout,
* we need to align the size of the allocation to the layout.
*
* "Memory is conceptually broken into equal-sized chunks,
* where the chunk size is a power of two that is greater than the page size.
* Chunks are always aligned to multiples of the chunk size.
* This alignment makes it possible to find metadata for user objects very quickly."
*
* From: https://linux.die.net/man/3/jemalloc
*/
if cfg!(feature = "enable-system-alloc") {
return std::alloc::System.alloc_zeroed(layout);
}
let size = (layout.size() + layout.align() - 1) & (!(layout.align() - 1));
let num_elements = size / layout.align();
match raw::RedisModule_Calloc {
Some(calloc) => calloc(num_elements, layout.align()).cast(),
None => allocation_free_panic(VALKEY_ALLOCATOR_NOT_AVAILABLE_MESSAGE),
}
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
if cfg!(feature = "enable-system-alloc") {
return std::alloc::System.dealloc(ptr, layout);
}
match raw::RedisModule_Free {
Some(f) => f(ptr.cast()),
None => allocation_free_panic(VALKEY_ALLOCATOR_NOT_AVAILABLE_MESSAGE),
};
}
}