chore: Improve and optimize in_memory implementation

Author: Harsh Dev Pathak

.github/workflows/build-and-test-in-memory.yml

@@ -25,7 +25,7 @@ jobs:
           [server_config]
           host = "127.0.0.1"
           port = 8080
-          store_type = "in_memory"
+          store_type = "in-memory"
           EOF
 
       - name: Build server

.github/workflows/ldk-node-integration.yml

@@ -102,7 +102,7 @@ jobs:
           [server_config]
           host = "127.0.0.1"
           port = 8080
-          store_type = "in_memory"
+          store_type = "in-memory"
           EOF
 
       - name: Build & Start VSS Server

rust/impls/src/in_memory_store.rs

@@ -1,6 +1,4 @@
-use crate::postgres_store::{
-	VssDbRecord, LIST_KEY_VERSIONS_MAX_PAGE_SIZE, MAX_PUT_REQUEST_ITEM_COUNT,
-};
+use crate::{VssDbRecord, LIST_KEY_VERSIONS_MAX_PAGE_SIZE, MAX_PUT_REQUEST_ITEM_COUNT};
 use api::error::VssError;
 use api::kv_store::{KvStore, GLOBAL_VERSION_KEY, INITIAL_RECORD_VERSION};
 use api::types::{
@@ -10,41 +8,38 @@ use api::types::{
 use async_trait::async_trait;
 use bytes::Bytes;
 use chrono::prelude::Utc;
-use std::collections::HashMap;
-use std::sync::Arc;
-use tokio::sync::RwLock;
+use std::collections::BTreeMap;
+use tokio::sync::Mutex;
 
 fn build_storage_key(user_token: &str, store_id: &str, key: &str) -> String {
 	format!("{}#{}#{}", user_token, store_id, key)
 }
 
 /// In-memory implementation of the VSS Store.
 pub struct InMemoryBackendImpl {
-	store: Arc<RwLock<HashMap<String, VssDbRecord>>>,
+	store: Mutex<BTreeMap<String, VssDbRecord>>,
 }
 
 impl InMemoryBackendImpl {
 	/// Creates an in-memory instance.
 	pub fn new() -> Self {
-		Self { store: Arc::new(RwLock::new(HashMap::new())) }
+		Self { store: Mutex::new(BTreeMap::new()) }
 	}
 
 	fn get_current_global_version(
-		&self, guard: &HashMap<String, VssDbRecord>, user_token: &str, store_id: &str,
+		&self, guard: &BTreeMap<String, VssDbRecord>, user_token: &str, store_id: &str,
 	) -> i64 {
 		let global_key = build_storage_key(user_token, store_id, GLOBAL_VERSION_KEY);
 		guard.get(&global_key).map(|r| r.version).unwrap_or(0)
 	}
 }
 
-// Validation functions - check if operations can succeed without modifying data
 fn validate_put_operation(
-	store: &HashMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: &KeyValue,
+	store: &BTreeMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: &KeyValue,
 ) -> Result<(), VssError> {
 	let key = build_storage_key(user_token, store_id, &key_value.key);
 
 	if key_value.version == -1 {
-		// Non-conditional upsert always succeeds
 		Ok(())
 	} else if key_value.version == 0 {
 		if store.contains_key(&key) {
@@ -75,12 +70,11 @@ fn validate_put_operation(
 }
 
 fn validate_delete_operation(
-	store: &HashMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: &KeyValue,
+	store: &BTreeMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: &KeyValue,
 ) -> Result<(), VssError> {
 	let key = build_storage_key(user_token, store_id, &key_value.key);
 
 	if key_value.version == -1 {
-		// Non-conditional delete always succeeds
 		Ok(())
 	} else {
 		if let Some(existing) = store.get(&key) {
@@ -101,20 +95,25 @@ fn validate_delete_operation(
 	}
 }
 
-fn execute_non_conditional_upsert(
-	store: &mut HashMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: KeyValue,
+fn execute_put_object(
+	store: &mut BTreeMap<String, VssDbRecord>, user_token: &str, store_id: &str,
+	key_value: KeyValue,
 ) {
 	let key = build_storage_key(user_token, store_id, &key_value.key);
 	let now = Utc::now();
 
 	match store.entry(key) {
-		std::collections::hash_map::Entry::Occupied(mut occ) => {
+		std::collections::btree_map::Entry::Occupied(mut occ) => {
 			let existing = occ.get_mut();
-			existing.version = INITIAL_RECORD_VERSION as i64;
+			existing.version = if key_value.version == -1 {
+				INITIAL_RECORD_VERSION as i64
+			} else {
+				existing.version.saturating_add(1)
+			};
 			existing.value = key_value.value.to_vec();
 			existing.last_updated_at = now;
 		},
-		std::collections::hash_map::Entry::Vacant(vac) => {
+		std::collections::btree_map::Entry::Vacant(vac) => {
 			let new_record = VssDbRecord {
 				user_token: user_token.to_string(),
 				store_id: store_id.to_string(),
@@ -129,51 +128,8 @@ fn execute_non_conditional_upsert(
 	}
 }
 
-fn execute_conditional_insert(
-	store: &mut HashMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: KeyValue,
-) {
-	let key = build_storage_key(user_token, store_id, &key_value.key);
-	let now = Utc::now();
-
-	let new_record = VssDbRecord {
-		user_token: user_token.to_string(),
-		store_id: store_id.to_string(),
-		key: key_value.key,
-		value: key_value.value.to_vec(),
-		version: INITIAL_RECORD_VERSION as i64,
-		created_at: now,
-		last_updated_at: now,
-	};
-	store.insert(key, new_record);
-}
-
-fn execute_conditional_update(
-	store: &mut HashMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: KeyValue,
-) {
-	let key = build_storage_key(user_token, store_id, &key_value.key);
-	let now = Utc::now();
-
-	if let Some(existing) = store.get_mut(&key) {
-		existing.version = key_value.version.saturating_add(1);
-		existing.value = key_value.value.to_vec();
-		existing.last_updated_at = now;
-	}
-}
-
-fn execute_put_object(
-	store: &mut HashMap<String, VssDbRecord>, user_token: &str, store_id: &str, key_value: KeyValue,
-) {
-	if key_value.version == -1 {
-		execute_non_conditional_upsert(store, user_token, store_id, key_value);
-	} else if key_value.version == 0 {
-		execute_conditional_insert(store, user_token, store_id, key_value);
-	} else {
-		execute_conditional_update(store, user_token, store_id, key_value);
-	}
-}
-
 fn execute_delete_object(
-	store: &mut HashMap<String, VssDbRecord>, user_token: &str, store_id: &str,
+	store: &mut BTreeMap<String, VssDbRecord>, user_token: &str, store_id: &str,
 	key_value: &KeyValue,
 ) {
 	let key = build_storage_key(user_token, store_id, &key_value.key);
@@ -185,8 +141,21 @@ impl KvStore for InMemoryBackendImpl {
 	async fn get(
 		&self, user_token: String, request: GetObjectRequest,
 	) -> Result<GetObjectResponse, VssError> {
+		if request.key == GLOBAL_VERSION_KEY {
+			let guard = self.store.lock().await;
+			let global_version =
+				self.get_current_global_version(&guard, &user_token, &request.store_id);
+			return Ok(GetObjectResponse {
+				value: Some(KeyValue {
+					key: GLOBAL_VERSION_KEY.to_string(),
+					value: Bytes::new(),
+					version: global_version,
+				}),
+			});
+		}
+
 		let key = build_storage_key(&user_token, &request.store_id, &request.key);
-		let guard = self.store.read().await;
+		let guard = self.store.lock().await;
 
 		if let Some(record) = guard.get(&key) {
 			Ok(GetObjectResponse {
@@ -196,14 +165,6 @@ impl KvStore for InMemoryBackendImpl {
 					version: record.version,
 				}),
 			})
-		} else if request.key == GLOBAL_VERSION_KEY {
-			Ok(GetObjectResponse {
-				value: Some(KeyValue {
-					key: GLOBAL_VERSION_KEY.to_string(),
-					value: Bytes::new(),
-					version: 0,
-				}),
-			})
 		} else {
 			Err(VssError::NoSuchKeyError("Requested key not found.".to_string()))
 		}
@@ -220,8 +181,8 @@ impl KvStore for InMemoryBackendImpl {
 			)));
 		}
 
-		let store_id = request.store_id.clone();
-		let mut guard = self.store.write().await;
+		let store_id = request.store_id;
+		let mut guard = self.store.lock().await;
 
 		if let Some(version) = request.global_version {
 			validate_put_operation(
@@ -267,8 +228,8 @@ impl KvStore for InMemoryBackendImpl {
 			VssError::InvalidRequestError("key_value missing in DeleteObjectRequest".to_string())
 		})?;
 
-		let store_id = request.store_id.clone();
-		let mut guard = self.store.write().await;
+		let store_id = request.store_id;
+		let mut guard = self.store.lock().await;
 
 		execute_delete_object(&mut guard, &user_token, &store_id, &key_value);
 
@@ -279,70 +240,42 @@ impl KvStore for InMemoryBackendImpl {
 		&self, user_token: String, request: ListKeyVersionsRequest,
 	) -> Result<ListKeyVersionsResponse, VssError> {
 		let store_id = request.store_id;
-		let key_prefix = request.key_prefix.unwrap_or("".to_string());
-		let page_token_option = request.page_token;
-		let page_size = request.page_size.unwrap_or(i32::MAX);
+		let key_prefix = request.key_prefix.clone().unwrap_or("".to_string());
+		let page_size = request.page_size.unwrap_or(LIST_KEY_VERSIONS_MAX_PAGE_SIZE);
 		let limit = std::cmp::min(page_size, LIST_KEY_VERSIONS_MAX_PAGE_SIZE) as usize;
 
-		let (keys_with_versions, global_version) = {
-			let guard = self.store.read().await;
+		let offset: usize =
+			request.page_token.as_ref().and_then(|s| s.parse::<usize>().ok()).unwrap_or(0);
 
-			let mut global_version: Option<i64> = None;
-			if page_token_option.is_none() {
-				global_version =
-					Some(self.get_current_global_version(&guard, &user_token, &store_id));
-			}
+		let guard = self.store.lock().await;
 
-			let storage_prefix = format!("{}#{}#", user_token, store_id);
-			let mut temp: Vec<(String, i64)> = Vec::new();
-
-			for (storage_key, r) in guard.iter() {
-				if !storage_key.starts_with(&storage_prefix) {
-					continue;
-				}
-				let key = &storage_key[storage_prefix.len()..];
-				if key == GLOBAL_VERSION_KEY {
-					continue;
-				}
-				if !key_prefix.is_empty() && !key.starts_with(&key_prefix) {
-					continue;
-				}
-				temp.push((key.to_string(), r.version));
-			}
-
-			(temp, global_version)
-		};
-
-		let mut keys_with_versions = keys_with_versions;
-		keys_with_versions.sort_by(|a, b| a.0.cmp(&b.0));
-
-		let start_idx = if page_token_option.is_none() {
-			0
-		} else if page_token_option.as_deref() == Some("") {
-			keys_with_versions.len()
+		let global_version = if offset == 0 {
+			Some(self.get_current_global_version(&guard, &user_token, &store_id))
 		} else {
-			let token = page_token_option.as_deref().unwrap();
-			keys_with_versions
-				.iter()
-				.position(|(k, _)| k.as_str() > token)
-				.unwrap_or(keys_with_versions.len())
+			None
 		};
 
-		let page_items: Vec<KeyValue> = keys_with_versions
+		let storage_prefix = build_storage_key(&user_token, &store_id, &key_prefix);
+		let global_key = build_storage_key(&user_token, &store_id, GLOBAL_VERSION_KEY);
+		let prefix_len = format!("{}#{}#", user_token, store_id).len();
+		let page_items: Vec<KeyValue> = guard
 			.iter()
-			.skip(start_idx)
+			.filter(|(storage_key, _)| {
+				storage_key.starts_with(&storage_prefix) && *storage_key != &global_key
+			})
+			.skip(offset)
 			.take(limit)
-			.map(|(key, version)| KeyValue {
-				key: key.clone(),
-				value: Bytes::new(),
-				version: *version,
+			.map(|(storage_key, record)| {
+				let key = &storage_key[prefix_len..];
+				KeyValue { key: key.to_string(), value: Bytes::new(), version: record.version }
 			})
 			.collect();
 
+		let next_offset = offset + page_items.len();
 		let next_page_token = if page_items.is_empty() {
 			Some("".to_string())
 		} else {
-			page_items.last().map(|kv| kv.key.clone())
+			Some(next_offset.to_string())
 		};
 
 		Ok(ListKeyVersionsResponse { key_versions: page_items, next_page_token, global_version })

rust/impls/src/lib.rs

@@ -11,11 +11,46 @@
 #![deny(rustdoc::private_intra_doc_links)]
 #![deny(missing_docs)]
 
+use chrono::Utc;
+
 /// 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;
 
+/// A record stored in the VSS database.
+struct VssDbRecord {
+	/// Token uniquely identifying the user that owns this record.
+	user_token: String,
+	/// Identifier for the store this record belongs to.
+	store_id: String,
+	/// Key under which the value is stored.
+	key: String,
+	/// Stored value as raw bytes.
+	value: Vec<u8>,
+	/// Version number for optimistic concurrency control.
+	version: i64,
+	/// Timestamp when the record was created (UTC).
+	created_at: chrono::DateTime<Utc>,
+	/// Timestamp when the record was last updated (UTC).
+	last_updated_at: chrono::DateTime<Utc>,
+}
+
+/// The maximum number of key versions that can be returned in a single page.
+///
+/// This constant helps control memory and bandwidth usage for list operations,
+/// preventing overly large payloads. If the number of results exceeds this limit,
+/// the response will be paginated.
+const LIST_KEY_VERSIONS_MAX_PAGE_SIZE: i32 = 100;
+
+/// The maximum number of items allowed in a single `PutObjectRequest`.
+///
+/// Setting an upper bound on the number of items helps ensure that
+/// each request stays within acceptable memory and performance limits.
+/// Exceeding this value will result in request rejection through [`VssError::InvalidRequestError`].
+const MAX_PUT_REQUEST_ITEM_COUNT: usize = 1000;
+
+mod migrations;
+
 #[macro_use]
 extern crate api;