Skip to content

Commit e56ed89

Browse files
gavriegavrie
committed
Store custom type in a key
1 parent 232abb8 commit e56ed89

File tree

4 files changed

+187
-45
lines changed

4 files changed

+187
-45
lines changed

‎examples/data_type.rs‎

Lines changed: 76 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
use std::os::raw::c_void;
2+
13
extern crate libc;
24

35
use libc::c_int;
@@ -9,12 +11,17 @@ use redismodule::Command;
911
use redismodule::raw;
1012
use redismodule::Redis;
1113
use redismodule::raw::module_init;
12-
use redismodule::types::RedisModuleType;
14+
use redismodule::types::RedisType;
1315

1416
const MODULE_NAME: &str = "alloc";
1517
const MODULE_VERSION: c_int = 1;
1618

17-
static MY_TYPE: RedisModuleType = RedisModuleType::new();
19+
#[allow(unused)]
20+
struct MyType {
21+
data: String,
22+
}
23+
24+
static MY_TYPE: RedisType = RedisType::new();
1825

1926
struct AllocSetCommand;
2027

@@ -40,26 +47,21 @@ impl Command for AllocSetCommand {
4047
let size = parse_i64(args[2])?;
4148

4249
// TODO:
43-
// 1. Open key
44-
// 2. Allocate data
45-
// 3. Set the key to the data
50+
// 1. Open key [OK]
51+
// 2. Allocate data [OK]
52+
// 3. Set the key to the data [OK]
4653
// 4. Activate custom allocator and compare Redis memory usage
4754

48-
//let data: Vec<u8> = Vec::with_capacity(size as usize);
49-
5055
let key = r.open_key_writable(key);
56+
key.check_type(&MY_TYPE)?;
5157

52-
//key.check_type(MY_TYPE)?;
53-
key.write(size.to_string().as_str())?;
54-
55-
/*
56-
raw::RedisModule_ModuleTypeSetValue.unwrap()(
57-
k,
58-
t,
59-
data,
60-
);
61-
*/
58+
let my = Box::into_raw(Box::new(
59+
MyType {
60+
data: "A".repeat(size as usize)
61+
}
62+
));
6263

64+
key.set_value(&MY_TYPE, my as *mut c_void)?;
6365
r.reply_integer(size)?;
6466

6567
Ok(())
@@ -77,6 +79,61 @@ pub extern "C" fn AllocSetCommand_Redis(
7779

7880
//////////////////////////////////////////////////////
7981

82+
struct AllocGetCommand;
83+
84+
impl Command for AllocGetCommand {
85+
fn name() -> &'static str { "alloc.get" }
86+
87+
fn external_command() -> raw::CommandFunc { AllocGetCommand_Redis }
88+
89+
fn str_flags() -> &'static str { "" }
90+
91+
// Run the command.
92+
fn run(r: Redis, args: &[&str]) -> Result<(), Error> {
93+
if args.len() != 2 {
94+
// FIXME: Use RedisModule_WrongArity instead. Return an ArityError here and
95+
// in the low-level implementation call RM_WrongArity.
96+
return Err(Error::generic(format!(
97+
"Usage: {} <key>", Self::name()
98+
).as_str()));
99+
}
100+
101+
// the first argument is command name (ignore it)
102+
let key = args[1];
103+
104+
let key = r.open_key(key);
105+
106+
key.check_type(&MY_TYPE)?;
107+
108+
let my = key.get_value() as *mut MyType;
109+
110+
if my.is_null() {
111+
r.reply_integer(0)?;
112+
return Ok(());
113+
}
114+
115+
let my = unsafe { &mut *my };
116+
let size = my.data.len();
117+
118+
r.reply_array(2)?;
119+
r.reply_integer(size as i64)?;
120+
r.reply_string(my.data.as_str())?;
121+
122+
Ok(())
123+
}
124+
}
125+
126+
#[allow(non_snake_case)]
127+
pub extern "C" fn AllocGetCommand_Redis(
128+
ctx: *mut raw::RedisModuleCtx,
129+
argv: *mut *mut raw::RedisModuleString,
130+
argc: c_int,
131+
) -> c_int {
132+
AllocGetCommand::execute(ctx, argv, argc).into()
133+
}
134+
135+
//////////////////////////////////////////////////////
136+
80137
struct AllocDelCommand;
81138

82139
impl Command for AllocDelCommand {
@@ -96,7 +153,7 @@ impl Command for AllocDelCommand {
96153
}
97154

98155
// the first argument is command name (ignore it)
99-
let key = args[1];
156+
let _key = args[1];
100157

101158
r.reply_string("OK")?;
102159

@@ -122,6 +179,7 @@ fn module_on_load(ctx: *mut raw::RedisModuleCtx) -> Result<(), &'static str> {
122179
MY_TYPE.create_data_type(ctx, "mytype123")?;
123180

124181
AllocSetCommand::create(ctx)?;
182+
AllocGetCommand::create(ctx)?;
125183
AllocDelCommand::create(ctx)?;
126184

127185
Ok(())

‎src/lib.rs‎

Lines changed: 45 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering::SeqCst};
4141
use std::ffi::c_void;
4242

4343
use crate::error::Error;
44-
use crate::types::RedisModuleType;
44+
use crate::types::RedisType;
4545

4646
struct RedisAlloc;
4747

@@ -249,6 +249,10 @@ pub enum KeyMode {
249249
ReadWrite,
250250
}
251251

252+
trait RedisKeyType {
253+
fn check_type(&self, redis_type: &RedisType) -> Result<(), Error>;
254+
}
255+
252256
/// `RedisKey` is an abstraction over a Redis key that allows readonly
253257
/// operations.
254258
///
@@ -289,6 +293,21 @@ impl RedisKey {
289293
};
290294
Ok(val)
291295
}
296+
297+
pub fn check_type(&self, redis_type: &RedisType) -> Result<(), Error> {
298+
match raw::check_key_type(
299+
self.ctx,
300+
self.key_inner,
301+
*redis_type.raw_type.borrow_mut(),
302+
) {
303+
Ok(_) => Ok(()),
304+
Err(s) => Err(Error::generic(s))
305+
}
306+
}
307+
308+
pub fn get_value(&self) -> *mut c_void {
309+
raw::module_type_get_value(self.key_inner)
310+
}
292311
}
293312

