use crate::multi_tenancy::{
access_control::AccessControl,
billing::{BillingEngine, BillingMetrics, BillingPeriod, PricingModel},
isolation::{IsolationLevel, IsolationStrategy, NamespaceManager},
quota::{QuotaEnforcer, QuotaLimits, RateLimiter},
tenant::{Tenant, TenantId, TenantMetadata, TenantStatus},
types::{
MultiTenancyError, MultiTenancyResult, TenantContext, TenantOperation, TenantStatistics,
},
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TenantConfig {
pub metadata: TenantMetadata,
pub isolation: IsolationStrategy,
pub quotas: QuotaLimits,
pub pricing: PricingModel,
pub rate_limit: Option<f64>,
}
impl TenantConfig {
pub fn free_tier(tenant_id: impl Into<String>, name: impl Into<String>) -> Self {
let tenant_id = tenant_id.into();
Self {
metadata: TenantMetadata::new(name, "free"),
isolation: IsolationStrategy::free_tier(),
quotas: QuotaLimits::free_tier(&tenant_id),
pricing: PricingModel::PerRequest {
cost_per_request: 0.001,
},
rate_limit: Some(10.0),
}
}
pub fn pro_tier(tenant_id: impl Into<String>, name: impl Into<String>) -> Self {
let tenant_id = tenant_id.into();
Self {
metadata: TenantMetadata::new(name, "pro"),
isolation: IsolationStrategy::pro_tier(),
quotas: QuotaLimits::pro_tier(&tenant_id),
pricing: PricingModel::PerComputeUnit {
cost_per_unit: 0.01,
},
rate_limit: Some(100.0),
}
}
pub fn enterprise_tier(tenant_id: impl Into<String>, name: impl Into<String>) -> Self {
let tenant_id = tenant_id.into();
Self {
metadata: TenantMetadata::new(name, "enterprise"),
isolation: IsolationStrategy::enterprise_tier(),
quotas: QuotaLimits::enterprise_tier(&tenant_id),
pricing: PricingModel::Subscription {
monthly_fee: 1000.0,
included_requests: 1_000_000,
overage_cost: 0.005,
},
rate_limit: None, }
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TenantManagerConfig {
pub default_isolation: IsolationLevel,
pub billing_period: BillingPeriod,
pub strict_quotas: bool,
pub enable_access_control: bool,
pub enable_billing: bool,
pub auto_suspend_on_quota_exceeded: bool,
}
impl TenantManagerConfig {
pub fn default_config() -> Self {
Self {
default_isolation: IsolationLevel::Namespace,
billing_period: BillingPeriod::Monthly,
strict_quotas: true,
enable_access_control: true,
enable_billing: true,
auto_suspend_on_quota_exceeded: false,
}
}
pub fn production() -> Self {
Self {
default_isolation: IsolationLevel::SeparateIndex,
billing_period: BillingPeriod::Monthly,
strict_quotas: true,
enable_access_control: true,
enable_billing: true,
auto_suspend_on_quota_exceeded: true,
}
}
}
pub struct MultiTenantManager {
config: TenantManagerConfig,
tenants: Arc<RwLock<HashMap<TenantId, Tenant>>>,
statistics: Arc<RwLock<HashMap<TenantId, TenantStatistics>>>,
namespace_manager: Arc<NamespaceManager>,
quota_enforcer: Arc<QuotaEnforcer>,
rate_limiter: Arc<RateLimiter>,
access_control: Arc<AccessControl>,
billing_engine: Arc<BillingEngine>,
}
impl MultiTenantManager {
pub fn new(config: TenantManagerConfig) -> Self {
let isolation_strategy = IsolationStrategy::new(config.default_isolation);
let namespace_manager = Arc::new(NamespaceManager::new(isolation_strategy));
let quota_enforcer = Arc::new(QuotaEnforcer::new());
let rate_limiter = Arc::new(RateLimiter::new());
let access_control = Arc::new(AccessControl::new());
let billing_engine = Arc::new(BillingEngine::new(config.billing_period));
Self {
config,
tenants: Arc::new(RwLock::new(HashMap::new())),
statistics: Arc::new(RwLock::new(HashMap::new())),
namespace_manager,
quota_enforcer,
rate_limiter,
access_control,
billing_engine,
}
}
pub fn with_defaults() -> Self {
Self::new(TenantManagerConfig::default_config())
}
pub fn create_tenant(
&self,
tenant_id: impl Into<String>,
config: TenantConfig,
) -> MultiTenancyResult<()> {
let tenant_id = tenant_id.into();
let tenant = Tenant::new(tenant_id.clone(), config.metadata);
self.namespace_manager.register_tenant(&tenant_id)?;
self.quota_enforcer.set_limits(config.quotas)?;
if let Some(rate) = config.rate_limit {
self.rate_limiter.set_rate(&tenant_id, rate)?;
}
self.billing_engine
.set_pricing(&tenant_id, config.pricing)?;
if self.config.enable_access_control {
self.access_control.create_default_policy(&tenant_id)?;
}
self.tenants
.write()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?
.insert(tenant_id.clone(), tenant);
self.statistics
.write()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?
.insert(tenant_id.clone(), TenantStatistics::new(tenant_id));
Ok(())
}
pub fn get_tenant(&self, tenant_id: &str) -> MultiTenancyResult<Tenant> {
self.tenants
.read()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?
.get(tenant_id)
.cloned()
.ok_or_else(|| MultiTenancyError::TenantNotFound {
tenant_id: tenant_id.to_string(),
})
}
pub fn update_tenant_status(
&self,
tenant_id: &str,
status: TenantStatus,
) -> MultiTenancyResult<()> {
let mut tenants = self
.tenants
.write()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?;
let tenant =
tenants
.get_mut(tenant_id)
.ok_or_else(|| MultiTenancyError::TenantNotFound {
tenant_id: tenant_id.to_string(),
})?;
tenant.set_status(status);
Ok(())
}
pub fn delete_tenant(&self, tenant_id: &str) -> MultiTenancyResult<()> {
self.tenants
.write()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?
.remove(tenant_id);
self.namespace_manager.unregister_tenant(tenant_id)?;
self.statistics
.write()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?
.remove(tenant_id);
Ok(())
}
pub fn check_operation(
&self,
context: &TenantContext,
_operation: TenantOperation,
resource_delta: Option<(crate::multi_tenancy::quota::ResourceType, u64)>,
) -> MultiTenancyResult<()> {
let tenant_id = &context.tenant_id;
let tenant = self.get_tenant(tenant_id)?;
if !tenant.is_operational() {
if tenant.status == TenantStatus::Suspended {
return Err(MultiTenancyError::TenantSuspended {
tenant_id: tenant_id.clone(),
});
}
return Err(MultiTenancyError::InternalError {
message: format!("Tenant {} is not operational", tenant_id),
});
}
if !self.rate_limiter.allow_request(tenant_id)? {
return Err(MultiTenancyError::RateLimitExceeded {
tenant_id: tenant_id.clone(),
});
}
if let Some((resource_type, amount)) = resource_delta {
if self.config.strict_quotas
&& !self
.quota_enforcer
.check_quota(tenant_id, resource_type, amount)?
{
if self.config.auto_suspend_on_quota_exceeded {
self.update_tenant_status(tenant_id, TenantStatus::Suspended)?;
}
return Err(MultiTenancyError::QuotaExceeded {
tenant_id: tenant_id.clone(),
resource: resource_type.name(),
});
}
}
Ok(())
}
pub fn execute_operation<F, R>(
&self,
context: &TenantContext,
operation: TenantOperation,
func: F,
) -> MultiTenancyResult<R>
where
F: FnOnce() -> MultiTenancyResult<R>,
{
self.check_operation(context, operation, None)?;
let start = chrono::Utc::now();
let result = func()?;
let latency_ms = (chrono::Utc::now() - start).num_milliseconds() as f64;
self.record_operation_completed(context, operation, latency_ms)?;
Ok(result)
}
fn record_operation_completed(
&self,
context: &TenantContext,
operation: TenantOperation,
latency_ms: f64,
) -> MultiTenancyResult<()> {
let tenant_id = &context.tenant_id;
let mut stats = self
.statistics
.write()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?;
stats
.entry(tenant_id.clone())
.or_insert_with(|| TenantStatistics::new(tenant_id))
.record_operation(operation);
if operation == TenantOperation::VectorSearch {
stats
.get_mut(tenant_id)
.expect("tenant stats entry was just inserted via or_insert_with")
.record_query(latency_ms);
}
if self.config.enable_billing {
self.billing_engine.record_usage(tenant_id, operation, 1)?;
}
Ok(())
}
pub fn get_statistics(&self, tenant_id: &str) -> MultiTenancyResult<TenantStatistics> {
self.statistics
.read()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?
.get(tenant_id)
.cloned()
.ok_or_else(|| MultiTenancyError::TenantNotFound {
tenant_id: tenant_id.to_string(),
})
}
pub fn get_billing_metrics(&self, tenant_id: &str) -> MultiTenancyResult<BillingMetrics> {
self.billing_engine.get_metrics(tenant_id)
}
pub fn list_tenants(&self) -> MultiTenancyResult<Vec<Tenant>> {
Ok(self
.tenants
.read()
.map_err(|e| MultiTenancyError::InternalError {
message: format!("Lock error: {}", e),
})?
.values()
.cloned()
.collect())
}
pub fn namespace_manager(&self) -> &NamespaceManager {
&self.namespace_manager
}
pub fn quota_enforcer(&self) -> &QuotaEnforcer {
&self.quota_enforcer
}
pub fn access_control(&self) -> &AccessControl {
&self.access_control
}
pub fn billing_engine(&self) -> &BillingEngine {
&self.billing_engine
}
}
#[cfg(test)]
mod tests {
type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
use super::*;
#[test]
fn test_tenant_config_tiers() {
let free = TenantConfig::free_tier("t1", "Free Tenant");
assert_eq!(free.metadata.tier, "free");
assert!(free.rate_limit.is_some());
let pro = TenantConfig::pro_tier("t2", "Pro Tenant");
assert_eq!(pro.metadata.tier, "pro");
let enterprise = TenantConfig::enterprise_tier("t3", "Enterprise");
assert_eq!(enterprise.metadata.tier, "enterprise");
assert!(enterprise.rate_limit.is_none()); }
#[test]
fn test_manager_creation() -> Result<()> {
let manager = MultiTenantManager::with_defaults();
assert_eq!(manager.list_tenants().expect("test value").len(), 0);
Ok(())
}
#[test]
fn test_create_and_get_tenant() -> Result<()> {
let manager = MultiTenantManager::with_defaults();
let config = TenantConfig::free_tier("tenant1", "Test Tenant");
manager.create_tenant("tenant1", config)?;
let tenant = manager.get_tenant("tenant1")?;
assert_eq!(tenant.id, "tenant1");
assert_eq!(tenant.metadata.name, "Test Tenant");
assert_eq!(tenant.status, TenantStatus::Active);
Ok(())
}
#[test]
fn test_tenant_operations() -> Result<()> {
let manager = MultiTenantManager::with_defaults();
let config = TenantConfig::free_tier("tenant1", "Test");
manager.create_tenant("tenant1", config)?;
let context = TenantContext::new("tenant1");
manager.check_operation(&context, TenantOperation::VectorSearch, None)?;
let result: MultiTenancyResult<i32> =
manager.execute_operation(&context, TenantOperation::VectorSearch, || Ok(42));
let __val = result?;
assert_eq!(__val, 42);
let stats = manager.get_statistics("tenant1")?;
assert_eq!(stats.total_queries, 1);
Ok(())
}
#[test]
fn test_tenant_status_changes() -> Result<()> {
let manager = MultiTenantManager::with_defaults();
let config = TenantConfig::free_tier("tenant1", "Test");
manager.create_tenant("tenant1", config)?;
manager.update_tenant_status("tenant1", TenantStatus::Suspended)?;
let tenant = manager.get_tenant("tenant1")?;
assert_eq!(tenant.status, TenantStatus::Suspended);
let context = TenantContext::new("tenant1");
assert!(manager
.check_operation(&context, TenantOperation::VectorSearch, None)
.is_err());
Ok(())
}
#[test]
fn test_delete_tenant() -> Result<()> {
let manager = MultiTenantManager::with_defaults();
let config = TenantConfig::free_tier("tenant1", "Test");
manager.create_tenant("tenant1", config)?;
assert!(manager.get_tenant("tenant1").is_ok());
manager.delete_tenant("tenant1")?;
assert!(manager.get_tenant("tenant1").is_err());
Ok(())
}
#[test]
fn test_list_tenants() -> Result<()> {
let manager = MultiTenantManager::with_defaults();
manager.create_tenant("tenant1", TenantConfig::free_tier("tenant1", "T1"))?;
manager.create_tenant("tenant2", TenantConfig::pro_tier("tenant2", "T2"))?;
manager.create_tenant("tenant3", TenantConfig::enterprise_tier("tenant3", "T3"))?;
let tenants = manager.list_tenants()?;
assert_eq!(tenants.len(), 3);
Ok(())
}
#[test]
fn test_billing_integration() -> Result<()> {
let manager = MultiTenantManager::with_defaults();
let config = TenantConfig::free_tier("tenant1", "Test");
manager.create_tenant("tenant1", config)?;
let context = TenantContext::new("tenant1");
for _ in 0..10 {
let _ = manager.execute_operation(&context, TenantOperation::VectorSearch, || Ok(()));
}
let metrics = manager.get_billing_metrics("tenant1")?;
assert_eq!(metrics.total_requests, 10);
assert!(metrics.total_cost > 0.0);
Ok(())
}
#[test]
fn test_manager_config() {
let config = TenantManagerConfig::default_config();
assert!(config.strict_quotas);
assert!(config.enable_access_control);
assert!(config.enable_billing);
let prod_config = TenantManagerConfig::production();
assert_eq!(prod_config.default_isolation, IsolationLevel::SeparateIndex);
assert!(prod_config.auto_suspend_on_quota_exceeded);
}
}