zhenyi
118 lines
3.4 KiB
Rust
118 lines
3.4 KiB
Rust
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(())
|
|
}
|
|
}
|