294313
impl Drop for RedisKey {
@@ -362,21 +381,36 @@ impl RedisKeyWritable {
362381
}
363382
}
364383

365-
pub fn check_type(&self, redis_type: RedisModuleType) -> Result<(), Error> {
366-
/*
367-
int type = RedisModule_KeyType(key);
368-
if (type != REDISMODULE_KEYTYPE_EMPTY &&
369-
RedisModule_ModuleTypeGetType(key) != HelloType)
370-
{
371-
return RedisModule_ReplyWithError(ctx,REDISMODULE_ERRORMSG_WRONGTYPE);
384+
pub fn check_type(&self, redis_type: &RedisType) -> Result<(), Error> {
385+
match raw::check_key_type(
386+
self.ctx,
387+
self.key_inner,
388+
*redis_type.raw_type.borrow_mut(),
389+
) {
390+
Ok(_) => Ok(()),
391+
Err(s) => Err(Error::generic(s))
372392
}
373-
*/
374-
//raw::key_type(key)
393+
}
394+
395+
pub fn set_value(&self, redis_type: &RedisType, value: *mut c_void) -> Result<(), Error> {
396+
raw::module_type_set_value(
397+
self.key_inner,
398+
*redis_type.raw_type.borrow_mut(),
399+
value,
400+
).into()
401+
}
402+
}
375403

