valkey_module/
redismodule.rs

1use std::borrow::Borrow;
2use std::convert::TryFrom;
3use std::ffi::CString;
4use std::fmt::Display;
5use std::ops::Deref;
6use std::os::raw::{c_char, c_int, c_void};
7use std::ptr::{null_mut, NonNull};
8use std::slice;
9use std::str;
10use std::str::Utf8Error;
11use std::string::FromUtf8Error;
12use std::{fmt, ptr};
13
14use serde::de::{Error, SeqAccess};
15
16pub use crate::raw;
17pub use crate::rediserror::ValkeyError;
18pub use crate::redisvalue::ValkeyValue;
19use crate::Context;
20
21/// A short-hand type that stores a [std::result::Result] with custom
22/// type and [RedisError].
23pub type ValkeyResult<T = ValkeyValue> = Result<T, ValkeyError>;
24/// A [RedisResult] with [ValkeyValue].
25pub type ValkeyValueResult = ValkeyResult<ValkeyValue>;
26
27impl From<ValkeyValue> for ValkeyValueResult {
28    fn from(v: ValkeyValue) -> Self {
29        Ok(v)
30    }
31}
32
33impl From<ValkeyError> for ValkeyValueResult {
34    fn from(v: ValkeyError) -> Self {
35        Err(v)
36    }
37}
38
39pub const VALKEY_OK: ValkeyValueResult = Ok(ValkeyValue::SimpleStringStatic("OK"));
40pub const TYPE_METHOD_VERSION: u64 = raw::REDISMODULE_TYPE_METHOD_VERSION as u64;
41pub const AUTH_HANDLED: i32 = raw::REDISMODULE_AUTH_HANDLED as i32;
42pub const AUTH_NOT_HANDLED: i32 = raw::REDISMODULE_AUTH_NOT_HANDLED as i32;
43
44pub trait NextArg {
45    fn next_arg(&mut self) -> Result<ValkeyString, ValkeyError>;
46    fn next_string(&mut self) -> Result<String, ValkeyError>;
47    fn next_str<'a>(&mut self) -> Result<&'a str, ValkeyError>;
48    fn next_i64(&mut self) -> Result<i64, ValkeyError>;
49    fn next_u64(&mut self) -> Result<u64, ValkeyError>;
50    fn next_f64(&mut self) -> Result<f64, ValkeyError>;
51    fn done(&mut self) -> Result<(), ValkeyError>;
52}
53
54impl<T> NextArg for T
55where
56    T: Iterator<Item = ValkeyString>,
57{
58    #[inline]
59    fn next_arg(&mut self) -> Result<ValkeyString, ValkeyError> {
60        self.next().ok_or(ValkeyError::WrongArity)
61    }
62
63    #[inline]
64    fn next_string(&mut self) -> Result<String, ValkeyError> {
65        self.next()
66            .map_or(Err(ValkeyError::WrongArity), |v| Ok(v.to_string_lossy()))
67    }
68
69    #[inline]
70    fn next_str<'a>(&mut self) -> Result<&'a str, ValkeyError> {
71        self.next()
72            .map_or(Err(ValkeyError::WrongArity), |v| v.try_as_str())
73    }
74
75    #[inline]
76    fn next_i64(&mut self) -> Result<i64, ValkeyError> {
77        self.next()
78            .map_or(Err(ValkeyError::WrongArity), |v| v.parse_integer())
79    }
80
81    #[inline]
82    fn next_u64(&mut self) -> Result<u64, ValkeyError> {
83        self.next()
84            .map_or(Err(ValkeyError::WrongArity), |v| v.parse_unsigned_integer())
85    }
86
87    #[inline]
88    fn next_f64(&mut self) -> Result<f64, ValkeyError> {
89        self.next()
90            .map_or(Err(ValkeyError::WrongArity), |v| v.parse_float())
91    }
92
93    /// Return an error if there are any more arguments
94    #[inline]
95    fn done(&mut self) -> Result<(), ValkeyError> {
96        self.next().map_or(Ok(()), |_| Err(ValkeyError::WrongArity))
97    }
98}
99
100#[allow(clippy::not_unsafe_ptr_arg_deref)]
101pub fn decode_args(
102    ctx: *mut raw::RedisModuleCtx,
103    argv: *mut *mut raw::RedisModuleString,
104    argc: c_int,
105) -> Vec<ValkeyString> {
106    if argv.is_null() {
107        return Vec::new();
108    }
109    unsafe { slice::from_raw_parts(argv, argc as usize) }
110        .iter()
111        .map(|&arg| ValkeyString::new(NonNull::new(ctx), arg))
112        .collect()
113}
114
115///////////////////////////////////////////////////
116
117#[derive(Debug)]
118pub struct ValkeyString {
119    ctx: *mut raw::RedisModuleCtx,
120    pub inner: *mut raw::RedisModuleString,
121}
122
123impl ValkeyString {
124    pub(crate) fn take(mut self) -> *mut raw::RedisModuleString {
125        let inner = self.inner;
126        self.inner = std::ptr::null_mut();
127        inner
128    }
129
130    pub fn new(
131        ctx: Option<NonNull<raw::RedisModuleCtx>>,
132        inner: *mut raw::RedisModuleString,
133    ) -> Self {
134        let ctx = ctx.map_or(std::ptr::null_mut(), |v| v.as_ptr());
135        raw::string_retain_string(ctx, inner);
136        Self { ctx, inner }
137    }
138
139    /// In general, [RedisModuleString] is none atomic ref counted object.
140    /// So it is not safe to clone it if Valkey GIL is not held.
141    /// [Self::safe_clone] gets a context reference which indicates that Valkey GIL is held.
142    pub fn safe_clone(&self, _ctx: &Context) -> Self {
143        // RedisString are *not* atomic ref counted, so we must get a lock indicator to clone them.
144        // Alos notice that Valkey allows us to create RedisModuleString with NULL context
145        // so we use [std::ptr::null_mut()] instead of the curren RedisString context.
146        // We do this because we can not promise the new RedisString will not outlive the current
147        // context and we want them to be independent.
148        raw::string_retain_string(ptr::null_mut(), self.inner);
149        Self {
150            ctx: ptr::null_mut(),
151            inner: self.inner,
152        }
153    }
154
155    #[allow(clippy::not_unsafe_ptr_arg_deref)]
156    pub fn create<T: Into<Vec<u8>>>(ctx: Option<NonNull<raw::RedisModuleCtx>>, s: T) -> Self {
157        let ctx = ctx.map_or(std::ptr::null_mut(), |v| v.as_ptr());
158        let str = CString::new(s).unwrap();
159        let inner = unsafe {
160            raw::RedisModule_CreateString.unwrap()(ctx, str.as_ptr(), str.as_bytes().len())
161        };
162
163        Self { ctx, inner }
164    }
165
166    #[allow(clippy::not_unsafe_ptr_arg_deref)]
167    pub fn create_from_slice(ctx: *mut raw::RedisModuleCtx, s: &[u8]) -> Self {
168        let inner = unsafe {
169            raw::RedisModule_CreateString.unwrap()(ctx, s.as_ptr().cast::<c_char>(), s.len())
170        };
171
172        Self { ctx, inner }
173    }
174
175    /// Creates a ValkeyString from a &str and retains it.  This is useful in cases where Modules need to pass ownership of a ValkeyString to the core engine without it being freed when we drop a ValkeyString
176    pub fn create_and_retain(arg: &str) -> ValkeyString {
177        let arg = ValkeyString::create(None, arg);
178        raw::string_retain_string(null_mut(), arg.inner);
179        arg
180    }
181
182    pub const fn from_redis_module_string(
183        ctx: *mut raw::RedisModuleCtx,
184        inner: *mut raw::RedisModuleString,
185    ) -> Self {
186        // Need to avoid string_retain_string
187        Self { ctx, inner }
188    }
189
190    #[allow(clippy::not_unsafe_ptr_arg_deref)]
191    pub fn from_ptr<'a>(ptr: *const raw::RedisModuleString) -> Result<&'a str, Utf8Error> {
192        str::from_utf8(Self::string_as_slice(ptr))
193    }
194
195    pub fn append(&mut self, s: &str) -> raw::Status {
196        raw::string_append_buffer(self.ctx, self.inner, s)
197    }
198
199    #[must_use]
200    pub fn len(&self) -> usize {
201        let mut len: usize = 0;
202        raw::string_ptr_len(self.inner, &mut len);
203        len
204    }
205
206    #[must_use]
207    pub fn is_empty(&self) -> bool {
208        let mut len: usize = 0;
209        raw::string_ptr_len(self.inner, &mut len);
210        len == 0
211    }
212
213    pub fn try_as_str<'a>(&self) -> Result<&'a str, ValkeyError> {
214        Self::from_ptr(self.inner).map_err(|_| ValkeyError::Str("Couldn't parse as UTF-8 string"))
215    }
216
217    #[must_use]
218    pub fn as_slice(&self) -> &[u8] {
219        Self::string_as_slice(self.inner)
220    }
221
222    #[allow(clippy::not_unsafe_ptr_arg_deref)]
223    pub fn string_as_slice<'a>(ptr: *const raw::RedisModuleString) -> &'a [u8] {
224        let mut len: libc::size_t = 0;
225        let bytes = unsafe { raw::RedisModule_StringPtrLen.unwrap()(ptr, &mut len) };
226
227        unsafe { slice::from_raw_parts(bytes.cast::<u8>(), len) }
228    }
229
230    /// Performs lossy conversion of a `RedisString` into an owned `String. This conversion
231    /// will replace any invalid UTF-8 sequences with U+FFFD REPLACEMENT CHARACTER, which
232    /// looks like this: �.
233    ///
234    /// # Panics
235    ///
236    /// Will panic if `RedisModule_StringPtrLen` is missing in redismodule.h
237    #[must_use]
238    pub fn to_string_lossy(&self) -> String {
239        String::from_utf8_lossy(self.as_slice()).into_owned()
240    }
241
242    pub fn parse_unsigned_integer(&self) -> Result<u64, ValkeyError> {
243        let val = self.parse_integer()?;
244        u64::try_from(val)
245            .map_err(|_| ValkeyError::Str("Couldn't parse negative number as unsigned integer"))
246    }
247
248    pub fn parse_integer(&self) -> Result<i64, ValkeyError> {
249        let mut val: i64 = 0;
250        match raw::string_to_longlong(self.inner, &mut val) {
251            raw::Status::Ok => Ok(val),
252            raw::Status::Err => Err(ValkeyError::Str("Couldn't parse as integer")),
253        }
254    }
255
256    pub fn parse_float(&self) -> Result<f64, ValkeyError> {
257        let mut val: f64 = 0.0;
258        match raw::string_to_double(self.inner, &mut val) {
259            raw::Status::Ok => Ok(val),
260            raw::Status::Err => Err(ValkeyError::Str("Couldn't parse as float")),
261        }
262    }
263
264    // TODO: Valkey allows storing and retrieving any arbitrary bytes.
265    // However rust's String and str can only store valid UTF-8.
266    // Implement these to allow non-utf8 bytes to be consumed:
267    // pub fn into_bytes(self) -> Vec<u8> {}
268    // pub fn as_bytes(&self) -> &[u8] {}
269}
270
271impl Drop for ValkeyString {
272    fn drop(&mut self) {
273        if !self.inner.is_null() {
274            unsafe {
275                raw::RedisModule_FreeString.unwrap()(self.ctx, self.inner);
276            }
277        }
278    }
279}
280
281impl PartialEq for ValkeyString {
282    fn eq(&self, other: &Self) -> bool {
283        self.cmp(other).is_eq()
284    }
285}
286
287impl Eq for ValkeyString {}
288
289impl PartialOrd for ValkeyString {
290    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
291        Some(self.cmp(other))
292    }
293}
294
295impl Ord for ValkeyString {
296    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
297        raw::string_compare(self.inner, other.inner)
298    }
299}
300
301impl core::hash::Hash for ValkeyString {
302    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
303        self.as_slice().hash(state);
304    }
305}
306
307impl Display for ValkeyString {
308    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
309        write!(f, "{}", self.to_string_lossy())
310    }
311}
312
313impl Borrow<str> for ValkeyString {
314    fn borrow(&self) -> &str {
315        // RedisString might not be UTF-8 safe
316        self.try_as_str().unwrap_or("<Invalid UTF-8 data>")
317    }
318}
319
320impl Clone for ValkeyString {
321    fn clone(&self) -> Self {
322        let inner =
323            // Valkey allows us to create RedisModuleString with NULL context
324            // so we use [std::ptr::null_mut()] instead of the curren RedisString context.
325            // We do this because we can not promise the new RedisString will not outlive the current
326            // context and we want them to be independent.
327            unsafe { raw::RedisModule_CreateStringFromString.unwrap()(ptr::null_mut(), self.inner) };
328        Self::from_redis_module_string(ptr::null_mut(), inner)
329    }
330}
331
332impl From<ValkeyString> for String {
333    fn from(rs: ValkeyString) -> Self {
334        rs.to_string_lossy()
335    }
336}
337
338impl Deref for ValkeyString {
339    type Target = [u8];
340
341    fn deref(&self) -> &Self::Target {
342        self.as_slice()
343    }
344}
345
346impl From<ValkeyString> for Vec<u8> {
347    fn from(rs: ValkeyString) -> Self {
348        rs.as_slice().to_vec()
349    }
350}
351
352impl serde::Serialize for ValkeyString {
353    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
354    where
355        S: serde::Serializer,
356    {
357        serializer.serialize_bytes(self.as_slice())
358    }
359}
360
361struct RedisStringVisitor;
362
363impl<'de> serde::de::Visitor<'de> for RedisStringVisitor {
364    type Value = ValkeyString;
365
366    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
367        formatter.write_str("A bytes buffer")
368    }
369
370    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
371    where
372        E: Error,
373    {
374        Ok(ValkeyString::create(None, v))
375    }
376
377    fn visit_seq<V>(self, mut visitor: V) -> Result<Self::Value, V::Error>
378    where
379        V: SeqAccess<'de>,
380    {
381        let mut v = if let Some(size_hint) = visitor.size_hint() {
382            Vec::with_capacity(size_hint)
383        } else {
384            Vec::new()
385        };
386        while let Some(elem) = visitor.next_element()? {
387            v.push(elem);
388        }
389
390        Ok(ValkeyString::create(None, v.as_slice()))
391    }
392}
393
394impl<'de> serde::Deserialize<'de> for ValkeyString {
395    fn deserialize<D>(deserializer: D) -> Result<ValkeyString, D::Error>
396    where
397        D: serde::Deserializer<'de>,
398    {
399        deserializer.deserialize_bytes(RedisStringVisitor)
400    }
401}
402
403///////////////////////////////////////////////////
404
405#[derive(Debug)]
406pub struct RedisBuffer {
407    buffer: *mut c_char,
408    len: usize,
409}
410
411impl RedisBuffer {
412    pub const fn new(buffer: *mut c_char, len: usize) -> Self {
413        Self { buffer, len }
414    }
415
416    pub fn to_string(&self) -> Result<String, FromUtf8Error> {
417        String::from_utf8(self.as_ref().to_vec())
418    }
419}
420
421impl AsRef<[u8]> for RedisBuffer {
422    fn as_ref(&self) -> &[u8] {
423        unsafe { slice::from_raw_parts(self.buffer as *const u8, self.len) }
424    }
425}
426
427impl Drop for RedisBuffer {
428    fn drop(&mut self) {
429        unsafe {
430            raw::RedisModule_Free.unwrap()(self.buffer.cast::<c_void>());
431        }
432    }
433}
valkey_module/redismodule.rs

valkey_module/
redismodule.rs
