Merge pull request #27 from nociza/magic_mount

Storage Macro: Chunk

Author: stneng

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "colink"
-version = "0.2.4"
+version = "0.2.5"
 edition = "2021"
 description = "CoLink Rust SDK"
 license = "MIT"
@@ -10,6 +10,7 @@ documentation = "https://docs.rs/colink"
 repository = "https://github.com/CoLearn-Dev/colink-sdk-rust-dev"
 
 [dependencies]
+async-recursion = "1.0.0"
 async-trait = "0.1"
 base64 = "0.13.0"
 chrono = "0.4"
@@ -33,10 +34,11 @@ prost-build = "0.10"
 tonic-build = "0.7"
 
 [features]
-default = ["extensions", "remote_storage", "variable_transfer", "registry", "policy_module", "instant_server"]
+default = ["extensions", "remote_storage", "variable_transfer", "registry", "policy_module", "instant_server", "storage_macro"]
 extensions = []
 remote_storage = ["extensions"]
 variable_transfer = ["extensions", "remote_storage"]
 registry = []
 policy_module = []
 instant_server = []
+storage_macro = []

README.md

@@ -9,7 +9,7 @@ CoLink SDK helps both application adnd protocol developers access the functional
 Add this to your Cargo.toml:
 ```toml
 [dependencies]
-colink = "0.2.4"
+colink = "0.2.5"
 ```
 
 ## Getting Started
@@ -91,6 +91,9 @@ cargo run --example user_stop_protocol_operator <address> <user_jwt> <instance_i
 ```
 cargo run --example user_wait_task <address> <user_jwt> <target_user_id>
 ```
