mod error;
mod r#impl;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
pub use error::{Error, Result};
use parking_lot::{Mutex, RawMutex, lock_api::RawMutex as _};
use smol_str::SmolStr;
pub const PRELOAD_SEC: u64 = 60;
pub const STEP_MIN: u64 = 1;
pub const STEP_MAX: u64 = 1000000;
pub const KVID_KEY: &str = "kvid";
#[derive(Debug, Clone, Copy)]
struct Seg {
id: u64,
max: u64,
}
#[derive(Debug)]
struct Slow {
next: Option<Seg>,
step: u64,
ts: u64,
}
impl Slow {
const fn new() -> Self {
Self {
next: None,
step: STEP_MIN,
ts: 0,
}
}
}
#[derive(Debug)]
struct Fast {
id: AtomicU64,
max: AtomicU64,
lock: AtomicBool,
}
impl Fast {
const fn new() -> Self {
Self {
id: AtomicU64::new(0),
max: AtomicU64::new(0),
lock: AtomicBool::new(false),
}
}
}
pub struct KvId {
pub name: SmolStr,
fast: Fast,
slow: Mutex<Slow>,
}
impl KvId {
pub const fn const_new(name: &'static str) -> Self {
Self {
name: SmolStr::new_inline(name),
fast: Fast::new(),
slow: Mutex::const_new(RawMutex::INIT, Slow::new()),
}
}
pub fn new(name: impl Into<SmolStr>) -> Self {
Self {
name: name.into(),
fast: Fast::new(),
slow: Mutex::new(Slow::new()),
}
}
fn spawn_fill(&'static self) {
if self.fast.lock.swap(true, Ordering::Acquire) {
return;
}
tokio::spawn(async move {
self.fill().await;
});
}
async fn fill(&self) {
let now = ts_::sec();
let step = {
let s = self.slow.lock();
if s.next.is_some() {
self.fast.lock.store(false, Ordering::Release);
return;
}
Self::calc_step(&s, now)
};
if let Ok(seg) = self.fetch(step).await {
let mut s = self.slow.lock();
if s.next.is_none() {
s.next = Some(seg);
}
}
self.fast.lock.store(false, Ordering::Release);
}
pub async fn next(&'static self) -> Result<u64> {
if let Some(id) = self.try_next() {
if !self.fast.lock.load(Ordering::Relaxed) {
self.spawn_fill();
}
return Ok(id);
}
{
let mut s = self.slow.lock();
if let Some(id) = self.try_next() {
return Ok(id);
}
if let Some(seg) = s.next.take() {
self.set_seg(seg);
let id = self.try_next().ok_or(Error::Empty)?;
self.spawn_fill();
return Ok(id);
}
}
let now = ts_::sec();
let step = {
let s = self.slow.lock();
Self::calc_step(&s, now).max(s.step)
};
let seg = self.fetch(step).await?;
{
let mut s = self.slow.lock();
self.set_seg(seg);
s.step = step;
s.ts = now;
}
let id = self.try_next().ok_or(Error::Empty)?;
self.spawn_fill();
Ok(id)
}
}