Add batch insert operations to storage system

modules/tests/test_db/harness/harness.sh

@@ -62,6 +62,10 @@ curl -d "delete_check_returned_data:some_key:a" http://$PLAID_LOCATION/webhook/$
 curl -d "insert:some_key:some_value" http://$PLAID_LOCATION/webhook/$URL
 curl -d "delete_check_returned_data:some_key:some_value" http://$PLAID_LOCATION/webhook/$URL
 # the DB is empty
+curl -d "insert_batch:key1=value1|key2=value2|key3=value3" http://$PLAID_LOCATION/webhook/$URL
+curl -d "get:key1" http://$PLAID_LOCATION/webhook/$URL
+curl -d "get:key2" http://$PLAID_LOCATION/webhook/$URL
+curl -d "get:key3" http://$PLAID_LOCATION/webhook/$URL
 
 sleep 2
 
@@ -83,8 +87,11 @@ kill $RH_PID 2>&1 > /dev/null
 # OK
 # OK
 # OK
+# value1
+# value2
+# value3
 
-echo -e "Empty\nEmpty\nfirst_value\nEmpty\nsecond_value\nEmpty\nEmpty\nEmpty\na\naaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\nOK\nOK\nOK\nOK\nOK" > expected.txt
+echo -e "Empty\nEmpty\nfirst_value\nEmpty\nsecond_value\nEmpty\nEmpty\nEmpty\na\naaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\nOK\nOK\nOK\nOK\nOK\nvalue1\nvalue2\nvalue3" > expected.txt
 diff expected.txt $FILE
 RESULT=$?
 

modules/tests/test_db/src/lib.rs

@@ -26,6 +26,20 @@ fn handle_post(log: &str) -> Result<(), i32> {
         "insert" => {
             plaid::storage::insert(parts[1], parts[2].as_bytes()).unwrap();
         }
+        "insert_batch" => {
+            // Format: insert_batch:key1=value1|key2=value2|...
+            let items: Vec<plaid::storage::Item> = parts[1]
+                .split('|')
+                .map(|pair| {
+                    let (key, value) = pair.split_once('=').expect("invalid key=value pair");
+                    plaid::storage::Item {
+                        key: key.to_string(),
+                        value: value.as_bytes().to_vec(),
+                    }
+                })
+                .collect();
+            plaid::storage::insert_batch(&items).unwrap();
+        }
         "delete" => {
             plaid::storage::delete(parts[1]).unwrap();
         }

modules/tests/test_shared_db_rule_1/harness/harness.sh

@@ -27,6 +27,12 @@ curl -d "delete and check" http://$PLAID_LOCATION/webhook/$URL2
 sleep 2
 curl -d "read after deletion" http://$PLAID_LOCATION/webhook/$URL1
 sleep 2
+curl -d "insert batch and check" http://$PLAID_LOCATION/webhook/$URL2
+sleep 2
+curl -d "read after batch insert" http://$PLAID_LOCATION/webhook/$URL1
+sleep 2
+curl -d "delete batch inserts" http://$PLAID_LOCATION/webhook/$URL2
+sleep 2
 curl -d "fill up the db" http://$PLAID_LOCATION/webhook/$URL2
 sleep 2
 curl -d "write to full db" http://$PLAID_LOCATION/webhook/$URL2
@@ -36,7 +42,7 @@ sleep 2
 
 kill $RH_PID 2>&1 > /dev/null
 
-echo -e "OK\nOK\nOK\nOK\nOK\nOK\nOK\nOK" > expected.txt
+echo -e "OK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK\nOK" > expected.txt
 diff expected.txt $FILE
 RESULT=$?
 

modules/tests/test_shared_db_rule_1/src/lib.rs

@@ -47,6 +47,19 @@ fn main(log: String, _: LogSource) -> Result<(), i32> {
             }
             make_named_request("test-response", "OK", HashMap::new()).unwrap();
         }
