use async_trait::async_trait;
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use std::sync::RwLock;
use crate::error::{Error, Result};
use crate::state::{StateStore, UploadInventory, UploadState, WriteMode};
pub struct MemoryStateStore {
states: RwLock<HashMap<String, UploadState>>,
}
impl MemoryStateStore {
pub fn new() -> Self {
Self {
states: RwLock::new(HashMap::new()),
}
}
#[cfg(test)]
pub fn len(&self) -> usize {
self.states.read().unwrap().len()
}
#[cfg(test)]
pub fn is_empty(&self) -> bool {
self.states.read().unwrap().is_empty()
}
#[cfg(test)]
pub fn clear(&self) {
self.states.write().unwrap().clear();
}
}
impl Default for MemoryStateStore {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for MemoryStateStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let states = self.states.read().unwrap();
f.debug_struct("MemoryStateStore")
.field("uploads", &states.len())
.finish()
}
}
fn validate_upload_id(id: &str) -> Result<()> {
id.parse::<crate::protocol::UploadId>()?;
Ok(())
}
#[async_trait]
impl StateStore for MemoryStateStore {
fn name(&self) -> &'static str {
"memory"
}
async fn set(&self, state: &UploadState, mode: WriteMode) -> Result<()> {
validate_upload_id(state.id())?;
let mut states = self.states.write().map_err(|_| poisoned())?;
if mode == WriteMode::CreateNew && states.contains_key(state.id()) {
return Err(Error::AlreadyExists(state.id().to_string()));
}
states.insert(state.id().to_string(), state.clone());
Ok(())
}
async fn get(&self, id: &str) -> Result<Option<UploadState>> {
validate_upload_id(id)?;
let states = self.states.read().map_err(|_| poisoned())?;
Ok(states.get(id).cloned())
}
async fn delete(&self, id: &str) -> Result<()> {
validate_upload_id(id)?;
self.states.write().map_err(|_| poisoned())?.remove(id);
Ok(())
}
async fn list_expired(&self, before: DateTime<Utc>) -> Result<Vec<String>> {
let states = self.states.read().map_err(|_| poisoned())?;
let expired: Vec<String> = states
.values()
.filter(|s| s.expires_before(before))
.map(|s| s.id().to_string())
.collect();
Ok(expired)
}
}
#[async_trait]
impl UploadInventory for MemoryStateStore {
async fn list_upload_ids(&self, limit: usize, offset: usize) -> Result<Vec<String>> {
let states = self.states.read().map_err(|_| poisoned())?;
let mut ids: Vec<String> = states.keys().cloned().collect();
ids.sort();
Ok(ids.into_iter().skip(offset).take(limit).collect())
}
}
fn poisoned() -> Error {
Error::Internal("memory state store lock poisoned".to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Duration;
#[tokio::test]
async fn state_store_conformance() {
let store = MemoryStateStore::new();
crate::state::conformance::assert_state_store_semantics(&store).await;
}
#[tokio::test]
async fn upload_inventory_conformance() {
let store = MemoryStateStore::new();
crate::state::conformance::assert_upload_inventory_semantics(&store).await;
}
#[tokio::test]
async fn test_set_and_get() {
let store = MemoryStateStore::new();
let state = UploadState::new("test-1").with_length(1000);
store.set(&state, WriteMode::CreateNew).await.unwrap();
let retrieved = store.get("test-1").await.unwrap().unwrap();
assert_eq!(retrieved.id(), "test-1");
assert_eq!(retrieved.length(), Some(1000));
}
#[tokio::test]
async fn test_create_duplicate() {
let store = MemoryStateStore::new();
let state = UploadState::new("test-1");
store.set(&state, WriteMode::CreateNew).await.unwrap();
let result = store.set(&state, WriteMode::CreateNew).await;
assert!(matches!(result, Err(Error::AlreadyExists(_))));
}
#[tokio::test]
async fn test_update_existing() {
let store = MemoryStateStore::new();
let state = UploadState::new("test-1").with_length(1000);
store.set(&state, WriteMode::CreateNew).await.unwrap();
let mut updated = state.clone();
updated.set_offset(500);
store.set(&updated, WriteMode::Update).await.unwrap();
let retrieved = store.get("test-1").await.unwrap().unwrap();
assert_eq!(retrieved.offset(), 500);
}
#[tokio::test]
async fn test_get_not_found() {
let store = MemoryStateStore::new();
let result = store.get("nonexistent").await.unwrap();
assert!(result.is_none());
}
#[tokio::test]
async fn test_delete() {
let store = MemoryStateStore::new();
let state = UploadState::new("test-1");
store.set(&state, WriteMode::CreateNew).await.unwrap();
assert_eq!(store.len(), 1);
store.delete("test-1").await.unwrap();
assert_eq!(store.len(), 0);
assert!(store.get("test-1").await.unwrap().is_none());
}
#[tokio::test]
async fn test_list_expired() {
let store = MemoryStateStore::new();
let state1 =
UploadState::new("not-expired").with_expiration(Utc::now() + Duration::hours(1));
store.set(&state1, WriteMode::CreateNew).await.unwrap();
let state2 = UploadState::new("expired").with_expiration(Utc::now() - Duration::hours(1));
store.set(&state2, WriteMode::CreateNew).await.unwrap();
let state3 = UploadState::new("no-expiration");
store.set(&state3, WriteMode::CreateNew).await.unwrap();
let expired = store.list_expired(Utc::now()).await.unwrap();
assert_eq!(expired.len(), 1);
assert!(expired.contains(&"expired".to_string()));
}
#[tokio::test]
async fn test_clear() {
let store = MemoryStateStore::new();
for i in 0..5 {
let state = UploadState::new(format!("upload-{}", i));
store.set(&state, WriteMode::CreateNew).await.unwrap();
}
assert_eq!(store.len(), 5);
store.clear();
assert_eq!(store.len(), 0);
}
}