feat: init

2026-06-07 11:30:56 +08:00
commit 563381c1ca
361 changed files with 41327 additions and 0 deletions
@@ -0,0 +1,117 @@
+use std::sync::Arc;
+use std::sync::atomic::{AtomicI64, Ordering};
+
+use dashmap::DashMap;
+use tokio::sync::Mutex;
+use uuid::Uuid;
+
+use crate::cache::redis::AppRedis;
+use crate::error::{AppError, AppResult};
+use ::redis::Cmd;
+
+use super::redis_keys::*;
+
+struct Segment {
+    end: i64,
+    next: AtomicI64,
+}
+
+pub struct SeqAllocator {
+    redis: AppRedis,
+    segments: DashMap<Uuid, Arc<Segment>>,
+    locks: DashMap<Uuid, Arc<Mutex<()>>>,
+    segment_size: u64,
+}
+
+const MAX_RETRIES: u32 = 3;
+
+impl SeqAllocator {
+    pub fn new(redis: AppRedis) -> Self {
+        Self {
+            redis,
+            segments: DashMap::new(),
+            locks: DashMap::new(),
+            segment_size: WS_SEQ_SEGMENT_SIZE,
+        }
+    }
+
+    pub async fn next(&self, channel_id: Uuid) -> AppResult<i64> {
+        for _ in 0..MAX_RETRIES {
+            if let Some(seq) = self.try_allocate(&channel_id) {
+                return Ok(seq);
+            }
+            let lock = self
+                .locks
+                .entry(channel_id)
+                .or_insert_with(|| Arc::new(Mutex::new(())))
+                .clone();
+            let _guard = lock.lock().await;
+            if let Some(seq) = self.try_allocate(&channel_id) {
+                return Ok(seq);
+            }
+            self.refresh(channel_id).await?;
+        }
+        Err(AppError::InternalServerError(
+            "seq allocation exhausted retries".into(),
+        ))
+    }
+
+    pub async fn bootstrap(&self, channel_id: Uuid, db_max: i64) -> AppResult<i64> {
+        let key = format!("{WS_SEQ_PREFIX}{channel_id}");
+        let mut conn = self.redis.get_connection()?;
+        let current: i64 = Cmd::new()
+            .arg("SET")
+            .arg(&key)
+            .arg(db_max)
+            .arg("NX")
+            .arg("EX")
+            .arg(86400)
+            .query::<Option<String>>(&mut *conn.inner_mut())
+            .map_err(AppError::Redis)?
+            .and_then(|v| v.parse().ok())
+            .unwrap_or_else(|| {
+                let existing: i64 = Cmd::new()
+                    .arg("GET")
+                    .arg(&key)
+                    .query(&mut *conn.inner_mut())
+                    .map_err(AppError::Redis)
+                    .unwrap_or(db_max);
+                if existing < db_max { db_max } else { existing }
+            });
+        self.segments.remove(&channel_id);
+        Ok(current)
+    }
+
+    fn try_allocate(&self, channel_id: &Uuid) -> Option<i64> {
+        let state = self.segments.get(channel_id)?;
+        let next = state.next.fetch_add(1, Ordering::Relaxed);
+        if next < state.end { Some(next) } else { None }
+    }
+
+    async fn refresh(&self, channel_id: Uuid) -> AppResult<()> {
+        let key = format!("{WS_SEQ_PREFIX}{channel_id}");
+        let mut conn = self.redis.get_connection()?;
+        let counter: i64 = Cmd::new()
+            .arg("INCRBY")
+            .arg(&key)
+            .arg(self.segment_size as i64)
+            .query(&mut *conn.inner_mut())
+            .map_err(AppError::Redis)?;
+
+        let start = counter - self.segment_size as i64 + 1;
+        let end = counter + 1;
+        self.segments.insert(
+            channel_id,
+            Arc::new(Segment {
+                end,
+                next: AtomicI64::new(start),
+            }),
+        );
+        let _ = Cmd::new()
+            .arg("EXPIRE")
+            .arg(&key)
+            .arg(86400_u64)
+            .query::<()>(&mut *conn.inner_mut());
+        Ok(())
+    }
+}