Merge pull request #59 from CoLearn-Dev/sm_chunk

Storage macro: chunk compatibility mode, read_keys
Bugfix: lock in vt_inbox

Author: stneng

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "colink"
-version = "0.3.5"
+version = "0.3.6"
 edition = "2021"
 description = "CoLink Rust SDK"
 license = "MIT"

README.md

@@ -9,7 +9,7 @@ CoLink SDK helps both application and protocol developers access the functionali
 Add this to your Cargo.toml:
 ```toml
 [dependencies]
-colink = "0.3.5"
+colink = "0.3.6"
 ```
 
 ## Getting Started

src/application.rs

@@ -203,7 +203,6 @@ impl CoLink {
     }
 
     pub async fn create_entry(&self, key_name: &str, payload: &[u8]) -> Result<String, Error> {
-        let mut client = self._grpc_connect(&self.core_addr).await?;
         if key_name.contains('$') {
             #[cfg(feature = "storage_macro")]
             return self._sm_create_entry(key_name, payload).await;
@@ -215,6 +214,7 @@ impl CoLink {
             .into());
         }
 
+        let mut client = self._grpc_connect(&self.core_addr).await?;
         let request = generate_request(
             &self.jwt,
             StorageEntry {
@@ -269,7 +269,6 @@ impl CoLink {
     }
 
     pub async fn update_entry(&self, key_name: &str, payload: &[u8]) -> Result<String, Error> {
-        let mut client = self._grpc_connect(&self.core_addr).await?;
         if key_name.contains('$') {
             #[cfg(feature = "storage_macro")]
             return self._sm_update_entry(key_name, payload).await;
@@ -281,6 +280,7 @@ impl CoLink {
             .into());
         }
 
+        let mut client = self._grpc_connect(&self.core_addr).await?;
         let request = generate_request(
             &self.jwt,
             StorageEntry {
@@ -295,7 +295,6 @@ impl CoLink {
     }
 
     pub async fn delete_entry(&self, key_name: &str) -> Result<String, Error> {
-        let mut client = self._grpc_connect(&self.core_addr).await?;
         if key_name.contains('$') {
             #[cfg(feature = "storage_macro")]
             return self._sm_delete_entry(key_name).await;
@@ -307,6 +306,7 @@ impl CoLink {
             .into());
         }
 
+        let mut client = self._grpc_connect(&self.core_addr).await?;
         let request = generate_request(
             &self.jwt,
             StorageEntry {
@@ -324,6 +324,17 @@ impl CoLink {
         prefix: &str,
         include_history: bool,
     ) -> Result<Vec<StorageEntry>, Error> {
+        if prefix.contains('$') {
+            #[cfg(feature = "storage_macro")]
+            return self._sm_read_keys(prefix, include_history).await;
+            #[cfg(not(feature = "storage_macro"))]
+            return Err(format!(
+                "Storage Macro feature not enabled, but found $ symbol in key name: {}",
+                key_name
+            )
+            .into());
+        }
+
         let mut client = self._grpc_connect(&self.core_addr).await?;
         let request = generate_request(
             &self.jwt,

src/extensions/storage_macro.rs

@@ -2,6 +2,7 @@ mod append;
 mod chunk;
 mod fs;
 mod redis;
+use crate::StorageEntry;
 
 type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 
@@ -53,6 +54,12 @@ impl crate::application::CoLink {
     }
 
     pub(crate) async fn _sm_read_entry(&self, key_name: &str) -> Result<Vec<u8>, Error> {
+        let key_name = if key_name.contains("::") {
+            &key_name
+                [key_name.find(':').unwrap() + 2..key_name.rfind('@').unwrap_or(key_name.len())]
+        } else {
+            key_name
+        };
         let (string_before, macro_type, string_after) = self._parse_macro(key_name);
         match macro_type.as_str() {
             "chunk" => self._read_entry_chunk(&string_before).await,
@@ -107,4 +114,39 @@ impl crate::application::CoLink {
             .into()),
         }
     }
+
+    pub(crate) async fn _sm_read_keys(
+        &self,
+        prefix: &str,
+        include_history: bool,
+    ) -> Result<Vec<StorageEntry>, Error> {
+        if include_history {
+            return Err("Storage Macro: include_history is not supported.".into());
+        }
+        if !prefix.starts_with(&format!("{}::", self.get_user_id()?)) {
+            return Err("prefix must start with the given user_id".into());
+        }
+        let key_name_prefix = &prefix[prefix.find(':').unwrap() + 2..];
+        let (string_before, macro_type, string_after) = self._parse_macro(key_name_prefix);
+        let key_list = match macro_type.as_str() {
+            "redis" => self._read_keys_redis(&string_before, &string_after).await?,
+            "fs" => self._read_keys_fs(&string_before, &string_after).await?,
+            _ => {
+                return Err(format!(
+                    "invalid storage macro, found {} in prefix {}",
+                    macro_type, key_name_prefix
+                )
+                .into())
+            }
+        };
+        let mut res: Vec<StorageEntry> = Vec::new();
+        for key in key_list {
+            res.push(StorageEntry {
+                key_name: Default::default(),
+                key_path: format!("{}:{}@0", prefix, key),
+                payload: Default::default(),
+            });
+        }
+        Ok(res)
+    }
 }

src/extensions/storage_macro/chunk.rs

@@ -97,13 +97,122 @@ impl crate::application::CoLink {
         Ok(chunk_paths_string)
     }
 
+    async fn _store_chunks_compatibility_mode(
+        &self,
+        payload: &[u8],
+        key_name: &str,
+    ) -> Result<i32, Error> {
+        let mut offset = 0;
+        let mut chunk_id = 0;
+        while offset < payload.len() {
+            let chunk_size = if offset + CHUNK_SIZE > payload.len() {
+                payload.len() - offset
+            } else {
+                CHUNK_SIZE
+            };
+            self.update_entry(
+                &format!("{}:{}", key_name, chunk_id),
+                &payload[offset..offset + chunk_size],
+            )
+            .await?;
+            offset += chunk_size;
+            chunk_id += 1;
+        }
+        Ok(chunk_id)
+    }
+
+    async fn _append_chunks_compatibility_mode(
+        &self,
+        chunk_len: i32,
+        payload: &[u8],
+        key_name: &str,
+    ) -> Result<i32, Error> {
+        let last_chunk_id = chunk_len - 1;
+        let mut last_chunk = self
+            .read_entry(&format!("{}:{}", key_name, last_chunk_id,))
+            .await?;
+        let mut offset = 0;
+        let mut chunk_id = chunk_len;
+        if last_chunk.len() < CHUNK_SIZE {
+            let chunk_size = if payload.len() < CHUNK_SIZE - last_chunk.len() {
+                payload.len()
+            } else {
+                CHUNK_SIZE - last_chunk.len()
+            };
+            last_chunk.append(&mut payload[..chunk_size].to_vec());
+            self.update_entry(&format!("{}:{}", key_name, last_chunk_id), &last_chunk)
+                .await?;
+            offset = chunk_size;
+        }
+        while offset < payload.len() {
+            let chunk_size = if offset + CHUNK_SIZE > payload.len() {
+                payload.len() - offset
+            } else {
+                CHUNK_SIZE
+            };
+            self.update_entry(
+                &format!("{}:{}", key_name, chunk_id),
+                &payload[offset..offset + chunk_size],
+            )
+            .await?;
+            offset += chunk_size;
+            chunk_id += 1;
+        }
+        Ok(chunk_id)
+    }
+
+    async fn _delete_chunks_compatibility_mode(&self, key_name: &str) -> Result<String, Error> {
+        let metadata_key = format!("{}:chunk_metadata", key_name);
+        let chunk_len = self.read_entry(&metadata_key).await?;
+        let chunk_len = String::from_utf8_lossy(&chunk_len).parse::<i32>()?;
+        let res = self.delete_entry(&metadata_key).await?;
+        for i in 0..chunk_len {
+            self.delete_entry(&format!("{}:{}", key_name, i)).await?;
+        }
+        Ok(res)
+    }
+
+    async fn _chunk_lock_compatibility_mode(&self, key_name: &str) -> Result<(), Error> {
+        loop {
+            if self
+                .create_entry(&format!("{}:chunk_lock", key_name), b"")
+                .await
+                .is_ok()
+            {
+                return Ok(());
+            }
+            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+        }
+    }
+
+    async fn _chunk_unlock_compatibility_mode(&self, key_name: &str) -> Result<(), Error> {
+        self.delete_entry(&format!("{}:chunk_lock", key_name))
+            .await?;
+        Ok(())
+    }
+
     #[async_recursion]
     pub(crate) async fn _create_entry_chunk(
         &self,
         key_name: &str,
         payload: &[u8],
     ) -> Result<String, Error> {
         let metadata_key = format!("{}:chunk_metadata", key_name);
+        if key_name.contains('$') {
+            self._chunk_lock_compatibility_mode(key_name).await?;
+            if let Err(e) = self.create_entry(&metadata_key, b"0").await {
+                self._chunk_unlock_compatibility_mode(key_name).await?;
+                return Err(e);
+            }
+            let chunk_len = self
+                ._store_chunks_compatibility_mode(payload, key_name)
+                .await?;
+            let res = self
+                .update_entry(&metadata_key, chunk_len.to_string().as_bytes())
+                .await?;
+            self._chunk_unlock_compatibility_mode(key_name).await?;
+            return Ok(res);
+        }
         // lock the metadata entry to prevent simultaneous writes
         let lock_token = self.lock(&metadata_key).await?;
         // use a closure to prevent locking forever caused by errors
@@ -114,7 +223,7 @@ impl crate::application::CoLink {
             let chunk_paths_string = self._check_chunk_paths_size(chunk_paths)?;
             // store the chunk paths in the metadata entry and update metadata
             let response = self
-                .create_entry(&metadata_key, &chunk_paths_string.into_bytes())
+                .create_entry(&metadata_key, chunk_paths_string.as_bytes())
                 .await?;
             Ok::<String, Error>(response)
         }
@@ -126,6 +235,22 @@ impl crate::application::CoLink {
     #[async_recursion]
     pub(crate) async fn _read_entry_chunk(&self, key_name: &str) -> Result<Vec<u8>, Error> {
         let metadata_key = format!("{}:chunk_metadata", key_name);
+        if key_name.contains('$') {
+            self._chunk_lock_compatibility_mode(key_name).await?;
+            let res = async {
+                let chunk_len = self.read_entry(&metadata_key).await?;
+                let chunk_len = String::from_utf8_lossy(&chunk_len).parse::<i32>()?;
+                let mut payload = Vec::new();
+                for i in 0..chunk_len {
+                    let mut res = self.read_entry(&format!("{}:{}", key_name, i)).await?;
+                    payload.append(&mut res);
+                }
+                Ok::<Vec<u8>, Error>(payload)
+            }
+            .await;
+            self._chunk_unlock_compatibility_mode(key_name).await?;
+            return res;
+        }
         let metadata_response = self.read_entry(&metadata_key).await?;
         let payload_string = String::from_utf8(metadata_response)?;
         let user_id = self.get_user_id()?;
@@ -149,6 +274,18 @@ impl crate::application::CoLink {
         payload: &[u8],
     ) -> Result<String, Error> {
         let metadata_key = format!("{}:chunk_metadata", key_name);
+        if key_name.contains('$') {
+            self._chunk_lock_compatibility_mode(key_name).await?;
+            let _ = self._delete_chunks_compatibility_mode(key_name).await;
+            let chunk_len = self
+                ._store_chunks_compatibility_mode(payload, key_name)
+                .await?;
+            let res = self
+                .update_entry(&metadata_key, chunk_len.to_string().as_bytes())
+                .await?;
+            self._chunk_unlock_compatibility_mode(key_name).await?;
+            return Ok(res);
+        }
         // lock the metadata entry to prevent simultaneous writes
         let lock_token = self.lock(&metadata_key).await?;
         // use a closure to prevent locking forever caused by errors
@@ -159,7 +296,7 @@ impl crate::application::CoLink {
             let chunk_paths_string = self._check_chunk_paths_size(chunk_paths)?;
             // update the metadata entry
             let response = self
-                .update_entry(&metadata_key, &chunk_paths_string.into_bytes())
+                .update_entry(&metadata_key, chunk_paths_string.as_bytes())
                 .await?;
             Ok::<String, Error>(response)
         }
@@ -175,6 +312,19 @@ impl crate::application::CoLink {
         payload: &[u8],
     ) -> Result<String, Error> {
         let metadata_key = format!("{}:chunk_metadata", key_name);
+        if key_name.contains('$') {
+            self._chunk_lock_compatibility_mode(key_name).await?;
+            let chunk_len = self.read_entry(&metadata_key).await?;
+            let chunk_len = String::from_utf8_lossy(&chunk_len).parse::<i32>()?;
+            let new_chunk_len = self
+                ._append_chunks_compatibility_mode(chunk_len, payload, key_name)
+                .await?;
+            let res = self
+                .update_entry(&metadata_key, new_chunk_len.to_string().as_bytes())
+                .await?;
+            self._chunk_unlock_compatibility_mode(key_name).await?;
+            return Ok(res);
+        }
         // lock the metadata entry to prevent simultaneous writes
         let lock_token = self.lock(&metadata_key).await?;
         // use a closure to prevent locking forever caused by errors
@@ -189,7 +339,7 @@ impl crate::application::CoLink {
             let chunk_paths_string = self._check_chunk_paths_size(chunk_paths)?;
             // update the metadata entry
             let response = self
-                .update_entry(&metadata_key, &chunk_paths_string.into_bytes())
+                .update_entry(&metadata_key, chunk_paths_string.as_bytes())
                 .await?;
             Ok::<String, Error>(response)
         }
@@ -200,6 +350,12 @@ impl crate::application::CoLink {
 
     #[async_recursion]
     pub(crate) async fn _delete_entry_chunk(&self, key_name: &str) -> Result<String, Error> {
+        if key_name.contains('$') {
+            self._chunk_lock_compatibility_mode(key_name).await?;
+            let res = self._delete_chunks_compatibility_mode(key_name).await;
+            self._chunk_unlock_compatibility_mode(key_name).await?;
+            return res;
+        }
         let metadata_key = format!("{}:chunk_metadata", key_name);
         let lock_token = self.lock(&metadata_key).await?;
         let res = self.delete_entry(&metadata_key).await;

src/extensions/storage_macro/fs.rs

@@ -17,7 +17,6 @@ impl crate::application::CoLink {
             path += "/";
             path += &path_suffix;
         }
-        println!("path: {}", path);
         let path = PathBuf::from(path);
         tokio::fs::create_dir_all(path.parent().unwrap()).await?;
         Ok(path)
@@ -97,4 +96,19 @@ impl crate::application::CoLink {
         tokio::fs::remove_file(path).await?;
         Ok("ok".to_string())
     }
+
+    #[async_recursion]
+    pub(crate) async fn _read_keys_fs(
+        &self,
+        path_key_name: &str,
+        prefix: &str,
+    ) -> Result<Vec<String>, Error> {
+        let path = self._sm_fs_get_path(path_key_name, prefix).await?;
+        let mut key_list: Vec<String> = Vec::new();
+        let mut dir = tokio::fs::read_dir(path).await?;
+        while let Some(entry) = dir.next_entry().await? {
+            key_list.push(entry.file_name().to_string_lossy().to_string());
+        }
+        Ok(key_list)
+    }
 }