use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use dashmap::DashMap;
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ModuleVersion {
pub major: u32,
pub minor: u32,
pub patch: u32,
}
impl ModuleVersion {
pub fn new(major: u32, minor: u32, patch: u32) -> Self {
Self { major, minor, patch }
}
pub fn parse(s: &str) -> Option<Self> {
let parts: Vec<&str> = s.split('.').collect();
if parts.len() != 3 {
return None;
}
Some(Self {
major: parts[0].parse().ok()?,
minor: parts[1].parse().ok()?,
patch: parts[2].parse().ok()?,
})
}
pub fn as_str(&self) -> String {
format!("{}.{}.{}", self.major, self.minor, self.patch)
}
pub fn is_compatible_with(&self, other: &ModuleVersion) -> bool {
self.major == other.major
}
pub fn is_greater_than(&self, other: &ModuleVersion) -> bool {
if self.major != other.major {
return self.major > other.major;
}
if self.minor != other.minor {
return self.minor > other.minor;
}
self.patch > other.patch
}
}
impl std::fmt::Display for ModuleVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
}
}
#[derive(Debug, Clone)]
pub struct VersionedModule {
pub name: String,
pub version: ModuleVersion,
pub path: String,
pub is_active: bool,
}
#[derive(Debug, Clone)]
pub struct TrafficRouting {
pub new_version_percentage: u32,
pub min_requests: u32,
pub error_threshold: f32,
}
impl Default for TrafficRouting {
fn default() -> Self {
Self {
new_version_percentage: 0,
min_requests: 100,
error_threshold: 0.01, }
}
}
#[derive(Debug)]
pub struct VersionManager {
module_name: String,
versions: DashMap<ModuleVersion, VersionedModule>,
active_version: std::sync::Mutex<Option<ModuleVersion>>,
routing: std::sync::Mutex<TrafficRouting>,
new_version_requests: AtomicU32,
new_version_errors: AtomicU32,
}
impl VersionManager {
pub fn new(module_name: String) -> Self {
Self {
module_name,
versions: DashMap::new(),
active_version: std::sync::Mutex::new(None),
routing: std::sync::Mutex::new(TrafficRouting::default()),
new_version_requests: AtomicU32::new(0),
new_version_errors: AtomicU32::new(0),
}
}
pub fn register_version(&self, version: ModuleVersion, path: String) {
let module = VersionedModule {
name: self.module_name.clone(),
version: version.clone(),
path,
is_active: false,
};
self.versions.insert(version, module);
}
pub fn activate_version(&self, version: &ModuleVersion) -> bool {
if !self.versions.contains_key(version) {
return false;
}
for entry in self.versions.iter() {
let mut module = entry.value().clone();
module.is_active = false;
}
*self.active_version.lock().unwrap() = Some(version.clone());
true
}
pub fn select_version(&self) -> Option<ModuleVersion> {
let routing = self.routing.lock().unwrap();
if routing.new_version_percentage == 0 {
return self.active_version.lock().unwrap().clone();
}
use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hasher};
let hasher = RandomState::new().build_hasher();
let roll = hasher.finish() % 100;
if roll < routing.new_version_percentage as u64 {
self.versions
.iter()
.max_by(|a, b| a.key().cmp(b.key()))
.map(|e| e.key().clone())
} else {
self.active_version.lock().unwrap().clone()
}
}
pub fn record_request(&self) {
self.new_version_requests.fetch_add(1, Ordering::Relaxed);
}
pub fn record_error(&self) {
self.new_version_errors.fetch_add(1, Ordering::Relaxed);
}
pub fn should_rollback(&self) -> bool {
let routing = self.routing.lock().unwrap();
let requests = self.new_version_requests.load(Ordering::Relaxed);
if requests < routing.min_requests {
return false;
}
let errors = self.new_version_errors.load(Ordering::Relaxed);
let error_rate = errors as f32 / requests as f32;
error_rate > routing.error_threshold
}
pub fn all_versions(&self) -> Vec<VersionedModule> {
self.versions.iter().map(|e| e.value().clone()).collect()
}
pub fn active_version(&self) -> Option<ModuleVersion> {
self.active_version.lock().unwrap().clone()
}
pub fn set_routing(&self, routing: TrafficRouting) {
*self.routing.lock().unwrap() = routing;
}
}
#[derive(Debug)]
pub struct VersionRegistry {
managers: DashMap<String, Arc<VersionManager>>,
}
impl VersionRegistry {
pub fn new() -> Self {
Self {
managers: DashMap::new(),
}
}
pub fn get_or_create(&self, module_name: &str) -> Arc<VersionManager> {
self.managers
.entry(module_name.to_string())
.or_insert_with(|| Arc::new(VersionManager::new(module_name.to_string())))
.clone()
}
pub fn module_count(&self) -> usize {
self.managers.len()
}
}
impl Default for VersionRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_version_parsing() {
let v = ModuleVersion::parse("1.2.3").unwrap();
assert_eq!(v.major, 1);
assert_eq!(v.minor, 2);
assert_eq!(v.patch, 3);
}
#[test]
fn test_version_comparison() {
let v1 = ModuleVersion::new(1, 0, 0);
let v2 = ModuleVersion::new(1, 0, 1);
let v3 = ModuleVersion::new(2, 0, 0);
assert!(v2.is_greater_than(&v1));
assert!(v3.is_greater_than(&v2));
assert!(v1.is_compatible_with(&v2));
assert!(!v1.is_compatible_with(&v3));
}
#[test]
fn test_version_manager() {
let manager = VersionManager::new("test".to_string());
manager.register_version(ModuleVersion::new(1, 0, 0), "/path/v1".to_string());
manager.register_version(ModuleVersion::new(2, 0, 0), "/path/v2".to_string());
manager.activate_version(&ModuleVersion::new(1, 0, 0));
assert_eq!(manager.active_version(), Some(ModuleVersion::new(1, 0, 0)));
assert_eq!(manager.all_versions().len(), 2);
}
}