+        "read after batch insert" => {
+            for i in 0..3 {
+                let key = format!("key{i}");
+                let expected_value = format!("value{i}");
+
+                let actual_value = plaid::storage::get_shared(SHARED_DB, &key).unwrap();
+
+                if expected_value.as_bytes() != actual_value {
+                    panic!("Return value does not match expected value for key {key}");
+                }
+            }
+            make_named_request("test-response", "OK", HashMap::new()).unwrap();
+        }
         _ => panic!("Got an unexpected log"),
     }
 

modules/tests/test_shared_db_rule_2/src/lib.rs

@@ -1,13 +1,22 @@
 use std::collections::HashMap;
 
-use plaid_stl::{entrypoint_with_source, messages::LogSource, network::make_named_request, plaid};
+use plaid_stl::{
+    entrypoint_with_source,
+    messages::LogSource,
+    network::make_named_request,
+    plaid::{self, storage::Item},
+};
 
 entrypoint_with_source!();
 
 const SHARED_DB: &str = "shared_db_1";
 const RULE_NAME: &str = "test_shared_db_rule_2";
 
+const BATCH_INSERT_SIZE: usize = 3;
+
 fn main(log: String, _: LogSource) -> Result<(), i32> {
+    let batch_insert_items = generate_batch_insert_items();
+
     // Depending on the value of "log", we do different things
     match log.as_str() {
         "write and check" => {
@@ -59,27 +68,93 @@ fn main(log: String, _: LogSource) -> Result<(), i32> {
                 "[{RULE_NAME}] Writing to a full shared DB, should fail..."
             ));
             match plaid::storage::insert_shared(SHARED_DB, "another_key", &vec![0u8]) {
-                Ok(_) => panic!("This should have failed"),
+                Ok(_) => panic!("Single insert on a full DB should have failed"),
+                Err(_) => {
+                    plaid::print_debug_string(&format!(
+                        "[{RULE_NAME}] Failed as expected on single item insert"
+                    ));
+                }
+            }
+            match plaid::storage::insert_batch_shared(SHARED_DB, &batch_insert_items) {
+                Ok(_) => panic!("Batch insert on a full DB should have failed"),
                 Err(_) => {
-                    plaid::print_debug_string(&format!("[{RULE_NAME}] Failed as expected"));
+                    plaid::print_debug_string(&format!(
+                        "[{RULE_NAME}] Failed as expected on batch insert"
+                    ));
                 }
             }
+
             make_named_request("test-response", "OK", HashMap::new()).unwrap();
         }
         "write to non-existing db" => {
             plaid::print_debug_string(&format!(
                 "[{RULE_NAME}] Writing to a non-existing shared DB, should fail..."
             ));
             match plaid::storage::insert_shared("this_does_not_exist", "some_key", &vec![0u8]) {
-                Ok(_) => panic!("This should have failed"),
+                Ok(_) => panic!("Single insert to non-existant DB should have failed"),
                 Err(_) => {
-                    plaid::print_debug_string(&format!("[{RULE_NAME}] Failed as expected"));
+                    plaid::print_debug_string(&format!(
+                        "[{RULE_NAME}] Failed as expected on single item write to non-existent DB"
+                    ));
+                }
+            }
+            match plaid::storage::insert_batch_shared("this_does_not_exist", &batch_insert_items) {
+                Ok(_) => panic!("Batch insert to non-existant DB should have failed"),
+                Err(_) => {
+                    plaid::print_debug_string(&format!(
+                        "[{RULE_NAME}] Failed as expected on batch item write to non-existent DB"
+                    ));
                 }
             }
             make_named_request("test-response", "OK", HashMap::new()).unwrap();
         }
+        "insert batch and check" => {
+            plaid::print_debug_string(&format!(
+                "[{RULE_NAME}] Writing {BATCH_INSERT_SIZE} items to DB..."
+            ));
+
+            plaid::storage::insert_batch_shared(SHARED_DB, &batch_insert_items).unwrap();
+
+            for item in batch_insert_items {
+                let returned_val = plaid::storage::get_shared(SHARED_DB, &item.key).unwrap();
+
+                if returned_val != item.value {
+                    panic!(
+                        "Returned value does not match expected value for key: {}",
+                        item.key
+                    )
+                }
+            }
+            make_named_request("test-response", "OK", HashMap::new()).unwrap();
+        }
+        "delete batch inserts" => {
+            plaid::print_debug_string(&format!(
+                "[{RULE_NAME}] Deleting {BATCH_INSERT_SIZE} items from DB..."
+            ));
+
+            for item in batch_insert_items {
+                plaid::storage::delete_shared(SHARED_DB, &item.key).unwrap();
+            }
+        }
         _ => panic!("Got an unexpected log"),
     }
 
     Ok(())
 }
