use crate::state::StreamState;
use parking_lot::{Mutex, RwLock};
use std::collections::HashMap;
use std::sync::Arc;
pub(crate) struct StreamStateMap {
inner: RwLock<HashMap<String, Arc<Mutex<StreamState>>>>,
}
impl StreamStateMap {
pub fn new(initial_streams: HashMap<String, StreamState>) -> Self {
let mut map = HashMap::new();
for (k, v) in initial_streams {
map.insert(k, Arc::new(Mutex::new(v)));
}
Self {
inner: RwLock::new(map),
}
}
pub fn get_or_create(&self, key: &str) -> Arc<Mutex<StreamState>> {
{
let map = self.inner.read();
if let Some(state) = map.get(key) {
return state.clone();
}
}
let mut map = self.inner.write();
if let Some(state) = map.get(key) {
return state.clone();
}
let new_state = Arc::new(Mutex::new(StreamState {
next_id: 0,
active_segment_start_id: 0,
}));
map.insert(key.to_string(), new_state.clone());
new_state
}
pub fn get(&self, key: &str) -> Option<Arc<Mutex<StreamState>>> {
let map = self.inner.read();
map.get(key).cloned()
}
pub fn contains(&self, key: &str) -> bool {
let map = self.inner.read();
map.contains_key(key)
}
}