use std::time::{SystemTime, UNIX_EPOCH};
pub use zentinel_common::TraceIdFormat;
#[inline]
pub fn generate_for_format(format: TraceIdFormat) -> String {
match format {
TraceIdFormat::TinyFlake => generate_tinyflake(),
TraceIdFormat::Uuid => generate_uuid(),
}
}
const BASE58_ALPHABET: &[u8; 58] = b"123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";
pub const TINYFLAKE_LENGTH: usize = 11;
const TIME_COMPONENT_LENGTH: usize = 3;
const RANDOM_COMPONENT_LENGTH: usize = 8;
const TIME_MODULO: u64 = 195_112;
pub fn generate_tinyflake() -> String {
let mut id = String::with_capacity(TINYFLAKE_LENGTH);
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let time_component = (now % TIME_MODULO) as usize;
encode_base58(time_component, TIME_COMPONENT_LENGTH, &mut id);
let random_bytes: [u8; 6] = rand::random();
let random_value = u64::from_le_bytes([
random_bytes[0],
random_bytes[1],
random_bytes[2],
random_bytes[3],
random_bytes[4],
random_bytes[5],
0,
0,
]) as usize;
encode_base58(random_value, RANDOM_COMPONENT_LENGTH, &mut id);
id
}
fn encode_base58(mut value: usize, width: usize, output: &mut String) {
let mut chars = Vec::with_capacity(width);
for _ in 0..width {
chars.push(BASE58_ALPHABET[value % 58] as char);
value /= 58;
}
for c in chars.into_iter().rev() {
output.push(c);
}
}
pub fn generate_uuid() -> String {
uuid::Uuid::new_v4().to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
#[test]
fn test_tinyflake_format() {
let id = generate_tinyflake();
assert_eq!(
id.len(),
TINYFLAKE_LENGTH,
"TinyFlake should be {} chars, got: {} ({})",
TINYFLAKE_LENGTH,
id.len(),
id
);
for c in id.chars() {
assert!(
BASE58_ALPHABET.contains(&(c as u8)),
"Invalid char '{}' in TinyFlake: {}",
c,
id
);
}
assert!(!id.contains('0'), "TinyFlake should not contain '0'");
assert!(!id.contains('O'), "TinyFlake should not contain 'O'");
assert!(!id.contains('I'), "TinyFlake should not contain 'I'");
assert!(!id.contains('l'), "TinyFlake should not contain 'l'");
}
#[test]
fn test_tinyflake_uniqueness() {
let mut ids = HashSet::new();
for _ in 0..10_000 {
let id = generate_tinyflake();
assert!(
ids.insert(id.clone()),
"Duplicate TinyFlake generated: {}",
id
);
}
}
#[test]
fn test_tinyflake_time_ordering() {
let id1 = generate_tinyflake();
let id2 = generate_tinyflake();
assert_eq!(
&id1[..TIME_COMPONENT_LENGTH],
&id2[..TIME_COMPONENT_LENGTH],
"Time prefix should match within same second: {} vs {}",
id1,
id2
);
}
#[test]
fn test_uuid_format() {
let id = generate_uuid();
assert_eq!(id.len(), 36, "UUID should be 36 chars, got: {}", id.len());
assert_eq!(
id.matches('-').count(),
4,
"UUID should have 4 dashes: {}",
id
);
assert!(
uuid::Uuid::parse_str(&id).is_ok(),
"Should be valid UUID: {}",
id
);
}
#[test]
fn test_trace_id_format_generate() {
let tinyflake = generate_for_format(TraceIdFormat::TinyFlake);
assert_eq!(tinyflake.len(), TINYFLAKE_LENGTH);
let uuid = generate_for_format(TraceIdFormat::Uuid);
assert_eq!(uuid.len(), 36);
}
#[test]
fn test_trace_id_format_from_str() {
assert_eq!(
TraceIdFormat::from_str_loose("tinyflake"),
TraceIdFormat::TinyFlake
);
assert_eq!(
TraceIdFormat::from_str_loose("TINYFLAKE"),
TraceIdFormat::TinyFlake
);
assert_eq!(TraceIdFormat::from_str_loose("uuid"), TraceIdFormat::Uuid);
assert_eq!(TraceIdFormat::from_str_loose("UUID"), TraceIdFormat::Uuid);
assert_eq!(TraceIdFormat::from_str_loose("uuid4"), TraceIdFormat::Uuid);
assert_eq!(TraceIdFormat::from_str_loose("uuidv4"), TraceIdFormat::Uuid);
assert_eq!(
TraceIdFormat::from_str_loose("unknown"),
TraceIdFormat::TinyFlake
); }
#[test]
fn test_trace_id_format_display() {
assert_eq!(TraceIdFormat::TinyFlake.to_string(), "tinyflake");
assert_eq!(TraceIdFormat::Uuid.to_string(), "uuid");
}
#[test]
fn test_encode_base58() {
let mut output = String::new();
encode_base58(0, 3, &mut output);
assert_eq!(output, "111");
output.clear();
encode_base58(57, 3, &mut output);
assert_eq!(output, "11z");
output.clear();
encode_base58(58, 3, &mut output);
assert_eq!(output, "121");
}
#[test]
fn test_base58_alphabet_is_correct() {
let alphabet_str = std::str::from_utf8(BASE58_ALPHABET).unwrap();
assert!(!alphabet_str.contains('0'));
assert!(!alphabet_str.contains('O'));
assert!(!alphabet_str.contains('I'));
assert!(!alphabet_str.contains('l'));
assert_eq!(BASE58_ALPHABET.len(), 58);
let unique: HashSet<u8> = BASE58_ALPHABET.iter().copied().collect();
assert_eq!(unique.len(), 58);
}
}