riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: resolve critical concurrency, data loss, schema, and search audit passes

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-07 07:09:28 +01:00
1 parent e4a0fe72df
commit d80915635f
7 files changed
+846 -450
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/store.rs
+102 -54
View File
@@ -8,6 +8,10 @@ pub const STORE_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("stor
enum DbOp {
Insert(Vec<u8>),
Delete,
Batch {
inserts: Vec<(String, Vec<u8>)>,
deletes: Vec<String>,
},
}
/// Internal write request dispatched to the single database writer actor.
@@ -83,6 +87,28 @@ impl DbWriteQueue {
);
}
}
DbOp::Batch { inserts, deletes } => {
for del_k in deletes {
if let Err(e) = table.remove(del_k.as_str()) {
tracing::error!(
"Failed to delete batch key '{}' from redb: {}",
del_k,
e
);
}
}
for (ins_k, ins_bytes) in inserts {
if let Err(e) =
table.insert(ins_k.as_str(), ins_bytes.as_slice())
{
tracing::error!(
"Failed to insert batch key '{}' into redb: {}",
ins_k,
e
);
}
}
}
}
}
}
@@ -130,6 +156,16 @@ impl DbWriteQueue {
self.push_op(key, DbOp::Delete, flushed_notifier)
}
pub fn push_batch(
&self,
key: String,
inserts: Vec<(String, Vec<u8>)>,
deletes: Vec<String>,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
self.push_op(key, DbOp::Batch { inserts, deletes }, flushed_notifier)
}
fn push_op(
&self,
key: String,
@@ -191,6 +227,17 @@ impl DbWriteQueue {
.await
}
pub async fn push_batch_async(
&self,
key: String,
inserts: Vec<(String, Vec<u8>)>,
deletes: Vec<String>,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
self.push_op_async(key, DbOp::Batch { inserts, deletes }, flushed_notifier)
.await
}
async fn push_op_async(
&self,
key: String,
@@ -222,14 +269,16 @@ pub struct Store<T> {
key: String,
queue: DbWriteQueue,
is_corrupted: bool,
known_granular_keys: Arc<RwLock<std::collections::HashSet<String>>>,
}
impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T> {
pub fn new(key: &str, db: Arc<Database>) -> Self {
let (initial_data, is_corrupted) = Self::load_from_db(key, &db);
let (initial_data, is_corrupted, known_keys) = Self::load_from_db(key, &db);
let cache = Arc::new(RwLock::new(initial_data));
let flushed = Arc::new(tokio::sync::Notify::new());
let queue = get_or_create_queue(db);
let known_granular_keys = Arc::new(RwLock::new(known_keys));
Self {
cache,
@@ -237,29 +286,38 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
key: key.to_string(),
queue,
is_corrupted,
known_granular_keys,
}
}
fn load_from_db(key: &str, db: &Database) -> (T, bool) {
fn load_from_db(key: &str, db: &Database) -> (T, bool, std::collections::HashSet<String>) {
let Ok(read_txn) = db.begin_read() else {
tracing::error!("Failed to begin read transaction for key '{}'", key);
return (T::default(), false);
return (T::default(), false, std::collections::HashSet::new());
};
let Ok(table) = read_txn.open_table(STORE_TABLE) else {
return (T::default(), false);
return (T::default(), false, std::collections::HashSet::new());
};
// 1. Check monolithic key first as the authoritative snapshot
match table.get(key) {
Ok(Some(value)) => match serde_json::from_slice::<T>(value.value()) {
Ok(parsed) => return (parsed, false),
Ok(parsed) => {
let mut known = std::collections::HashSet::new();
if let Ok(val) = serde_json::to_value(&parsed) {
for k in Self::extract_granular_keys(key, &val) {
known.insert(k);
}
}
return (parsed, false, known);
}
Err(e) => {
tracing::error!(
"CRITICAL: Corrupted data for key '{}' in database: {}. Quarantine mode active: state initialized to empty default without overwriting DB key.",
key,
e
);
return (T::default(), true);
return (T::default(), true, std::collections::HashSet::new());
}
},
Ok(None) => {}
@@ -273,6 +331,7 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
let mut items_array = Vec::new();
let mut items_map = serde_json::Map::new();
let mut found_granular = false;
let mut known = std::collections::HashSet::new();
if let Ok(range) = table.range(prefix.as_str()..) {
for entry in range {
@@ -282,6 +341,7 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
break;
}
found_granular = true;
known.insert(k_str.to_string());
if let Ok(val) = serde_json::from_slice::<serde_json::Value>(v.value()) {
let sub_key = &k_str[prefix.len()..];
items_array.push(val.clone());
@@ -293,14 +353,14 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
if found_granular {
if let Ok(parsed) = serde_json::from_value::<T>(serde_json::Value::Array(items_array)) {
return (parsed, false);
return (parsed, false, known);
}
if let Ok(parsed) = serde_json::from_value::<T>(serde_json::Value::Object(items_map)) {
return (parsed, false);
return (parsed, false, known);
}
}
(T::default(), false)
(T::default(), false, std::collections::HashSet::new())
}
fn extract_granular_keys(base_key: &str, val: &serde_json::Value) -> Vec<String> {
@@ -377,49 +437,43 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
return;
}
// Fast mutation under critical lock section, then immediately release the RwLock guard
let (old_snapshot, new_snapshot) = {
// Fast mutation under critical lock section, only ONE clone taken, then immediately release the RwLock guard
let new_snapshot = {
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
let old = (*lock).clone();
f(&mut lock);
let new = (*lock).clone();
(old, new)
(*lock).clone()
};
// Expensive serialization and granular extraction run completely unblocked outside the lock
let old_keys = serde_json::to_value(&old_snapshot)
.map(|val| Self::extract_granular_keys(&self.key, &val))
.unwrap_or_default();
let full_bytes_res = serde_json::to_vec(&new_snapshot);
let granular_entries = serde_json::to_value(&new_snapshot)
.map(|val| Self::extract_granular_entries(&self.key, &val))
.unwrap_or_default();
let new_keys: std::collections::HashSet<&str> =
granular_entries.iter().map(|(k, _)| k.as_str()).collect();
let new_keys: std::collections::HashSet<String> =
granular_entries.iter().map(|(k, _)| k.clone()).collect();
let mut removed_keys = Vec::new();
for old_k in &old_keys {
if !new_keys.contains(old_k.as_str()) {
removed_keys.push(old_k.clone());
{
let mut known = self.known_granular_keys.write().unwrap_or_else(|e| e.into_inner());
for old_k in known.iter() {
if !new_keys.contains(old_k) {
removed_keys.push(old_k.clone());
}
}
*known = new_keys;
}
match full_bytes_res {
Ok(data) => {
// Delete removed granular entries so they don't resurrect on restart
for del_key in removed_keys {
self.queue.push_delete(del_key, self.flushed.clone());
}
// Queue granular entries
let mut batch_inserts = Vec::with_capacity(granular_entries.len() + 1);
for (g_key, g_bytes) in granular_entries {
self.queue.push(g_key, g_bytes, self.flushed.clone());
batch_inserts.push((g_key, g_bytes));
}
batch_inserts.push((self.key.clone(), data.clone()));
if self
.queue
.push(self.key.clone(), data.clone(), self.flushed.clone())
.push_batch(self.key.clone(), batch_inserts.clone(), removed_keys.clone(), self.flushed.clone())
.is_none()
{
tracing::warn!(
@@ -433,7 +487,7 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
handle.spawn(async move {
let _ = tokio::time::timeout(
std::time::Duration::from_secs(10),
queue.push_async(key, data, flushed),
queue.push_batch_async(key, batch_inserts, removed_keys, flushed),
)
.await;
});
@@ -460,47 +514,41 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
return;
}
let (old_snapshot, new_snapshot) = {
let new_snapshot = {
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
let old = (*lock).clone();
f(&mut lock);
let new = (*lock).clone();
(old, new)
(*lock).clone()
};
let old_keys = serde_json::to_value(&old_snapshot)
.map(|val| Self::extract_granular_keys(&self.key, &val))
.unwrap_or_default();
let full_bytes_res = serde_json::to_vec(&new_snapshot);
let granular_entries = serde_json::to_value(&new_snapshot)
.map(|val| Self::extract_granular_entries(&self.key, &val))
.unwrap_or_default();
let new_keys: std::collections::HashSet<&str> =
granular_entries.iter().map(|(k, _)| k.as_str()).collect();
let new_keys: std::collections::HashSet<String> =
granular_entries.iter().map(|(k, _)| k.clone()).collect();
let mut removed_keys = Vec::new();
for old_k in &old_keys {
if !new_keys.contains(old_k.as_str()) {
removed_keys.push(old_k.clone());
{
let mut known = self.known_granular_keys.write().unwrap_or_else(|e| e.into_inner());
for old_k in known.iter() {
if !new_keys.contains(old_k) {
removed_keys.push(old_k.clone());
}
}
*known = new_keys;
}
match full_bytes_res {
Ok(data) => {
for del_key in removed_keys {
self.queue
.push_delete_async(del_key, self.flushed.clone())
.await;
}
let mut batch_inserts = Vec::with_capacity(granular_entries.len() + 1);
for (g_key, g_bytes) in granular_entries {
self.queue
.push_async(g_key, g_bytes, self.flushed.clone())
.await;
batch_inserts.push((g_key, g_bytes));
}
batch_inserts.push((self.key.clone(), data));
if let Some(rx) = self
.queue
.push_async(self.key.clone(), data, self.flushed.clone())
.push_batch_async(self.key.clone(), batch_inserts, removed_keys, self.flushed.clone())
.await
{
let _ = rx.await;