use std::collections::HashMap;
use std::sync::Mutex;
use std::sync::atomic::{AtomicI64, Ordering};
use serde::{Deserialize, Serialize};
use crate::error::LiteError;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiteSequenceDef {
pub name: String,
pub start_value: i64,
pub increment: i64,
pub min_value: i64,
pub max_value: i64,
pub cycle: bool,
}
struct LiteSeqHandle {
counter: AtomicI64,
called: std::sync::atomic::AtomicBool,
def: LiteSequenceDef,
}
pub struct LiteSequenceRegistry {
sequences: Mutex<HashMap<String, LiteSeqHandle>>,
}
impl LiteSequenceRegistry {
pub fn new() -> Self {
Self {
sequences: Mutex::new(HashMap::new()),
}
}
pub fn register(&self, def: LiteSequenceDef) {
let initial = def.start_value - def.increment;
let mut map = self.sequences.lock().unwrap_or_else(|p| p.into_inner());
map.insert(
def.name.clone(),
LiteSeqHandle {
counter: AtomicI64::new(initial),
called: std::sync::atomic::AtomicBool::new(false),
def,
},
);
}
pub fn nextval(&self, name: &str) -> Result<i64, LiteError> {
let map = self.sequences.lock().unwrap_or_else(|p| p.into_inner());
let handle = map
.get(name)
.ok_or_else(|| LiteError::Query(format!("sequence \"{name}\" does not exist")))?;
let inc = handle.def.increment;
let prev = handle.counter.fetch_add(inc, Ordering::Relaxed);
let new_val = prev + inc;
handle.called.store(true, Ordering::Relaxed);
if inc > 0 && new_val > handle.def.max_value {
if handle.def.cycle {
handle
.counter
.store(handle.def.min_value, Ordering::Relaxed);
return Ok(handle.def.min_value);
}
handle.counter.store(prev, Ordering::Relaxed);
return Err(LiteError::Query(format!("sequence \"{name}\" exhausted")));
}
if inc < 0 && new_val < handle.def.min_value {
if handle.def.cycle {
handle
.counter
.store(handle.def.max_value, Ordering::Relaxed);
return Ok(handle.def.max_value);
}
handle.counter.store(prev, Ordering::Relaxed);
return Err(LiteError::Query(format!("sequence \"{name}\" exhausted")));
}
Ok(new_val)
}
pub fn currval(&self, name: &str) -> Result<i64, LiteError> {
let map = self.sequences.lock().unwrap_or_else(|p| p.into_inner());
let handle = map
.get(name)
.ok_or_else(|| LiteError::Query(format!("sequence \"{name}\" does not exist")))?;
if !handle.called.load(Ordering::Relaxed) {
return Err(LiteError::Query(format!(
"currval of \"{name}\" not yet defined"
)));
}
Ok(handle.counter.load(Ordering::Relaxed))
}
pub fn setval(&self, name: &str, value: i64) -> Result<i64, LiteError> {
let map = self.sequences.lock().unwrap_or_else(|p| p.into_inner());
let handle = map
.get(name)
.ok_or_else(|| LiteError::Query(format!("sequence \"{name}\" does not exist")))?;
if value < handle.def.min_value || value > handle.def.max_value {
return Err(LiteError::Query(format!(
"setval: {value} outside [{}, {}]",
handle.def.min_value, handle.def.max_value
)));
}
handle.counter.store(value, Ordering::Relaxed);
handle.called.store(true, Ordering::Relaxed);
Ok(value)
}
pub fn restart_collection_sequences(&self, collection: &str) -> Vec<String> {
let prefix = format!("{collection}_");
let map = self.sequences.lock().unwrap_or_else(|p| p.into_inner());
let mut restarted = Vec::new();
for (name, handle) in map.iter() {
if !(name.starts_with(&prefix) && name.ends_with("_seq")) {
continue;
}
handle.counter.store(
handle.def.start_value - handle.def.increment,
Ordering::Relaxed,
);
handle.called.store(false, Ordering::Relaxed);
restarted.push(name.clone());
}
restarted
}
pub fn names(&self) -> Vec<String> {
let map = self.sequences.lock().unwrap_or_else(|p| p.into_inner());
map.keys().cloned().collect()
}
pub fn remove(&self, name: &str) -> bool {
let mut map = self.sequences.lock().unwrap_or_else(|p| p.into_inner());
map.remove(name).is_some()
}
}
impl Default for LiteSequenceRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_def(name: &str) -> LiteSequenceDef {
LiteSequenceDef {
name: name.to_string(),
start_value: 1,
increment: 1,
min_value: 1,
max_value: i64::MAX,
cycle: false,
}
}
#[test]
fn basic_nextval() {
let reg = LiteSequenceRegistry::new();
reg.register(make_def("s1"));
assert_eq!(reg.nextval("s1").unwrap(), 1);
assert_eq!(reg.nextval("s1").unwrap(), 2);
assert_eq!(reg.currval("s1").unwrap(), 2);
}
#[test]
fn setval_and_nextval() {
let reg = LiteSequenceRegistry::new();
reg.register(make_def("s1"));
reg.setval("s1", 50).unwrap();
assert_eq!(reg.nextval("s1").unwrap(), 51);
}
#[test]
fn restart_collection_sequences_makes_start_the_next_value() {
let reg = LiteSequenceRegistry::new();
reg.register(make_def("orders_id_seq"));
reg.register(make_def("orders_line_seq"));
reg.register(make_def("orders_archive_id_seq"));
reg.register(make_def("users_id_seq"));
for _ in 0..3 {
reg.nextval("orders_id_seq").unwrap();
reg.nextval("orders_line_seq").unwrap();
reg.nextval("orders_archive_id_seq").unwrap();
reg.nextval("users_id_seq").unwrap();
}
let mut restarted = reg.restart_collection_sequences("orders");
restarted.sort();
assert_eq!(
restarted,
vec![
"orders_archive_id_seq".to_string(),
"orders_id_seq".to_string(),
"orders_line_seq".to_string()
]
);
assert!(
reg.currval("orders_id_seq").is_err(),
"restart clears currval"
);
assert_eq!(reg.nextval("orders_id_seq").unwrap(), 1);
assert_eq!(reg.nextval("orders_line_seq").unwrap(), 1);
assert_eq!(reg.nextval("users_id_seq").unwrap(), 4);
}
#[test]
fn not_found() {
let reg = LiteSequenceRegistry::new();
assert!(reg.nextval("nope").is_err());
}
}