gitdataai/libs/transport/seq.rs

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<SeqAllocatorInner>);

struct SeqAllocatorInner {
    cache: AppCache,
    db: AppDatabase,
    segments: DashMap<RoomId, Arc<SegmentState>>,
    refresh_locks: DashMap<RoomId, Arc<Mutex<()>>>,
    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<i64, AppTransportError> {
        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<i64, AppTransportError> {
        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<RoomId>,
    ) -> Result<HashMap<RoomId, i64>, 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<i64> {
        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<Mutex<()>> {
        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<i64, AppTransportError> {
        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))
    }
}