Optimized InMemoryBackendImpl

Author: Harshdev098

rust/Cargo.lock

@@ -25,7 +25,7 @@ cargo build --release
    cargo run -- server/vss-server-config.toml
    ```
 
-   **Note:** For testing puropose you can edit vss-server-config.toml to use `store_type` as in-memory instead of postgresql `store_type = "memory"`
+   **Note:** For testing purposes you can edit `vss-server-config.toml` to use `store_type` as in-memory instead of PostgreSQL: `store_type = "in_memory"`
 4. VSS endpoint should be reachable at `http://localhost:8080/vss`.
 
 ### Configuration

rust/README.md

@@ -0,0 +1,367 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use chrono::prelude::Utc;
+
+use crate::postgres_store::{
+	VssDbRecord, LIST_KEY_VERSIONS_MAX_PAGE_SIZE, MAX_PUT_REQUEST_ITEM_COUNT,
+};
+use api::error::VssError;
+use api::kv_store::{KvStore, GLOBAL_VERSION_KEY};
+use api::types::{
+	DeleteObjectRequest, DeleteObjectResponse, GetObjectRequest, GetObjectResponse, KeyValue,
+	ListKeyVersionsRequest, ListKeyVersionsResponse, PutObjectRequest, PutObjectResponse,
+};
+
+/// In-memory backend for VSS, for testing purposes only.
+pub struct InMemoryBackendImpl {
+	store: Arc<Mutex<HashMap<String, VssDbRecord>>>,
+}
+
+impl InMemoryBackendImpl {
+	/// Creates a new in-memory backend.
+	pub fn new() -> Self {
+		Self { store: Arc::new(Mutex::new(HashMap::new())) }
+	}
+
+	fn build_vss_record(&self, user_token: String, store_id: String, kv: KeyValue) -> VssDbRecord {
+		let now = Utc::now();
+		VssDbRecord {
+			user_token,
+			store_id,
+			key: kv.key,
+			value: kv.value.to_vec(),
+			version: kv.version,
+			created_at: now,
+			last_updated_at: now,
+		}
+	}
+
+	fn build_key(user_token: &str, store_id: &str, key: &str) -> String {
+		format!("{}#{}#{}", user_token, store_id, key)
+	}
+
+	fn get_current_global_version(
+		&self, guard: &HashMap<String, VssDbRecord>, user_token: &str, store_id: &str,
+	) -> i64 {
+		let global_key = Self::build_key(user_token, store_id, GLOBAL_VERSION_KEY);
+		guard.get(&global_key).map(|r| r.version).unwrap_or(0)
+	}
+
+	fn set_global_version(
+		&self, guard: &mut HashMap<String, VssDbRecord>, user_token: String, store_id: String,
+		new_version: i64,
+	) {
+		let global_key = Self::build_key(&user_token, &store_id, GLOBAL_VERSION_KEY);
+		let now = Utc::now();
+
+		let entry = guard.entry(global_key);
+		match entry {
+			std::collections::hash_map::Entry::Occupied(mut occ) => {
+				let rec = occ.get_mut();
+				rec.version = new_version;
+				rec.last_updated_at = now;
+			},
+			std::collections::hash_map::Entry::Vacant(vac) => {
+				let record = VssDbRecord {
+					user_token,
+					store_id,
+					key: GLOBAL_VERSION_KEY.to_string(),
+					value: vec![],
+					version: new_version,
+					created_at: now,
+					last_updated_at: now,
+				};
+				vac.insert(record);
+			},
+		}
+	}
+}
+
+fn execute_put_object(
+	store: &mut HashMap<String, VssDbRecord>, record: VssDbRecord,
+) -> Result<(), VssError> {
+	let key = format!("{}#{}#{}", record.user_token, record.store_id, record.key);
+	let now = Utc::now();
+
+	let entry = store.entry(key);
+	match entry {
+		std::collections::hash_map::Entry::Occupied(mut occ) => {
+			let existing = occ.get_mut();
+			if existing.version >= record.version {
+				return Err(VssError::ConflictError(format!(
+					"Version conflict on put for key {}",
+					record.key
+				)));
+			}
+			existing.version = record.version;
+			existing.value = record.value;
+			existing.last_updated_at = now;
+		},
+		std::collections::hash_map::Entry::Vacant(vac) => {
+			vac.insert(record);
+		},
+	}
+	Ok(())
+}
+
+fn execute_delete_object(
+	store: &mut HashMap<String, VssDbRecord>, record: VssDbRecord,
+) -> Result<(), VssError> {
+	let key = format!("{}#{}#{}", record.user_token, record.store_id, record.key);
+	if let Some(existing) = store.get(&key) {
+		if existing.version != record.version {
+			return Ok(());
+		}
+		store.remove(&key);
+	}
+	Ok(())
+}
+
+#[async_trait]
+impl KvStore for InMemoryBackendImpl {
+	async fn get(
+		&self, user_token: String, request: GetObjectRequest,
+	) -> Result<GetObjectResponse, VssError> {
+		let key = Self::build_key(&user_token, &request.store_id, &request.key);
+		let guard = self.store.lock().unwrap();
+
+		if let Some(record) = guard.get(&key) {
+			Ok(GetObjectResponse {
+				value: Some(KeyValue {
+					key: record.key.clone(),
+					value: Bytes::from(record.value.clone()),
+					version: record.version,
+				}),
+			})
+		} else if request.key == GLOBAL_VERSION_KEY {
+			let current_global =
+				self.get_current_global_version(&guard, &user_token, &request.store_id);
+			Ok(GetObjectResponse {
+				value: Some(KeyValue {
+					key: GLOBAL_VERSION_KEY.to_string(),
+					value: Bytes::new(),
+					version: current_global,
+				}),
+			})
+		} else {
+			Err(VssError::NoSuchKeyError("Requested key not found.".to_string()))
+		}
+	}
+
+	async fn put(
+		&self, user_token: String, request: PutObjectRequest,
+	) -> Result<PutObjectResponse, VssError> {
+		if request.transaction_items.len() + request.delete_items.len() > MAX_PUT_REQUEST_ITEM_COUNT
+		{
+			return Err(VssError::InvalidRequestError(format!(
+				"Number of write items per request should be less than equal to {}",
+				MAX_PUT_REQUEST_ITEM_COUNT
+			)));
+		}
+
+		let store_id = request.store_id.clone();
+		let mut guard = self.store.lock().unwrap();
+
+		// Handling global_version precondition
+		let current_global = self.get_current_global_version(&guard, &user_token, &store_id);
+		if let Some(expected_global) = request.global_version {
+			if current_global != expected_global {
+				return Err(VssError::ConflictError(format!(
+					"Global version conflict: expected {}, current {}",
+					expected_global, current_global
+				)));
+			}
+		}
+
+		// Check for conflicts on puts
+		for kv in &request.transaction_items {
+			let key = Self::build_key(&user_token, &store_id, &kv.key);
+			if let Some(existing) = guard.get(&key) {
+				if existing.version >= kv.version {
+					return Err(VssError::ConflictError(format!(
+						"Version conflict on put for key {}",
+						kv.key
+					)));
+				}
+			}
+		}
+
+		// Apply updates
+		let vss_put_records: Vec<VssDbRecord> = request
+			.transaction_items
+			.into_iter()
+			.map(|kv| self.build_vss_record(user_token.clone(), store_id.clone(), kv))
+			.collect();
+
+		let vss_delete_records: Vec<VssDbRecord> = request
+			.delete_items
+			.into_iter()
+			.map(|kv| self.build_vss_record(user_token.clone(), store_id.clone(), kv))
+			.collect();
+
+		let mut mutated = false;
+		for vss_record in vss_put_records {
+			execute_put_object(&mut guard, vss_record)?;
+			mutated = true;
+		}
+		for vss_record in vss_delete_records {
+			// execute_delete_object is no-op on version mismatch, but call anyway
+			execute_delete_object(&mut guard, vss_record)?;
+			mutated = true;
+		}
+
+		if mutated || request.global_version.is_some() {
+			let new_global = current_global + 1;
+			self.set_global_version(&mut guard, user_token.clone(), store_id.clone(), new_global);
+		}
+
+		Ok(PutObjectResponse {})
+	}
+
+	async fn delete(
+		&self, user_token: String, request: DeleteObjectRequest,
+	) -> Result<DeleteObjectResponse, VssError> {
+		let key_value = request.key_value.ok_or_else(|| {
+			VssError::InvalidRequestError("key_value missing in DeleteObjectRequest".to_string())
+		})?;
+		let store_id = request.store_id.clone();
+		let mut guard = self.store.lock().unwrap();
+
+		let current_global = self.get_current_global_version(&guard, &user_token, &store_id);
+
+		// Check for version conflict
+		let key = Self::build_key(&user_token, &store_id, &key_value.key);
+		let mutated = if let Some(existing) = guard.get(&key) {
+			if existing.version != key_value.version {
+				false
+			} else {
+				true
+			}
+		} else {
+			false
+		};
+
+		if mutated {
+			let vss_record = self.build_vss_record(user_token.clone(), store_id.clone(), key_value);
+			execute_delete_object(&mut guard, vss_record)?;
+			let new_global = current_global + 1;
+			self.set_global_version(&mut guard, user_token.clone(), store_id.clone(), new_global);
+		}
+
+		Ok(DeleteObjectResponse {})
+	}
+
+	async fn list_key_versions(
+		&self, user_token: String, request: ListKeyVersionsRequest,
+	) -> Result<ListKeyVersionsResponse, VssError> {
+		let store_id = request.store_id;
+		let key_prefix = request.key_prefix.unwrap_or_default();
+		let page_token = request.page_token.unwrap_or_default();
+		let page_size = request.page_size.unwrap_or(i32::MAX);
+		let limit = std::cmp::min(page_size, LIST_KEY_VERSIONS_MAX_PAGE_SIZE) as usize;
+
+		let mut global_version = None;
+		if page_token.is_empty() {
+			let get_global_version_request = GetObjectRequest {
+				store_id: store_id.clone(),
+				key: GLOBAL_VERSION_KEY.to_string(),
+			};
+			let get_response = self.get(user_token.clone(), get_global_version_request).await?;
+			global_version = Some(get_response.value.unwrap().version);
+		}
+
+		let key_versions: Vec<KeyValue> = {
+			let guard = self.store.lock().unwrap();
+			let mut key_versions: Vec<KeyValue> = guard
+				.iter()
+				.filter(|(k, _)| {
+					let parts: Vec<&str> = k.split('#').collect();
+					if parts.len() < 3 {
+						return false;
+					}
+					parts[0] == user_token.as_str()
+						&& parts[1] == store_id.as_str()
+						&& parts[2].starts_with(&key_prefix)
+						&& parts[2] > page_token.as_str()
+						&& parts[2] != GLOBAL_VERSION_KEY
+				})
+				.map(|(_, record)| KeyValue {
+					key: record.key.clone(),
+					value: Bytes::new(),
+					version: record.version,
+				})
+				.collect();
+
+			key_versions.sort_by(|a, b| a.key.cmp(&b.key));
+			key_versions.into_iter().take(limit).collect()
+		};
+
+		let next_page_token = if key_versions.len() == limit {
+			key_versions.last().map(|kv| kv.key.clone())
+		} else {
+			None
+		};
+
+		Ok(ListKeyVersionsResponse { key_versions, next_page_token, global_version })
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use api::kv_store::INITIAL_RECORD_VERSION;
+	use bytes::Bytes;
+	use tokio::test;
+
+	#[test]
+	async fn test_in_memory_crud() {
+		let store = InMemoryBackendImpl::new();
+		let user_token = "test_user".to_string();
+		let store_id = "test_store".to_string();
+
+		// Put
+		let put_request = PutObjectRequest {
+			store_id: store_id.clone(),
+			transaction_items: vec![KeyValue {
+				key: "key1".to_string(),
+				value: Bytes::from("value1"),
+				version: INITIAL_RECORD_VERSION as i64,
+			}],
+			delete_items: vec![],
+			global_version: None,
+		};
+		store.put(user_token.clone(), put_request).await.unwrap();
+
+		// Get
+		let get_request = GetObjectRequest { store_id: store_id.clone(), key: "key1".to_string() };
+		let response = store.get(user_token.clone(), get_request).await.unwrap();
+		let key_value = response.value.unwrap();
+		assert_eq!(key_value.value, Bytes::from("value1"));
+		let current_version = key_value.version;
+
+		// List
+		let list_request = ListKeyVersionsRequest {
+			store_id: store_id.clone(),
+			key_prefix: None,
+			page_size: Some(1),
+			page_token: None,
+		};
+		let response = store.list_key_versions(user_token.clone(), list_request).await.unwrap();
+		assert_eq!(response.key_versions.len(), 1);
+		assert_eq!(response.key_versions[0].key, "key1");
+
+		// Delete
+		let delete_request = DeleteObjectRequest {
+			store_id: store_id.clone(),
+			key_value: Some(KeyValue {
+				key: "key1".to_string(),
+				value: Bytes::new(),
+				version: current_version,
+			}),
+		};
+		store.delete(user_token.clone(), delete_request).await.unwrap();
+	}
+}

rust/impls/src/in_memory_store.rs

@@ -11,11 +11,11 @@
 #![deny(rustdoc::private_intra_doc_links)]
 #![deny(missing_docs)]
 
+/// Contains in-memory backend implementation for VSS, for testing purposes only.
+pub mod in_memory_store;
 mod migrations;
 /// Contains [PostgreSQL](https://www.postgresql.org/) based backend implementation for VSS.
 pub mod postgres_store;
-/// Contains in-memory backend implementation for VSS, for testing purposes only.
-pub mod memory_store;
 
 #[macro_use]
 extern crate api;