Reallocate when setting existing key

Author: gavrie

examples/data_type.rs

@@ -51,18 +51,36 @@ impl Command for AllocSetCommand {
         // 2. Allocate data [OK]
         // 3. Set the key to the data [OK]
         // 4. Activate custom allocator and compare Redis memory usage [OK]
-        // 5. Handle deallocation of existing value
+        // 5. Handle deallocation of existing value [OK]
 
         let key = r.open_key_writable(key);
-        key.check_type(&MY_TYPE)?;
-
-        // TODO: If there is an existing value, reuse it or deallocate it.
-
-        let my = Box::into_raw(Box::new(
-            MyType {
-                data: "A".repeat(size as usize)
+        let key_type = key.verify_and_get_type(&MY_TYPE)?;
+
+        let my = match key_type {
+            raw::KeyType::Empty => {
+                // Create a new value
+                Box::new(
+                    MyType {
+                        data: "A".repeat(size as usize)
+                    }
+                )
+            }
+            _ => {
+                // There is an existing value, reuse it
+                let my = key.get_value() as *mut MyType;
+
+                if my.is_null() {
+                    r.reply_integer(0)?;
+                    return Ok(());
+                }
+
+                let mut my = unsafe { Box::from_raw(my) };
+                my.data = "B".repeat(size as usize);
+                my
             }
-        ));
+        };
+
+        let my = Box::into_raw(my);
 
         key.set_value(&MY_TYPE, my as *mut c_void)?;
         r.reply_integer(size)?;
@@ -105,9 +123,7 @@ impl Command for AllocGetCommand {
         let key = args[1];
 
         let key = r.open_key(key);
-
-        key.check_type(&MY_TYPE)?;
-
+        key.verify_and_get_type(&MY_TYPE)?;
         let my = key.get_value() as *mut MyType;
 
         if my.is_null() {
@@ -179,7 +195,7 @@ pub extern "C" fn AllocDelCommand_Redis(
 fn module_on_load(ctx: *mut raw::RedisModuleCtx) -> Result<(), &'static str> {
     module_init(ctx, MODULE_NAME, MODULE_VERSION)?;
 
-    // TODO: Call this from inside module_init
+// TODO: Call this from inside module_init
     redismodule::use_redis_alloc();
 
     MY_TYPE.create_data_type(ctx, "mytype123")?;

src/lib.rs

@@ -4,6 +4,7 @@
 extern crate bitflags;
 
 pub mod redisraw;
+pub mod error;
 pub mod raw;
 pub mod types;
 
@@ -34,7 +35,6 @@ extern crate num_traits;
 #[macro_use]
 mod macros;
 
-pub mod error;
 
 use std::alloc::{System, GlobalAlloc, Layout};
 use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering::SeqCst};
@@ -254,10 +254,6 @@ pub enum KeyMode {
     ReadWrite,
 }
 
-trait RedisKeyType {
-    fn check_type(&self, redis_type: &RedisType) -> Result<(), Error>;
-}
-
 /// `RedisKey` is an abstraction over a Redis key that allows readonly
 /// operations.
 ///
@@ -299,15 +295,12 @@ impl RedisKey {
         Ok(val)
     }
 
-    pub fn check_type(&self, redis_type: &RedisType) -> Result<(), Error> {
-        match raw::check_key_type(
+    pub fn verify_and_get_type(&self, redis_type: &RedisType) -> Result<raw::KeyType, Error> {
+        raw::verify_and_get_type(
             self.ctx,
             self.key_inner,
             *redis_type.raw_type.borrow_mut(),
-        ) {
-            Ok(_) => Ok(()),
-            Err(s) => Err(Error::generic(s))
-        }
+        )
     }
 
     pub fn get_value(&self) -> *mut c_void {
@@ -386,15 +379,16 @@ impl RedisKeyWritable {
         }
     }
 
-    pub fn check_type(&self, redis_type: &RedisType) -> Result<(), Error> {
-        match raw::check_key_type(
+    pub fn verify_and_get_type(&self, redis_type: &RedisType) -> Result<raw::KeyType, Error> {
+        raw::verify_and_get_type(
             self.ctx,
             self.key_inner,
             *redis_type.raw_type.borrow_mut(),
-        ) {
-            Ok(_) => Ok(()),
-            Err(s) => Err(Error::generic(s))
-        }
+        )
+    }
+
+    pub fn get_value(&self) -> *mut c_void {
+        raw::module_type_get_value(self.key_inner)
     }
 
     pub fn set_value(&self, redis_type: &RedisType, value: *mut c_void) -> Result<(), Error> {

src/raw.rs

@@ -15,6 +15,7 @@ use libc::size_t;
 
 pub use crate::redisraw::bindings::*;
 use crate::types::redis_log;
+use crate::error::Error;
 
 bitflags! {
     pub struct KeyMode: c_int {
@@ -332,11 +333,11 @@ pub fn module_type_set_value(
     }.into()
 }
 
-pub fn check_key_type(
+pub fn verify_and_get_type(
     ctx: *mut RedisModuleCtx,
     key: *mut RedisModuleKey,
     redis_type: *mut RedisModuleType,
-) -> Result<(), &'static str> {
+) -> Result<KeyType, Error> {
 
     let key_type = key_type(key);
 
@@ -346,12 +347,12 @@ pub fn check_key_type(
     if key_type != KeyType::Empty {
         let raw_type = module_type_get_type(key);
         if raw_type != redis_type{
-            return Err("Key has existing value with wrong Redis type");
+            return Err(Error::generic("Key has existing value with wrong Redis type"));
         }
         redis_log(ctx, "Existing key has the correct type");
     }
 
     redis_log(ctx, "All OK");
 
-    Ok(())
+    Ok(key_type)
 }