valkey_module/
alloc.rs

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

valkey_module/
alloc.rs
