Keep removed objects cached on persist failure

Check for the object first, remove it from the backing store, and only then delete it from the in-memory map.

This finding was discovered by Project Loupe

AI-Assisted-By: OpenAI Codex

Author: benthecarman

src/data_store.rs

@@ -113,8 +113,8 @@ where
 
 	pub(crate) async fn remove(&self, id: &SO::Id) -> Result<(), Error> {
 		let _guard = self.mutation_lock.lock().await;
-		let removed = { self.objects.lock().expect("lock").remove(id).is_some() };
-		if removed {
+		let should_remove = { self.objects.lock().expect("lock").contains_key(id) };
+		if should_remove {
 			let store_key = id.encode_to_hex_str();
 			KVStore::remove(
 				&*self.kv_store,
@@ -135,6 +135,7 @@ where
 				);
 				Error::PersistenceFailed
 			})?;
+			self.objects.lock().expect("lock").remove(id);
 		}
 		Ok(())
 	}
@@ -421,4 +422,14 @@ mod tests {
 		assert_eq!(Err(Error::PersistenceFailed), data_store.update(update).await);
 		assert_eq!(Some(object), data_store.get(&id));
 	}
+
+	#[tokio::test]
+	async fn remove_does_not_mutate_memory_if_persist_fails() {
+		let id = TestObjectId { id: [42u8; 4] };
+		let object = TestObject { id, data: [23u8; 3] };
+		let data_store = new_failing_data_store(vec![object]);
+
+		assert_eq!(Err(Error::PersistenceFailed), data_store.remove(&id).await);
+		assert_eq!(Some(object), data_store.get(&id));
+	}
 }