+```
+cargo run --example storage_macro_chunk <address> <user_jwt> <payload_size>
+```
 
 ### Protocol
 ```

examples/storage_macro_chunk.rs

@@ -0,0 +1,64 @@
+use colink::{decode_jwt_without_validation, CoLink};
+use rand::Rng;
+use std::env;
+
+const CHUNK_SIZE: usize = 1024 * 1024; // use 1MB chunks
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
+    let args = env::args().skip(1).collect::<Vec<_>>();
+    let addr = &args[0];
+    let jwt = &args[1];
+    let length = if args.len() > 2 {
+        args[2].parse::<usize>().unwrap()
+    } else {
+        5e6 as usize // default to 5 * 10^6 bytes
+    };
+    let user_id = decode_jwt_without_validation(jwt).unwrap().user_id;
+    println!("user_id: {}", user_id);
+    let cl = CoLink::new(addr, jwt);
+    let key_name = "storage_macro_demo:$chunk";
+    let payload = rand::thread_rng()
+        .sample_iter(&rand::distributions::Standard)
+        .take(length)
+        .collect::<Vec<u8>>();
+
+    // create
+    println!("Creating entry...");
+    let response = cl.create_entry(key_name, &payload.clone()).await?;
+    println!("created entry at key name: {}", response);
+
+    // read
+    println!("Reading entry...");
+    let data = cl.read_entry(key_name).await?;
+    assert_eq!(data, payload);
+    println!(
+        "Read payload of {}MB ({} bytes), verified to be same as bytes written",
+        payload.len() as f32 / CHUNK_SIZE as f32,
+        payload.len()
+    );
+
+    // update
+    println!("Updating entry...");
+    let new_payload = rand::thread_rng()
+        .sample_iter(&rand::distributions::Standard)
+        .take(length / 2)
+        .collect::<Vec<u8>>();
+    let response = cl.update_entry(key_name, &new_payload.clone()).await?;
+    println!("updated entry at key name: {}", response);
+
+    // read again to verify
+    println!("Reading entry again...");
+    let data = cl.read_entry(key_name).await?;
+    assert_eq!(data, new_payload);
+    println!(
+        "Read payload of {}MB ({} bytes), verified to be same as the updated payload bytes",
+        new_payload.len() as f32 / CHUNK_SIZE as f32,
+        new_payload.len()
+    );
+
+    // delete
+    println!("Deleting entry...");
+    cl.delete_entry(key_name).await?;
+    Ok(())
+}

src/application.rs

@@ -159,6 +159,17 @@ 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;
+            #[cfg(not(feature = "storage_macro"))]
+            return Err(format!(
+                "Storage Macro feature not enabled, but found $ symbol in key name: {}",
+                key_name
+            )
+            .into());
+        }
+
         let request = generate_request(
             &self.jwt,
             StorageEntry {
@@ -186,6 +197,17 @@ impl CoLink {
     }
 
     pub async fn read_entry(&self, key: &str) -> Result<Vec<u8>, Error> {
+        if key.contains('$') {
+            #[cfg(feature = "storage_macro")]
+            return self._sm_read_entry(key).await;
+            #[cfg(not(feature = "storage_macro"))]
+            return Err(format!(
+                "Storage Macro feature not enabled, but found $ symbol in key name: {}",
+                key
+            )
+            .into());
+        }
+
         let storage_entry = if key.contains("::") {
             StorageEntry {
                 key_path: key.to_string(),
@@ -203,6 +225,17 @@ 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;
+            #[cfg(not(feature = "storage_macro"))]
+            return Err(format!(
+                "Storage Macro feature not enabled, but found $ symbol in key name: {}",
+                key_name
+            )
+            .into());
+        }
+
         let request = generate_request(
             &self.jwt,
             StorageEntry {
@@ -218,6 +251,17 @@ 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;
+            #[cfg(not(feature = "storage_macro"))]
+            return Err(format!(
+                "Storage Macro feature not enabled, but found $ symbol in key name: {}",
+                key_name
+            )
+            .into());
+        }
+
         let request = generate_request(
             &self.jwt,
             StorageEntry {

src/extensions.rs

@@ -12,6 +12,8 @@ mod read_or_wait;
 pub mod registry;
 #[cfg(feature = "remote_storage")]
 mod remote_storage;
+#[cfg(feature = "storage_macro")]
+mod storage_macro;
 #[cfg(feature = "extensions")]
 mod switch_to_generated_user;
 #[cfg(feature = "variable_transfer")]

src/extensions/storage_macro.rs

@@ -0,0 +1,89 @@
+mod chunk;
+
+type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
+
+impl crate::application::CoLink {
+    fn _parse_macro(&self, key_name: &str) -> (String, String, String) {
+        let split_key = key_name.split(':').collect::<Vec<&str>>();
+        let mut macro_type = String::new();
+        for s in split_key.iter().rev() {
+            if s.contains('$') {
+                macro_type = s.to_string();
+                break;
+            }
+        }
+        let macro_type_splitter = format!(":{}", macro_type);
+        let split_by_macro = key_name.split(&macro_type_splitter).collect::<Vec<&str>>();
+        let mut string_after = split_by_macro[split_by_macro.len() - 1].to_string();
+        if string_after.starts_with(':') {
+            string_after = string_after[1..].to_string();
+        }
+        (
+            split_by_macro[0..(split_by_macro.len() - 1)].join(&macro_type_splitter),
+            macro_type.replace('$', ""),
+            string_after,
+        )
+    }
+
+    pub(crate) async fn _sm_create_entry(
+        &self,
+        key_name: &str,
+        payload: &[u8],
+    ) -> Result<String, Error> {
+        let (string_before, macro_type, _) = self._parse_macro(key_name);
+        match macro_type.as_str() {
+            "chunk" => {
+                self._create_entry_chunk(string_before.as_str(), payload)
+                    .await
+            }
+            _ => Err(format!(
+                "invalid storage macro, found {} in key name {}",
+                macro_type, key_name
+            )
+            .into()),
+        }
+    }
+
+    pub(crate) async fn _sm_read_entry(&self, key_name: &str) -> Result<Vec<u8>, Error> {
+        let (string_before, macro_type, _) = self._parse_macro(key_name);
+        match macro_type.as_str() {
+            "chunk" => self._read_entry_chunk(string_before.as_str()).await,
+            _ => Err(format!(
+                "invalid storage macro, found {} in key name {}",
+                macro_type, key_name
+            )
+            .into()),
+        }
+    }
+
+    pub(crate) async fn _sm_update_entry(
+        &self,
+        key_name: &str,
+        payload: &[u8],
+    ) -> Result<String, Error> {
+        let (string_before, macro_type, _) = self._parse_macro(key_name);
+        match macro_type.as_str() {
+            "chunk" => {
+                self._update_entry_chunk(string_before.as_str(), payload)
+                    .await
+            }
+            _ => Err(format!(
+                "invalid storage macro, found {} in key name {}",
+                macro_type, key_name
+            )
+            .into()),
+        }
+    }
+
+    pub(crate) async fn _sm_delete_entry(&self, key_name: &str) -> Result<String, Error> {
+        let (string_before, macro_type, _) = self._parse_macro(key_name);
+        match macro_type.as_str() {
+            "chunk" => self._delete_entry_chunk(string_before.as_str()).await,
+            _ => Err(format!(
+                "invalid storage macro, found {} in key name {}",
+                macro_type, key_name
+            )
+            .into()),
+        }
+    }
+}

src/extensions/storage_macro/chunk.rs

@@ -0,0 +1,114 @@
+use crate::decode_jwt_without_validation;
+use async_recursion::async_recursion;
+
+const CHUNK_SIZE: usize = 1024 * 1024; // use 1MB chunks
+
+type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
+
+impl crate::application::CoLink {
+    #[async_recursion]
+    async fn _store_chunks(&self, payload: &[u8], key_name: &str) -> Result<Vec<String>, Error> {
+        let mut offset = 0;
+        let mut chunk_id = 0;
+        let mut chunk_paths = Vec::new();
+        while offset < payload.len() {
+            let chunk_size = if offset + CHUNK_SIZE > payload.len() {
+                payload.len() - offset
+            } else {
+                CHUNK_SIZE
+            };
+            let response = self
+                .update_entry(
+                    &format!("{}:{}", key_name, chunk_id),
+                    &payload[offset..offset + chunk_size],
+                )
+                .await?;
+            chunk_paths.push(response.split('@').last().unwrap().to_string()); // only store the timestamps
+            offset += chunk_size;
+            chunk_id += 1;
+        }
+        Ok(chunk_paths)
+    }
+
+    fn _check_chunk_paths_size(&self, chunk_paths: Vec<String>) -> Result<String, Error> {
+        let chunk_paths_string = chunk_paths.join(";");
+        if chunk_paths_string.len() > CHUNK_SIZE {
+            return Err(format!(
+                "File too large: failed to store {} chunks references in metadata",
+                chunk_paths.len()
+            )
+            .into());
+        }
+        Ok(chunk_paths_string)
+    }
+
+    #[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);
+        // lock the metadata entry to prevent simultaneous writes
+        let lock_token = self.lock(&metadata_key.clone()).await?;
+        // create the chunks and store them
+        let chunk_paths = self._store_chunks(payload, key_name).await?;
+        // make sure that the chunk paths are smaller than the maximum entry size
+        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.clone(), &chunk_paths_string.into_bytes())
+            .await?;
+        self.unlock(lock_token).await?;
+        Ok(response)
+    }
+
+    #[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);
+        let metadata_response = self.read_entry(&metadata_key.clone()).await?;
+        let payload_string = String::from_utf8(metadata_response.clone())?;
+        let user_id = decode_jwt_without_validation(&self.jwt).unwrap().user_id;
+
+        // read the chunks into a single vector
+        let chunks_paths = payload_string.split(';').collect::<Vec<&str>>();
+        let mut payload = Vec::new();
+        for (i, timestamp) in chunks_paths.iter().enumerate() {
+            let response = self
+                .read_entry(&format!("{}::{}:{}@{}", user_id, key_name, i, timestamp))
+                .await?;
+            payload.append(&mut response.clone());
+        }
+        Ok(payload)
+    }
+
+    #[async_recursion]
+    pub(crate) async fn _update_entry_chunk(
+        &self,
+        key_name: &str,
+        payload: &[u8],
+    ) -> Result<String, Error> {
+        let metadata_key = format!("{}:chunk_metadata", key_name);
+        // lock the metadata entry to prevent simultaneous writes
+        let lock_token = self.lock(&metadata_key.clone()).await?;
+        // split payload into chunks and update the chunks
+        let chunk_paths = self._store_chunks(payload, key_name).await?;
+        // make sure that the chunk paths are smaller than the maximum entry size
+        let chunk_paths_string = self._check_chunk_paths_size(chunk_paths)?;
+        // update the metadata entry
+        let response = self
+            .update_entry(&metadata_key.clone(), &chunk_paths_string.into_bytes())
+            .await?;
+        self.unlock(lock_token).await?;
+        Ok(response)
+    }
+
+    #[async_recursion]
+    pub(crate) async fn _delete_entry_chunk(&self, key_name: &str) -> Result<String, Error> {
+        let metadata_key = format!("{}:chunk_metadata", key_name);
+        let lock_token = self.lock(&metadata_key.clone()).await?;
+        let response = self.delete_entry(&metadata_key.clone()).await?;
+        self.unlock(lock_token).await?;
+        Ok(response)
+    }
+}