+
+fn generate_batch_insert_items() -> Vec<Item> {
+    let mut items = vec![];
+    for i in 0..BATCH_INSERT_SIZE {
+        let key = format!("key{i}");
+        let value = format!("value{i}");
+
+        let item = Item {
+            key,
+            value: value.as_bytes().to_vec(),
+        };
+
+        items.push(item)
+    }
+
+    items
+}

runtime/plaid-stl/src/plaid/storage.rs

@@ -1,11 +1,20 @@
 use std::fmt::Display;
 
+use serde::{Deserialize, Serialize};
+
 use crate::PlaidFunctionError;
 
 pub enum StorageError {
     BufferSizeMismatch,
 }
 
+/// A key/value pair for use with batch insert operations.
+#[derive(Serialize, Deserialize)]
+pub struct Item {
+    pub key: String,
+    pub value: Vec<u8>,
+}
+
 pub fn insert(key: &str, value: &[u8]) -> Result<Vec<u8>, PlaidFunctionError> {
     extern "C" {
         /// Send a request to store this data in whatever persistence system Plaid has configured.
@@ -122,6 +131,55 @@ pub fn insert_shared(
     }
 }
 
+/// Insert multiple key/value pairs in a single atomic batch operation.
+pub fn insert_batch(items: &[Item]) -> Result<(), PlaidFunctionError> {
+    extern "C" {
+        fn storage_insert_batch(items_buf: *const u8, items_buf_len: usize) -> i32;
+    }
+
+    let items_json =
+        serde_json::to_vec(items).map_err(|_| PlaidFunctionError::ErrorCouldNotSerialize)?;
+
+    let result = unsafe { storage_insert_batch(items_json.as_ptr(), items_json.len()) };
+
+    if result == 0 {
+        Ok(())
+    } else {
+        Err(result.into())
+    }
+}
+
+/// Insert multiple key/value pairs into a shared namespace in a single atomic batch operation.
+pub fn insert_batch_shared(namespace: &str, items: &[Item]) -> Result<(), PlaidFunctionError> {
+    extern "C" {
+        fn storage_insert_batch_shared(
+            namespace: *const u8,
+            namespace_len: usize,
+            items_buf: *const u8,
+            items_buf_len: usize,
+        ) -> i32;
+    }
+
+    let namespace_bytes = namespace.as_bytes();
+    let items_json =
+        serde_json::to_vec(items).map_err(|_| PlaidFunctionError::ErrorCouldNotSerialize)?;
+
+    let result = unsafe {
+        storage_insert_batch_shared(
+            namespace_bytes.as_ptr(),
+            namespace_bytes.len(),
+            items_json.as_ptr(),
+            items_json.len(),
+        )
+    };
+
+    if result == 0 {
+        Ok(())
+    } else {
+        Err(result.into())
+    }
+}
+
 pub fn get(key: &str) -> Result<Vec<u8>, PlaidFunctionError> {
     extern "C" {
         fn storage_get(key: *const u8, key_len: usize, data: *const u8, data_len: usize) -> i32;

runtime/plaid/src/functions/api.rs

@@ -666,6 +666,12 @@ pub fn to_api_function(
         "storage_insert_shared" => {
             Function::new_typed_with_env(&mut store, &env, super::storage::insert_shared)
         }
+        "storage_insert_batch" => {
+            Function::new_typed_with_env(&mut store, &env, super::storage::insert_batch)
+        }
+        "storage_insert_batch_shared" => {
+            Function::new_typed_with_env(&mut store, &env, super::storage::insert_batch_shared)
+        }
         "storage_get" => Function::new_typed_with_env(&mut store, &env, super::storage::get),
         "storage_get_shared" => {
             Function::new_typed_with_env(&mut store, &env, super::storage::get_shared)