mod error;
mod r#impl;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
pub use error::{Error, Result};
use hipstr::HipStr;
use parking_lot::{Mutex, RawMutex, lock_api::RawMutex as _};
use ts_::sec as ts_sec;
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,
}
}
}
#[repr(align(64))]
#[derive(Debug)]
struct CachePadded<T> {
value: T,
}
impl<T> CachePadded<T> {
const fn new(value: T) -> Self {
Self { value }
}
}
impl<T> std::ops::Deref for CachePadded<T> {
type Target = T;
#[inline]
fn deref(&self) -> &Self::Target {
&self.value
}
}
#[derive(Debug)]
struct Fast {
id: CachePadded<AtomicU64>,
max: CachePadded<AtomicU64>,
lock: CachePadded<AtomicBool>,
}
impl Fast {
const fn new() -> Self {
Self {
id: CachePadded::new(AtomicU64::new(0)),
max: CachePadded::new(AtomicU64::new(0)),
lock: CachePadded::new(AtomicBool::new(false)),
}
}
}
struct LockGuard<'a>(&'a AtomicBool);
impl Drop for LockGuard<'_> {
#[inline]
fn drop(&mut self) {
self.0.store(false, Ordering::Release);
}
}
pub struct KvId {
pub name: HipStr<'static>,
fast: Fast,
slow: Mutex<Slow>,
fetch_lock: tokio::sync::Mutex<()>,
}
impl KvId {
pub const fn const_new(name: &'static str) -> Self {
Self {
name: HipStr::borrowed(name),
fast: Fast::new(),
slow: Mutex::const_new(RawMutex::INIT, Slow::new()),
fetch_lock: tokio::sync::Mutex::const_new(()),
}
}
pub fn new(name: impl Into<HipStr<'static>>) -> Self {
Self {
name: name.into(),
fast: Fast::new(),
slow: Mutex::new(Slow::new()),
fetch_lock: tokio::sync::Mutex::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 _fetch_guard = self.fetch_lock.lock().await;
let _guard = LockGuard(&self.fast.lock);
let now = ts_sec();
let step = {
let s = self.slow.lock();
if s.next.is_some() {
return;
}
Self::calc_step(s.step, s.ts, now)
};
if let Ok(seg) = self.fetch(step).await {
let mut s = self.slow.lock();
if s.next.is_none() {
s.next = Some(seg);
}
}
}
fn try_switch_seg(&self) -> Option<u64> {
let seg = {
let mut s = self.slow.lock();
if self.fast.id.load(Ordering::Acquire) < self.fast.max.load(Ordering::Acquire) {
None
} else {
s.next.take()
}
};
if let Some(seg) = seg {
self.set_seg(seg);
self.try_next()
} else {
self.try_next()
}
}
pub async fn next(&'static self) -> Result<u64> {
loop {
if let Some(id) = self.try_next() {
return Ok(id);
}
if let Some(id) = self.try_switch_seg() {
self.spawn_fill();
return Ok(id);
}
let _fetch_guard = self.fetch_lock.lock().await;
if let Some(id) = self.try_switch_seg() {
self.spawn_fill();
return Ok(id);
}
let now = ts_sec();
let step = {
let s = self.slow.lock();
Self::calc_step(s.step, s.ts, now).max(s.step)
};
let seg = self.fetch(step).await?;
self.set_seg(seg);
{
let mut s = self.slow.lock();
s.step = step;
s.ts = now;
}
self.spawn_fill();
}
}
}