376-
Ok(())
404+
impl From<raw::Status> for Result<(), Error> {
405+
fn from(s: raw::Status) -> Self {
406+
match s {
407+
raw::Status::Ok => Ok(()),
408+
raw::Status::Err => Err(Error::generic("Generic error")),
409+
}
377410
}
378411
}
379412

413+
380414
impl Drop for RedisKeyWritable {
381415
// Frees resources appropriately as a RedisKey goes out of scope.
382416
fn drop(&mut self) {

‎src/raw.rs‎

Lines changed: 56 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,18 @@
33
#![allow(dead_code)]
44

55
use std::ffi::CString;
6-
use std::os::raw::{c_char, c_int, c_long, c_longlong};
6+
use std::os::raw::{c_char, c_int, c_long, c_longlong, c_void};
77

88
extern crate libc;
99
extern crate enum_primitive_derive;
1010
extern crate num_traits;
1111

12-
use num_traits::{FromPrimitive, ToPrimitive};
12+
use num_traits::FromPrimitive;
1313

1414
use libc::size_t;
1515

1616
pub use crate::redisraw::bindings::*;
17+
use crate::types::redis_log;
1718

1819
bitflags! {
1920
pub struct KeyMode: c_int {
@@ -22,7 +23,7 @@ bitflags! {
2223
}
2324
}
2425

25-
#[derive(Primitive)]
26+
#[derive(Primitive, Debug, PartialEq)]
2627
pub enum KeyType {
2728
Empty = REDISMODULE_KEYTYPE_EMPTY as isize,
2829
String = REDISMODULE_KEYTYPE_STRING as isize,
@@ -42,10 +43,10 @@ impl From<c_int> for KeyType {
4243
// Casting to isize inside the enum definition breaks the derive(Primitive) macro.
4344
const REDISMODULE_REPLY_UNKNOWN_ISIZE: isize = REDISMODULE_REPLY_UNKNOWN as isize;
4445
const REDISMODULE_REPLY_STRING_ISIZE: isize = REDISMODULE_REPLY_STRING as isize;
45-
const REDISMODULE_REPLY_ERROR_ISIZE: isize = REDISMODULE_REPLY_ERROR as isize;
46-
const REDISMODULE_REPLY_INTEGER_ISIZE: isize = REDISMODULE_REPLY_INTEGER as isize;
47-
const REDISMODULE_REPLY_ARRAY_ISIZE: isize = REDISMODULE_REPLY_ARRAY as isize;
48-
const REDISMODULE_REPLY_NULL_ISIZE: isize = REDISMODULE_REPLY_NULL as isize;
46+
const REDISMODULE_REPLY_ERROR_ISIZE: isize = REDISMODULE_REPLY_ERROR as isize;
47+
const REDISMODULE_REPLY_INTEGER_ISIZE: isize = REDISMODULE_REPLY_INTEGER as isize;
48+
const REDISMODULE_REPLY_ARRAY_ISIZE: isize = REDISMODULE_REPLY_ARRAY as isize;
49+
const REDISMODULE_REPLY_NULL_ISIZE: isize = REDISMODULE_REPLY_NULL as isize;
4950

5051
#[derive(Primitive, Debug, PartialEq)]
5152
pub enum ReplyType {
@@ -306,3 +307,51 @@ pub fn string_set(key: *mut RedisModuleKey, str: *mut RedisModuleString) -> Stat
306307
pub fn key_type(key: *mut RedisModuleKey) -> KeyType {
307308
unsafe { RedisModule_KeyType.unwrap()(key) }.into()
308309
}
310+
311+
pub fn module_type_get_type(key: *mut RedisModuleKey) -> *mut RedisModuleType {
312+
unsafe { RedisModule_ModuleTypeGetType.unwrap()(key) }
313+
}
314+
315+
pub fn module_type_get_value(key: *mut RedisModuleKey) -> *mut c_void {
316+
unsafe {
317+
RedisModule_ModuleTypeGetValue.unwrap()(key)
318+
}.into()
319+
}
320+
321+
pub fn module_type_set_value(
322+
key: *mut RedisModuleKey,
323+
redis_type: *mut RedisModuleType,
324+
value: *mut c_void,
325+
) -> Status {
326+
unsafe {
327+
RedisModule_ModuleTypeSetValue.unwrap()(
328+
key,
329+
redis_type,
330+
value,
331+
)
332+
}.into()
333+
}
334+
335+
pub fn check_key_type(
336+
ctx: *mut RedisModuleCtx,
337+
key: *mut RedisModuleKey,
338+
redis_type: *mut RedisModuleType,
339+
) -> Result<(), &'static str> {
340+
341+
let key_type = key_type(key);
342+
343+
// TODO: Make log() a method of the Redis and Key structs.
344+
redis_log(ctx, format!("key type: {:?}", key_type).as_str());
345+
346+
if key_type != KeyType::Empty {
347+
let raw_type = module_type_get_type(key);
348+
if raw_type != redis_type{
349+
return Err("Key has existing value with wrong Redis type");
350+
}
351+
redis_log(ctx, "Existing key has the correct type");
352+
}
353+
354+
redis_log(ctx, "All OK");
355+
356+
Ok(())
357+
}

‎src/types.rs‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,17 @@ use std::os::raw::{c_int, c_void};
55

66
use crate::raw;
77

8-
pub struct RedisModuleType {
9-
raw_type: RefCell<*mut raw::RedisModuleType>,
8+
pub struct RedisType {
9+
pub raw_type: RefCell<*mut raw::RedisModuleType>,
1010
}
1111

1212
// We want to be able to create static instances of this type,
1313
// which means we need to implement Sync.
14-
unsafe impl Sync for RedisModuleType {}
14+
unsafe impl Sync for RedisType {}
1515

16-
impl RedisModuleType {
16+
impl RedisType {
1717
pub const fn new() -> Self {
18-
RedisModuleType {
18+
RedisType {
1919
raw_type: RefCell::new(ptr::null_mut()),
2020
}
2121
}
@@ -70,7 +70,7 @@ impl RedisModuleType {
7070

7171
// FIXME: Generate these methods with a macro, since we need a set for each custom data type.
7272

73-
#[allow(non_snake_case)]
73+
#[allow(non_snake_case,unused)]
7474
#[no_mangle] // FIXME This should be unneeded
7575
pub unsafe extern "C" fn MyTypeRdbLoad(
7676
rdb: *mut raw::RedisModuleIO,
@@ -80,7 +80,7 @@ pub unsafe extern "C" fn MyTypeRdbLoad(
8080
ptr::null_mut()
8181
}
8282

83-
#[allow(non_snake_case)]
83+
#[allow(non_snake_case,unused)]
8484
#[no_mangle]
8585
pub unsafe extern "C" fn MyTypeRdbSave(
8686
rdb: *mut raw::RedisModuleIO,
@@ -89,7 +89,7 @@ pub unsafe extern "C" fn MyTypeRdbSave(
8989
// eprintln!("MyTypeRdbSave");
9090
}
9191

92-
#[allow(non_snake_case)]
92+
#[allow(non_snake_case,unused)]
9393
#[no_mangle]
9494
pub unsafe extern "C" fn MyTypeAofRewrite(
9595
aof: *mut raw::RedisModuleIO,
@@ -99,14 +99,15 @@ pub unsafe extern "C" fn MyTypeAofRewrite(
9999
// eprintln!("MyTypeAofRewrite");
100100
}
101101

102-
#[allow(non_snake_case)]
102+
#[allow(non_snake_case,unused)]
103103
#[no_mangle]
104104
pub unsafe extern "C" fn MyTypeFree(
105105
value: *mut c_void,
106106
) {
107107
// eprintln!("MyTypeFree");
108108
}
109109

110+
// TODO: Move to raw
110111
pub fn redis_log(
111112
ctx: *mut raw::RedisModuleCtx,
112113
msg: &str,

0 commit comments

Comments
 (0)