perf(su): new database impl partially complete #987

servers/su/src/domain/clients/local_store.rs

@@ -0,0 +1,349 @@
+use super::super::core::dal::{
+ DataStore, Log, Message, PaginatedMessages, Process, StoreErrorType,
+};
+use super::super::SuLog;
+use crate::domain::config::AoConfig;
+use async_trait::async_trait;
+use rocksdb::{Options, DB};
+use std::sync::Arc;
+
+pub struct LocalStoreClient {
+ logger: Arc<dyn Log>,
+ /*
+ A RocksDB instance that is a key value store
+ of ANS-104 bundles
+ */
+ file_db: DB,
+ /*
+ A RocksDB instance that is an index for querying
+ and ordering of Process's and Messages
+ */
+ index_db: DB,
+}
+
+impl From<serde_json::Error> for StoreErrorType {
+ fn from(error: serde_json::Error) -> Self {
+ StoreErrorType::JsonError(format!("data store json error: {}", error))
+ }
+}
+
+impl From<StoreErrorType> for String {
+ fn from(error: StoreErrorType) -> Self {
+ format!("{:?}", error)
+ }
+}
+
+impl From<String> for StoreErrorType {
+ fn from(error: String) -> Self {
+ StoreErrorType::DatabaseError(format!("{:?}", error))
+ }
+}
+
+impl From<rocksdb::Error> for StoreErrorType {
+ fn from(err: rocksdb::Error) -> StoreErrorType {
+ StoreErrorType::DatabaseError(format!("RocksDB error: {:?}", err))
+ }
+}
+
+impl From<std::string::FromUtf8Error> for StoreErrorType {
+ fn from(err: std::string::FromUtf8Error) -> StoreErrorType {
+ StoreErrorType::JsonError(format!("UTF-8 conversion error: {:?}", err))
+ }
+}
+
+impl LocalStoreClient {
+ pub fn new() -> Result<Self, StoreErrorType> {
+ let config = AoConfig::new(Some("su".to_string())).expect("Failed to read configuration");
+ let logger = SuLog::init();
+
+ let mut opts = Options::default();
+ opts.create_if_missing(true);
+ opts.set_enable_blob_files(true);
+ opts.set_blob_file_size(5 * 1024 * 1024 * 1024); // 5GB max
+ opts.set_min_blob_size(1024); // Low value ensures it is used
+
+ let file_db = DB::open(&opts, &config.su_file_db_dir)?;
+
+ let prefix_extractor = rocksdb::SliceTransform::create("id_extractor", 
+ |key| {
+ // Find the first and second colon and return the id part
+ if let Some(first_colon_pos) = key.iter().position(|&b| b == b':') {
+ if let Some(second_colon_pos) = key[first_colon_pos + 1..].iter().position(|&b| b == b':') {
+ &key[first_colon_pos + 1..first_colon_pos + 1 + second_colon_pos]
+ } else {
+ &key[first_colon_pos + 1..] // If no second colon, return rest of the key
+ }
+ } else {
+ key // Return the entire key if no colon is found
+ }
+ },
+ None // Optional in-domain function for validating if the key is within the domain
+ );
+
+ let mut opts_index = Options::default();
+ opts_index.create_if_missing(true);
+ opts_index.set_prefix_extractor(prefix_extractor);
+
+ let index_db = DB::open(&opts_index, &config.su_index_db_dir)?;
+
+ Ok(LocalStoreClient {
+ logger,
+ file_db,
+ index_db,
+ })
+ }
+}
+
+#[async_trait]
+impl DataStore for LocalStoreClient {
+ fn save_process(&self, process: &Process, bundle: &[u8]) -> Result<String, StoreErrorType> {
+ let process_id = &process.process.process_id;
+ let assignment_id = process.assignment_id()?;
+ let timestamp = process.timestamp()?;
+ let epoch = process.epoch()?;
+ let nonce = process.nonce()?;
+
+ // Store the process data by assignment_id (main key)
+ let assignment_key = format!("process_assignments:{}", assignment_id);
+ let process_data = serde_json::to_string(process)?;
+
+ // Save the process data by assignment_id
+ self.index_db
+ .put(assignment_key.as_bytes(), process_data.as_bytes())?;
+
+ // Save by process_id, but only store a reference to the assignment_id
+ let process_key = format!("process:{}", process_id);
+ self.index_db
+ .put(process_key.as_bytes(), assignment_id.as_bytes())?;
+
+ // Store by process_id, epoch, nonce, and timestamp, storing a reference to assignment_id
+ let process_order_key = format!(
+ "processes:{}:{}:{}:{}:{}",
+ process_id, epoch, nonce, timestamp, assignment_id
+ );
+
+ self.index_db
+ .put(process_order_key.as_bytes(), assignment_id.as_bytes())?;
+
+ // Store the binary bundle in file_db (by assignment_id)
+ self.file_db
+ .put(assignment_key.as_bytes(), bundle)?;
+
+ Ok("Process saved".to_string())
+ }
+
+ async fn get_process(&self, tx_id: &str) -> Result<Process, StoreErrorType> {
+ // First, try to fetch the process by assignment_id directly
+ let assignment_key = format!("process_assignments:{}", tx_id);
+ if let Some(process_data) = self.index_db.get(assignment_key.as_bytes())? {
+ // Found the process by assignment_id, deserialize and return it
+ let process: Process = serde_json::from_slice(&process_data)?;
+ return Ok(process);
+ }
+
+ // If not found by assignment_id, assume tx_id is a process_id
+ let process_key = format!("process:{}", tx_id);
+ if let Some(assignment_id_bytes) = self.index_db.get(process_key.as_bytes())? {
+ let assignment_id = String::from_utf8(assignment_id_bytes.to_vec())?;
+
+ // Now fetch the process by assignment_id
+ let assignment_key = format!("process_assignments:{}", assignment_id);
+ if let Some(process_data) = self.index_db.get(assignment_key.as_bytes())? {
+ let process: Process = serde_json::from_slice(&process_data)?;
+ return Ok(process);
+ }
+ }
+
+ Err(StoreErrorType::NotFound("Process not found".to_string()))
+ }
+
+ fn check_existing_message(&self, message_id: &String) -> Result<(), StoreErrorType> {
+ if let Ok(_message) = self.get_message(message_id) {
+ Err(StoreErrorType::MessageExists(
+ "Message already exists".to_string(),
+ ))
+ } else {
+ Ok(())
+ }
+ }
+
+ async fn save_message(
+ &self,
+ message: &Message,
+ bundle_in: &[u8],
+ ) -> Result<String, StoreErrorType> {
+ let process_id = message.process_id()?;
+ let message_id = message.message_id()?;
+ let assignment_id = message.assignment_id()?;
+ let timestamp = message.timestamp()?;
+ let epoch = message.epoch()?;
+ let nonce = message.nonce()?;
+
+ // Store the message data in the index_db, keyed by assignment_id (main key)
+ let assignment_key = format!("assignments:{}", assignment_id);
+ let message_data = serde_json::to_string(message)?;
+
+ // Save the actual message data by assignment_id (main unique key)
+ self.index_db
+ .put(assignment_key.as_bytes(), message_data.as_bytes())?;
+
+ // Save by message_id, but only store a reference to the assignment_id
+ let message_composite_key = format!("messages:{}:{}", message_id, assignment_id);
+ self.index_db
+ .put(message_composite_key.as_bytes(), assignment_id.as_bytes())?;
+
+ // Store by process_id, epoch, nonce, and timestamp, storing a reference to assignment_id
+ let process_order_key = format!(
+ "process_messages:{}:{}:{}:{}:{}",
+ process_id, epoch, nonce, timestamp, assignment_id
+ );
+ self.index_db
+ .put(process_order_key.as_bytes(), assignment_id.as_bytes())?;
+
+ // Store the binary bundle in file_db (by assignment_id)
+ self.file_db
+ .put(assignment_key.as_bytes(), bundle_in)?;
+
+ Ok("Message saved".to_string())
+ }
+
+ fn get_message(&self, tx_id: &str) -> Result<Message, StoreErrorType> {
+ // Try to fetch the message directly by assignment_id
+ let assignment_key = format!("assignments:{}", tx_id);
+ if let Some(message_data) = self.index_db.get(assignment_key.as_bytes())? {
+ // Found the message by assignment_id, deserialize and return it
+ let message: Message = serde_json::from_slice(&message_data)?;
+ return Ok(message);
+ }
+
+ // If not found by assignment_id, assume tx_id is a message_id
+ let message_key_prefix = format!("messages:{}:", tx_id);
+ let mut iter = self.index_db.prefix_iterator(message_key_prefix.as_bytes());
+
+ // Look for the assignment_id for this message_id
+ if let Some(result) = iter.next() {
+ let (key, assignment_id_bytes) = result?;
+ let assignment_id = String::from_utf8(assignment_id_bytes.to_vec())?;
+
+ // Now fetch the message by assignment_id
+ let assignment_key = format!("assignments:{}", assignment_id);
+ if let Some(message_data) = self.index_db.get(assignment_key.as_bytes())? {
+ let message: Message = serde_json::from_slice(&message_data)?;
+ return Ok(message);
+ }
+ }
+
+ Err(StoreErrorType::NotFound("Message not found".to_string()))
+ }
+
+ async fn get_messages(
+ &self,
+ process_in: &Process,
+ from: &Option<String>,
+ to: &Option<String>,
+ limit: &Option<i32>,
+ ) -> Result<PaginatedMessages, StoreErrorType> {
+ let process_id = &process_in.process.process_id;
+ let limit_val = limit.unwrap_or(100) as usize;
+
+ let mut messages = Vec::new();
+ let mut actual_limit = limit_val;
+
+ // Check if the process has an assignment and should be included as the first message
+ let include_process = process_in.assignment.is_some()
+ && match from {
+ Some(ref from_nonce) => from_nonce == &process_in.nonce()?.to_string(),
+ None => true, // No 'from' means it's the first page
+ };
+
+ if include_process {
+ // Add the process as the first message
+ let process_message = Message::from_process(process_in.clone())?;
+ messages.push(process_message);
+ actual_limit -= 1; // Adjust the limit since we're adding the process
+ }
+
+ let process_key_prefix = format!("process_messages:{}:", process_id);
+
+ // Use an iterator to scan through the keys for the Process
+ let iter = self.index_db.prefix_iterator(process_key_prefix.as_bytes());
+
+ let mut keys_with_ids = Vec::new();
+
+ // Collect all keys into memory
+ for item in iter {
+ let (key, assignment_id_bytes) = item?;
+ let key_str = String::from_utf8(key.to_vec())?;
+
+ // Extract nonce, epoch, and timestamp from the key for sorting
+ let parts: Vec<&str> = key_str.split(':').collect();
+ if parts.len() >= 4 {
+ // Store the key along with the assignment_id
+ let assignment_id = String::from_utf8(assignment_id_bytes.to_vec())?;
+ keys_with_ids.push((key_str, assignment_id));
+ }
+ }
+
+ // Sort the keys by epoch, nonce, and timestamp
+ keys_with_ids.sort_by(|(key_a, _), (key_b, _)| {
+ let parts_a: Vec<&str> = key_a.split(':').collect();
+ let parts_b: Vec<&str> = key_b.split(':').collect();
+
+ // Assuming keys are formatted as "process_messages:<process_id>:<nonce>:<epoch>:<timestamp>"
+ let epoch_a = parts_a.get(3).unwrap_or(&"0").parse::<i64>().unwrap_or(0);
+ let nonce_a = parts_a.get(2).unwrap_or(&"0").parse::<i64>().unwrap_or(0);
+ let timestamp_a = parts_a.get(4).unwrap_or(&"0").parse::<i64>().unwrap_or(0);
+
+ let epoch_b = parts_b.get(3).unwrap_or(&"0").parse::<i64>().unwrap_or(0);
+ let nonce_b = parts_b.get(2).unwrap_or(&"0").parse::<i64>().unwrap_or(0);
+ let timestamp_b = parts_b.get(4).unwrap_or(&"0").parse::<i64>().unwrap_or(0);
+
+ // Compare as integers
+ epoch_a
+ .cmp(&epoch_b)
+ .then(nonce_a.cmp(&nonce_b))
+ .then(timestamp_a.cmp(&timestamp_b))
+ });
+
+ // Apply pagination by slicing the sorted keys
+ let start_index = 0;
+ let end_index = actual_limit.min(keys_with_ids.len());
+ let paginated_keys: Vec<(String, String)> = keys_with_ids[start_index..end_index].to_vec();
+
+ // Fetch the messages for each paginated key
+ for (_, assignment_id) in paginated_keys {
+ let message = self.get_message(&assignment_id)?;
+ messages.push(message);
+ }
+
+ // Determine if there's a next page
+ let has_next_page = keys_with_ids.len() > actual_limit;
+
+ // Return the paginated result with the calculated `has_next_page`
+ Ok(PaginatedMessages::from_messages(messages, has_next_page)?)
+ }
+
+
+
+ fn get_latest_message(&self, process_id: &str) -> Result<Option<Message>, StoreErrorType> {
+ // Create the prefix key to search for messages by process_id
+ let prefix_key = format!("process_messages:{}:", process_id);
+
+ // Use a prefix iterator to scan through messages in sequential order
+ let mut iter = self.index_db.prefix_iterator(prefix_key.as_bytes());
+
+ let mut latest_message: Option<Message> = None;
+
+ // Iterate through all the messages and get the last one
+ for result in iter {
+ let (_key, assignment_id_bytes) = result?;
+
+ // Fetch the message by assignment_id
+ let assignment_id = String::from_utf8(assignment_id_bytes.to_vec())?;
+ latest_message = Some(self.get_message(&assignment_id)?);
+ }
+
+ // Return the latest message if found
+ Ok(latest_message)
+ }
+}