use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicI64, Ordering}; use dashmap::DashMap; use tokio::sync::Mutex; use crate::error::AppTransportError; use db::cache::AppCache; use db::database::AppDatabase; use models::RoomId; const SEQ_KEY_PREFIX: &str = "room:seq:"; const DEFAULT_SEGMENT_SIZE: u64 = 1024; const MAX_REFRESH_RETRIES: u32 = 3; const BOOTSTRAP_SCRIPT: &str = r#" local key = KEYS[1] local db_max = tonumber(ARGV[1]) local current = tonumber(redis.call('GET', key) or '0') if current < db_max then redis.call('SET', key, db_max) end return tonumber(redis.call('GET', key)) "#; struct SegmentState { end: i64, next: AtomicI64, } pub struct SeqAllocator(Arc); struct SeqAllocatorInner { cache: AppCache, db: AppDatabase, segments: DashMap>, refresh_locks: DashMap>>, segment_size: u64, } impl Clone for SeqAllocator { fn clone(&self) -> Self { Self(Arc::clone(&self.0)) } } impl SeqAllocator { pub fn new(cache: AppCache, db: AppDatabase) -> Self { Self::with_segment_size(cache, db, DEFAULT_SEGMENT_SIZE) } pub fn with_segment_size(cache: AppCache, db: AppDatabase, size: u64) -> Self { Self(Arc::new(SeqAllocatorInner { cache, db, segments: DashMap::new(), refresh_locks: DashMap::new(), segment_size: if size > 0 { size } else { DEFAULT_SEGMENT_SIZE }, })) } pub async fn seq(&self, room: RoomId) -> Result { for _ in 0..MAX_REFRESH_RETRIES { if let Some(seq) = self.try_allocate(&room) { return Ok(seq); } let lock = self.get_refresh_lock(room); let _guard = lock.lock().await; if let Some(seq) = self.try_allocate(&room) { return Ok(seq); } self.refresh_segment(room).await?; } Err(AppTransportError::Internal) } pub async fn bootstrap(&self, room: RoomId) -> Result { let db_max = self.db_max_seq(room).await?; if db_max == 0 { return Ok(0); } let key = format!("{}{}", SEQ_KEY_PREFIX, room); let mut conn = self .0 .cache .conn() .await .map_err(|_| AppTransportError::Internal)?; let current: i64 = redis::Cmd::new() .arg("EVAL") .arg(BOOTSTRAP_SCRIPT) .arg(1) .arg(&key) .arg(db_max) .query_async(&mut conn) .await .map_err(|_| AppTransportError::Internal)?; self.0.segments.remove(&room); Ok(current) } pub async fn bootstrap_all( &self, rooms: Vec, ) -> Result, AppTransportError> { let mut results = HashMap::with_capacity(rooms.len()); for room in rooms { results.insert(room, self.bootstrap(room).await?); } Ok(results) } fn try_allocate(&self, room: &RoomId) -> Option { let state = self.0.segments.get(room)?; let next = state.next.fetch_add(1, Ordering::Relaxed); if next < state.end { Some(next) } else { None } } fn get_refresh_lock(&self, room: RoomId) -> Arc> { Arc::clone( self.0 .refresh_locks .entry(room) .or_insert_with(|| Arc::new(Mutex::new(()))) .value(), ) } async fn refresh_segment(&self, room: RoomId) -> Result<(), AppTransportError> { let key = format!("{}{}", SEQ_KEY_PREFIX, room); let mut conn = self .0 .cache .conn() .await .map_err(|_| AppTransportError::Internal)?; let counter: i64 = redis::Cmd::new() .arg("INCRBY") .arg(&key) .arg(self.0.segment_size as i64) .query_async(&mut conn) .await .map_err(|_| AppTransportError::Internal)?; let start = counter - self.0.segment_size as i64 + 1; let end = counter + 1; self.0.segments.insert( room, Arc::new(SegmentState { end, next: AtomicI64::new(start), }), ); Ok(()) } async fn db_max_seq(&self, room: RoomId) -> Result { use models::rooms::room_message; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter, QueryOrder, QuerySelect}; let last_msg = room_message::Entity::find() .filter(room_message::Column::Room.eq(room)) .order_by_desc(room_message::Column::Seq) .limit(1) .one(&self.0.db) .await .map_err(|_| AppTransportError::Internal)?; Ok(last_msg.map(|m| m.seq).unwrap_or(0)) } }