use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct SessionInfo {
pub id: String,
pub created_at: Instant,
pub last_connected_at: Instant,
pub props: HashMap<String, serde_json::Value>,
}
pub struct SessionManagerConfig {
pub ttl: Duration,
pub generate_id: Option<Box<dyn Fn() -> String + Send + Sync>>,
}
impl Default for SessionManagerConfig {
fn default() -> Self {
Self {
ttl: Duration::from_secs(3600),
generate_id: None,
}
}
}
pub struct PendingSession {
pub info: SessionInfo,
pub saved_state: serde_json::Value,
cancel: Arc<Mutex<bool>>,
}
pub struct SessionManager {
ttl: Duration,
generate_id: Box<dyn Fn() -> String + Send + Sync>,
inner: Mutex<Inner>,
}
struct Inner {
active: HashMap<String, SessionInfo>,
pending: HashMap<String, PendingEntry>,
connections: HashMap<String, HashSet<u64>>,
}
struct PendingEntry {
info: SessionInfo,
saved_state: serde_json::Value,
cancel: Arc<Mutex<bool>>,
}
fn default_generate_id() -> String {
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::SystemTime;
static COUNTER: AtomicU64 = AtomicU64::new(0);
let ns = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap_or_default()
.as_nanos() as u64;
let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
format!("{:016x}{:04x}", ns, seq & 0xFFFF)
}
impl SessionManager {
pub fn new(config: SessionManagerConfig) -> Self {
let generate_id = config
.generate_id
.unwrap_or_else(|| Box::new(default_generate_id));
Self {
ttl: config.ttl,
generate_id,
inner: Mutex::new(Inner {
active: HashMap::new(),
pending: HashMap::new(),
connections: HashMap::new(),
}),
}
}
pub fn create_session(
&self,
props: HashMap<String, serde_json::Value>,
) -> SessionInfo {
let mut inner = self.inner.lock().unwrap();
let mut id = (self.generate_id)();
for _ in 0..10 {
if !inner.active.contains_key(&id) && !inner.pending.contains_key(&id) {
break;
}
id = (self.generate_id)();
}
let now = Instant::now();
let info = SessionInfo {
id: id.clone(),
created_at: now,
last_connected_at: now,
props,
};
inner.active.insert(id, info.clone());
info
}
pub fn get_active_session(&self, id: &str) -> Option<SessionInfo> {
self.inner.lock().unwrap().active.get(id).cloned()
}
pub fn suspend_session<F>(&self, id: &str, saved_state: serde_json::Value, on_expire: F) -> bool
where
F: FnOnce() + Send + 'static,
{
let mut inner = self.inner.lock().unwrap();
let info = match inner.active.remove(id) {
Some(s) => s,
None => return false,
};
let cancel = Arc::new(Mutex::new(false));
let entry = PendingEntry {
info: info.clone(),
saved_state,
cancel: Arc::clone(&cancel),
};
inner.pending.insert(id.to_string(), entry);
drop(inner);
let ttl = self.ttl;
let id_owned = id.to_string();
let inner_ref = &self.inner as *const Mutex<Inner>;
let inner_ptr = inner_ref as usize;
std::thread::spawn(move || {
std::thread::sleep(ttl);
if *cancel.lock().unwrap() {
return;
}
let inner: &Mutex<Inner> = unsafe { &*(inner_ptr as *const Mutex<Inner>) };
let mut guard = inner.lock().unwrap();
if guard.pending.remove(&id_owned).is_some() {
drop(guard);
on_expire();
}
});
true
}
pub fn resume_session(&self, id: &str) -> Option<PendingSession> {
let mut inner = self.inner.lock().unwrap();
let entry = inner.pending.remove(id)?;
*entry.cancel.lock().unwrap() = true;
let mut info = entry.info;
info.last_connected_at = Instant::now();
inner.active.insert(id.to_string(), info.clone());
Some(PendingSession {
info,
saved_state: entry.saved_state,
cancel: entry.cancel,
})
}
pub fn destroy_session(&self, id: &str) {
let mut inner = self.inner.lock().unwrap();
inner.active.remove(id);
if let Some(entry) = inner.pending.remove(id) {
*entry.cancel.lock().unwrap() = true;
}
inner.connections.remove(id);
}
pub fn track_connection(&self, session_id: &str, conn_id: u64) {
let mut inner = self.inner.lock().unwrap();
inner
.connections
.entry(session_id.to_string())
.or_default()
.insert(conn_id);
}
pub fn untrack_connection(&self, session_id: &str, conn_id: u64) {
let mut inner = self.inner.lock().unwrap();
if let Some(conns) = inner.connections.get_mut(session_id) {
conns.remove(&conn_id);
}
}
pub fn connection_count(&self, session_id: &str) -> usize {
self.inner
.lock()
.unwrap()
.connections
.get(session_id)
.map(|c| c.len())
.unwrap_or(0)
}
pub fn shutdown(&self) {
let mut inner = self.inner.lock().unwrap();
for (_, entry) in inner.pending.drain() {
*entry.cancel.lock().unwrap() = true;
}
inner.active.clear();
inner.connections.clear();
}
}