use bson::Document;
use rigatoni_core::state::{StateStore, StateStoreError};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, trace, warn};
#[derive(Debug, Clone)]
pub struct MemoryStore {
tokens: Arc<RwLock<HashMap<String, Document>>>,
}
impl MemoryStore {
#[must_use]
pub fn new() -> Self {
debug!("Creating new in-memory state store");
Self {
tokens: Arc::new(RwLock::new(HashMap::new())),
}
}
#[must_use]
pub fn with_tokens(tokens: HashMap<String, Document>) -> Self {
debug!(
token_count = tokens.len(),
"Creating in-memory state store with initial tokens"
);
Self {
tokens: Arc::new(RwLock::new(tokens)),
}
}
pub async fn len(&self) -> usize {
self.tokens.read().await.len()
}
pub async fn is_empty(&self) -> bool {
self.tokens.read().await.is_empty()
}
pub async fn clear(&self) {
let mut tokens = self.tokens.write().await;
let count = tokens.len();
tokens.clear();
debug!(
cleared_count = count,
"Cleared all tokens from memory store"
);
}
}
impl Default for MemoryStore {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl StateStore for MemoryStore {
async fn save_resume_token(
&self,
collection: &str,
token: &Document,
) -> Result<(), StateStoreError> {
trace!(
collection = collection,
token = ?token,
"Saving resume token to memory"
);
let mut tokens = self.tokens.write().await;
tokens.insert(collection.to_string(), token.clone());
debug!(
collection = collection,
total_tokens = tokens.len(),
"Saved resume token to memory"
);
Ok(())
}
async fn get_resume_token(
&self,
collection: &str,
) -> Result<Option<Document>, StateStoreError> {
trace!(
collection = collection,
"Retrieving resume token from memory"
);
let tokens = self.tokens.read().await;
let token = tokens.get(collection).cloned();
if token.is_some() {
debug!(collection = collection, "Found resume token in memory");
} else {
debug!(collection = collection, "No resume token found in memory");
}
Ok(token)
}
async fn delete_resume_token(&self, collection: &str) -> Result<(), StateStoreError> {
trace!(collection = collection, "Deleting resume token from memory");
let mut tokens = self.tokens.write().await;
let removed = tokens.remove(collection);
if removed.is_some() {
debug!(
collection = collection,
remaining_tokens = tokens.len(),
"Deleted resume token from memory"
);
} else {
warn!(
collection = collection,
"Attempted to delete non-existent resume token"
);
}
Ok(())
}
async fn list_resume_tokens(&self) -> Result<HashMap<String, Document>, StateStoreError> {
trace!("Listing all resume tokens from memory");
let tokens = self.tokens.read().await;
let token_count = tokens.len();
debug!(
token_count = token_count,
"Listed all resume tokens from memory"
);
Ok(tokens.clone())
}
async fn close(&self) -> Result<(), StateStoreError> {
debug!("Closing in-memory state store (no-op)");
Ok(())
}
async fn try_acquire_lock(
&self,
_key: &str,
_owner_id: &str,
_ttl: Duration,
) -> Result<bool, StateStoreError> {
trace!("MemoryStore: try_acquire_lock (no-op, always succeeds)");
Ok(true)
}
async fn refresh_lock(
&self,
_key: &str,
_owner_id: &str,
_ttl: Duration,
) -> Result<bool, StateStoreError> {
trace!("MemoryStore: refresh_lock (no-op, always succeeds)");
Ok(true)
}
async fn release_lock(&self, _key: &str, _owner_id: &str) -> Result<bool, StateStoreError> {
trace!("MemoryStore: release_lock (no-op, always succeeds)");
Ok(true)
}
async fn is_locked(&self, _key: &str) -> Result<bool, StateStoreError> {
trace!("MemoryStore: is_locked (no-op, always returns false)");
Ok(false)
}
}
#[cfg(test)]
mod tests {
use super::*;
use bson::doc;
#[tokio::test]
async fn test_new_store_is_empty() {
let store = MemoryStore::new();
assert!(store.is_empty().await);
assert_eq!(store.len().await, 0);
}
#[tokio::test]
async fn test_save_and_retrieve_token() {
let store = MemoryStore::new();
let token = doc! { "_data": "test_token_123" };
store
.save_resume_token("users", &token)
.await
.expect("Failed to save token");
let retrieved = store
.get_resume_token("users")
.await
.expect("Failed to retrieve token");
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap(), token);
assert_eq!(store.len().await, 1);
}
#[tokio::test]
async fn test_retrieve_nonexistent_token() {
let store = MemoryStore::new();
let retrieved = store
.get_resume_token("nonexistent")
.await
.expect("Failed to retrieve token");
assert!(retrieved.is_none());
}
#[tokio::test]
async fn test_update_token() {
let store = MemoryStore::new();
let token1 = doc! { "_data": "token_v1" };
let token2 = doc! { "_data": "token_v2" };
store
.save_resume_token("users", &token1)
.await
.expect("Failed to save token");
store
.save_resume_token("users", &token2)
.await
.expect("Failed to update token");
let retrieved = store.get_resume_token("users").await.unwrap();
assert_eq!(retrieved.unwrap(), token2);
assert_eq!(store.len().await, 1); }
#[tokio::test]
async fn test_delete_token() {
let store = MemoryStore::new();
let token = doc! { "_data": "test_token" };
store.save_resume_token("users", &token).await.unwrap();
assert_eq!(store.len().await, 1);
store.delete_resume_token("users").await.unwrap();
assert_eq!(store.len().await, 0);
let retrieved = store.get_resume_token("users").await.unwrap();
assert!(retrieved.is_none());
}
#[tokio::test]
async fn test_delete_nonexistent_token() {
let store = MemoryStore::new();
store.delete_resume_token("nonexistent").await.unwrap();
assert_eq!(store.len().await, 0);
}
#[tokio::test]
async fn test_list_tokens() {
let store = MemoryStore::new();
let token1 = doc! { "_data": "token1" };
let token2 = doc! { "_data": "token2" };
let token3 = doc! { "_data": "token3" };
store.save_resume_token("users", &token1).await.unwrap();
store.save_resume_token("orders", &token2).await.unwrap();
store.save_resume_token("products", &token3).await.unwrap();
let all_tokens = store.list_resume_tokens().await.unwrap();
assert_eq!(all_tokens.len(), 3);
assert_eq!(all_tokens.get("users"), Some(&token1));
assert_eq!(all_tokens.get("orders"), Some(&token2));
assert_eq!(all_tokens.get("products"), Some(&token3));
}
#[tokio::test]
async fn test_list_empty_store() {
let store = MemoryStore::new();
let tokens = store.list_resume_tokens().await.unwrap();
assert!(tokens.is_empty());
}
#[tokio::test]
async fn test_clear() {
let store = MemoryStore::new();
store
.save_resume_token("users", &doc! { "_data": "token1" })
.await
.unwrap();
store
.save_resume_token("orders", &doc! { "_data": "token2" })
.await
.unwrap();
assert_eq!(store.len().await, 2);
store.clear().await;
assert!(store.is_empty().await);
assert_eq!(store.len().await, 0);
}
#[tokio::test]
async fn test_with_tokens() {
let mut initial = HashMap::new();
initial.insert("users".to_string(), doc! { "_data": "token1" });
initial.insert("orders".to_string(), doc! { "_data": "token2" });
let store = MemoryStore::with_tokens(initial);
assert_eq!(store.len().await, 2);
let token = store.get_resume_token("users").await.unwrap();
assert!(token.is_some());
}
#[tokio::test]
async fn test_close() {
let store = MemoryStore::new();
store.close().await.expect("Failed to close store");
}
#[tokio::test]
async fn test_concurrent_access() {
use std::sync::Arc;
let store = Arc::new(MemoryStore::new());
let mut handles = vec![];
for i in 0..10 {
let store_clone = Arc::clone(&store);
let handle = tokio::spawn(async move {
let collection = format!("collection_{i}");
let token = doc! { "_data": format!("token_{i}") };
store_clone
.save_resume_token(&collection, &token)
.await
.unwrap();
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap();
}
assert_eq!(store.len().await, 10);
}
#[tokio::test]
async fn test_clone() {
let store1 = MemoryStore::new();
store1
.save_resume_token("users", &doc! { "_data": "token" })
.await
.unwrap();
let store2 = store1.clone();
assert_eq!(store2.len().await, 1);
store2
.save_resume_token("orders", &doc! { "_data": "token2" })
.await
.unwrap();
assert_eq!(store1.len().await, 2);
}
}