use crate::error::{LighterError, Result};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
#[derive(Debug)]
pub struct NonceManager {
counter: AtomicU64,
last_timestamp: AtomicU64,
}
impl NonceManager {
pub fn new() -> Self {
Self {
counter: AtomicU64::new(0),
last_timestamp: AtomicU64::new(0),
}
}
pub fn generate(&self) -> Result<i64> {
let current_timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|e| LighterError::Nonce(format!("Time error: {}", e)))?
.as_millis() as u64;
loop {
let last_timestamp = self.last_timestamp.load(Ordering::Acquire);
if current_timestamp > last_timestamp {
if self
.last_timestamp
.compare_exchange(
last_timestamp,
current_timestamp,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
self.counter.store(0, Ordering::Release);
return Ok((current_timestamp * 1000) as i64);
} else {
continue;
}
} else {
let counter = self.counter.fetch_add(1, Ordering::AcqRel);
if counter >= 999 {
return Err(LighterError::Nonce(
"Too many nonces generated in the same millisecond".to_string(),
));
}
return Ok((last_timestamp * 1000 + counter + 1) as i64);
}
}
}
}
impl Default for NonceManager {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
use std::sync::Arc;
use std::thread;
#[test]
fn test_nonce_uniqueness() {
let manager = NonceManager::new();
let mut nonces = HashSet::new();
for _ in 0..100 {
let nonce = manager.generate().unwrap();
assert!(nonces.insert(nonce), "Duplicate nonce generated: {}", nonce);
}
}
#[test]
fn test_nonce_increasing() {
let manager = NonceManager::new();
let mut previous = 0;
for _ in 0..50 {
let nonce = manager.generate().unwrap();
assert!(
nonce > previous,
"Nonce {} is not greater than previous {}",
nonce,
previous
);
previous = nonce;
}
}
#[test]
fn test_nonce_concurrent_generation() {
let manager = Arc::new(NonceManager::new());
let mut handles = vec![];
let nonces = Arc::new(std::sync::Mutex::new(HashSet::new()));
for _ in 0..10 {
let manager_clone = Arc::clone(&manager);
let nonces_clone = Arc::clone(&nonces);
let handle = thread::spawn(move || {
for _ in 0..10 {
let nonce = manager_clone.generate().unwrap();
let mut nonces_guard = nonces_clone.lock().unwrap();
assert!(
nonces_guard.insert(nonce),
"Duplicate nonce in concurrent generation: {}",
nonce
);
}
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
let final_nonces = nonces.lock().unwrap();
assert_eq!(final_nonces.len(), 100);
}
#[test]
fn test_nonce_timestamp_format() {
let manager = NonceManager::new();
let nonce = manager.generate().unwrap();
assert!(
nonce > 1_000_000_000_000,
"Nonce should be at least 13 digits"
);
let current_time_micros = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_millis() as i64
* 1000;
assert!(nonce <= current_time_micros + 1_000_000);
assert!(nonce >= current_time_micros - 1_000_000);
}
#[test]
fn test_nonce_counter_reset() {
let manager = NonceManager::new();
let first_nonce = manager.generate().unwrap();
std::thread::sleep(std::time::Duration::from_millis(2));
let second_nonce = manager.generate().unwrap();
assert!(second_nonce > first_nonce);
assert!(second_nonce - first_nonce >= 1000);
}
#[test]
fn test_nonce_rapid_generation() {
let manager = NonceManager::new();
let mut nonces = Vec::new();
for _ in 0..100 {
match manager.generate() {
Ok(nonce) => nonces.push(nonce),
Err(e) => {
if let LighterError::Nonce(msg) = e {
assert!(msg.contains("Too many nonces"));
break;
} else {
panic!("Unexpected error: {}", e);
}
}
}
}
let unique_nonces: HashSet<_> = nonces.iter().collect();
assert_eq!(unique_nonces.len(), nonces.len());
}
#[test]
fn test_nonce_default_impl() {
let manager1 = NonceManager::new();
let manager2 = NonceManager::default();
assert!(manager1.generate().is_ok());
assert!(manager2.generate().is_ok());
}
}