rustack_core/
state.rs

1//! Multi-account, multi-region state management.
2//!
3//! Provides [`AccountRegionStore`], a thread-safe concurrent store that
4//! partitions state by AWS account ID and region, matching LocalStack's
5//! `AccountRegionBundle` pattern.
6
7use std::sync::Arc;
8
9use dashmap::DashMap;
10
11use crate::types::{AccountId, AwsRegion};
12
13/// Thread-safe, multi-account, multi-region state store.
14///
15/// Each (account, region) pair gets its own isolated state instance of type `T`.
16/// Uses `DashMap` for lock-free concurrent access.
17///
18/// # Examples
19///
20/// ```
21/// use rustack_core::{AccountRegionStore, AccountId, AwsRegion};
22///
23/// #[derive(Debug, Default)]
24/// struct MyServiceState {
25///     counter: std::sync::atomic::AtomicU64,
26/// }
27///
28/// let store = AccountRegionStore::<MyServiceState>::new();
29/// let state = store.get_or_create(&AccountId::default(), &AwsRegion::default());
30/// state.counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
31/// ```
32#[derive(Debug)]
33pub struct AccountRegionStore<T: Default + Send + Sync> {
34    inner: DashMap<(AccountId, AwsRegion), Arc<T>>,
35}
36
37impl<T: Default + Send + Sync> AccountRegionStore<T> {
38    /// Create a new empty store.
39    #[must_use]
40    pub fn new() -> Self {
41        Self {
42            inner: DashMap::new(),
43        }
44    }
45
46    /// Get or create the state for the given account and region.
47    ///
48    /// If the state does not exist, a new default instance is created atomically.
49    #[must_use]
50    pub fn get_or_create(&self, account: &AccountId, region: &AwsRegion) -> Arc<T> {
51        self.inner
52            .entry((account.clone(), region.clone()))
53            .or_insert_with(|| Arc::new(T::default()))
54            .clone()
55    }
56
57    /// Get the state for the given account and region, if it exists.
58    #[must_use]
59    pub fn get(&self, account: &AccountId, region: &AwsRegion) -> Option<Arc<T>> {
60        self.inner
61            .get(&(account.clone(), region.clone()))
62            .map(|v| v.clone())
63    }
64
65    /// Remove the state for the given account and region.
66    #[must_use]
67    pub fn remove(&self, account: &AccountId, region: &AwsRegion) -> Option<Arc<T>> {
68        self.inner
69            .remove(&(account.clone(), region.clone()))
70            .map(|(_, v)| v)
71    }
72
73    /// Reset all state in the store.
74    pub fn reset(&self) {
75        self.inner.clear();
76    }
77
78    /// Number of (account, region) entries.
79    #[must_use]
80    pub fn len(&self) -> usize {
81        self.inner.len()
82    }
83
84    /// Whether the store is empty.
85    #[must_use]
86    pub fn is_empty(&self) -> bool {
87        self.inner.is_empty()
88    }
89}
90
91impl<T: Default + Send + Sync> Default for AccountRegionStore<T> {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[derive(Debug, Default)]
102    struct TestState {
103        value: std::sync::atomic::AtomicU64,
104    }
105
106    #[test]
107    fn test_should_create_state_on_first_access() {
108        let store = AccountRegionStore::<TestState>::new();
109        let account = AccountId::default();
110        let region = AwsRegion::default();
111
112        assert!(store.is_empty());
113        let state = store.get_or_create(&account, &region);
114        assert_eq!(store.len(), 1);
115        assert_eq!(state.value.load(std::sync::atomic::Ordering::Relaxed), 0);
116    }
117
118    #[test]
119    fn test_should_return_same_state_on_subsequent_access() {
120        let store = AccountRegionStore::<TestState>::new();
121        let account = AccountId::default();
122        let region = AwsRegion::default();
123
124        let state1 = store.get_or_create(&account, &region);
125        state1.value.store(42, std::sync::atomic::Ordering::Relaxed);
126
127        let state2 = store.get_or_create(&account, &region);
128        assert_eq!(state2.value.load(std::sync::atomic::Ordering::Relaxed), 42);
129    }
130
131    #[test]
132    fn test_should_isolate_different_regions() {
133        let store = AccountRegionStore::<TestState>::new();
134        let account = AccountId::default();
135        let us_east = AwsRegion::new("us-east-1");
136        let eu_west = AwsRegion::new("eu-west-1");
137
138        let state_us = store.get_or_create(&account, &us_east);
139        state_us
140            .value
141            .store(1, std::sync::atomic::Ordering::Relaxed);
142
143        let state_eu = store.get_or_create(&account, &eu_west);
144        assert_eq!(state_eu.value.load(std::sync::atomic::Ordering::Relaxed), 0);
145        assert_eq!(store.len(), 2);
146    }
147
148    #[test]
149    fn test_should_reset_all_state() {
150        let store = AccountRegionStore::<TestState>::new();
151        let _ = store.get_or_create(&AccountId::default(), &AwsRegion::default());
152        let _ = store.get_or_create(&AccountId::default(), &AwsRegion::new("eu-west-1"));
153
154        assert_eq!(store.len(), 2);
155        store.reset();
156        assert!(store.is_empty());
157    }
158}
rustack_core/state.rs

rustack_core/
state.rs
