use std::collections::BTreeMap;
use std::sync::Arc;
use async_trait::async_trait;
use rskit_errors::{AppError, AppResult, ErrorCode};
use crate::config::WorkloadConfig;
use crate::manager::Manager;
#[async_trait]
pub trait ManagerFactory: Send + Sync {
async fn create(&self, config: &WorkloadConfig) -> AppResult<Arc<dyn Manager>>;
}
#[derive(Default)]
pub struct WorkloadRegistry {
factories: BTreeMap<String, Arc<dyn ManagerFactory>>,
}
impl WorkloadRegistry {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn register(
&mut self,
name: impl Into<String>,
factory: Arc<dyn ManagerFactory>,
) -> AppResult<()> {
let name = name.into().trim().to_owned();
if name.is_empty() {
return Err(AppError::new(
ErrorCode::InvalidInput,
"workload provider name is required",
));
}
if self.factories.contains_key(&name) {
return Err(AppError::new(
ErrorCode::AlreadyExists,
format!("workload provider '{name}' is already registered"),
));
}
self.factories.insert(name, factory);
Ok(())
}
#[must_use]
pub fn contains(&self, name: &str) -> bool {
self.factories.contains_key(name.trim())
}
#[must_use]
pub fn len(&self) -> usize {
self.factories.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.factories.is_empty()
}
#[must_use]
pub fn names(&self) -> Vec<String> {
self.factories.keys().cloned().collect()
}
pub async fn build(&self, config: &WorkloadConfig) -> AppResult<Arc<dyn Manager>> {
config.validate()?;
let provider = config.provider.trim();
self.factories
.get(provider)
.ok_or_else(|| {
AppError::new(
ErrorCode::NotFound,
format!("workload provider '{provider}' is not registered"),
)
})?
.create(config)
.await
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use super::*;
use crate::test_support::FakeManager;
struct CountingFactory {
calls: Arc<AtomicUsize>,
}
#[async_trait]
impl ManagerFactory for CountingFactory {
async fn create(&self, _config: &WorkloadConfig) -> AppResult<Arc<dyn Manager>> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(Arc::new(FakeManager))
}
}
fn factory(calls: &Arc<AtomicUsize>) -> Arc<CountingFactory> {
Arc::new(CountingFactory {
calls: Arc::clone(calls),
})
}
#[tokio::test]
async fn registers_and_builds_selected_provider() {
let calls = Arc::new(AtomicUsize::new(0));
let mut registry = WorkloadRegistry::new();
assert!(registry.is_empty());
registry.register(" docker ", factory(&calls)).unwrap();
assert!(registry.contains("docker"));
assert!(registry.contains(" docker "));
assert_eq!(registry.len(), 1);
assert_eq!(registry.names(), vec!["docker".to_string()]);
let config = WorkloadConfig {
provider: "docker".to_string(),
..Default::default()
};
registry.build(&config).await.unwrap();
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn rejects_empty_and_duplicate_names() {
let calls = Arc::new(AtomicUsize::new(0));
let mut registry = WorkloadRegistry::new();
assert_eq!(
registry.register(" ", factory(&calls)).unwrap_err().code(),
ErrorCode::InvalidInput
);
registry.register("docker", factory(&calls)).unwrap();
assert_eq!(
registry
.register("docker", factory(&calls))
.unwrap_err()
.code(),
ErrorCode::AlreadyExists
);
}
#[tokio::test]
async fn build_reports_missing_and_unregistered_providers() {
let registry = WorkloadRegistry::new();
let empty = WorkloadConfig {
provider: String::new(),
..Default::default()
};
assert_eq!(
registry.build(&empty).await.map(|_| ()).unwrap_err().code(),
ErrorCode::MissingField
);
let unknown = WorkloadConfig {
provider: "podman".to_string(),
..Default::default()
};
assert_eq!(
registry
.build(&unknown)
.await
.map(|_| ())
.unwrap_err()
.code(),
ErrorCode::NotFound
);
}
}