1use crate::multi_tenancy::{
4 access_control::AccessControl,
5 billing::{BillingEngine, BillingMetrics, BillingPeriod, PricingModel},
6 isolation::{IsolationLevel, IsolationStrategy, NamespaceManager},
7 quota::{QuotaEnforcer, QuotaLimits, RateLimiter},
8 sla::SlaClass,
9 tenant::{Tenant, TenantId, TenantMetadata, TenantStatus},
10 types::{
11 MultiTenancyError, MultiTenancyResult, TenantContext, TenantOperation, TenantStatistics,
12 },
13};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16use std::sync::{Arc, RwLock};
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct TenantConfig {
21 pub metadata: TenantMetadata,
23
24 pub isolation: IsolationStrategy,
26
27 pub quotas: QuotaLimits,
29
30 pub pricing: PricingModel,
32
33 pub rate_limit: Option<f64>,
35
36 pub sla_class: SlaClass,
40}
41
42impl TenantConfig {
43 pub fn free_tier(tenant_id: impl Into<String>, name: impl Into<String>) -> Self {
45 let tenant_id = tenant_id.into();
46 Self {
47 metadata: TenantMetadata::new(name, "free"),
48 isolation: IsolationStrategy::free_tier(),
49 quotas: QuotaLimits::free_tier(&tenant_id),
50 pricing: PricingModel::PerRequest {
51 cost_per_request: 0.001,
52 },
53 rate_limit: Some(10.0),
54 sla_class: SlaClass::Bronze,
55 }
56 }
57
58 pub fn pro_tier(tenant_id: impl Into<String>, name: impl Into<String>) -> Self {
60 let tenant_id = tenant_id.into();
61 Self {
62 metadata: TenantMetadata::new(name, "pro"),
63 isolation: IsolationStrategy::pro_tier(),
64 quotas: QuotaLimits::pro_tier(&tenant_id),
65 pricing: PricingModel::PerComputeUnit {
66 cost_per_unit: 0.01,
67 },
68 rate_limit: Some(100.0),
69 sla_class: SlaClass::Gold,
70 }
71 }
72
73 pub fn enterprise_tier(tenant_id: impl Into<String>, name: impl Into<String>) -> Self {
75 let tenant_id = tenant_id.into();
76 Self {
77 metadata: TenantMetadata::new(name, "enterprise"),
78 isolation: IsolationStrategy::enterprise_tier(),
79 quotas: QuotaLimits::enterprise_tier(&tenant_id),
80 pricing: PricingModel::Subscription {
81 monthly_fee: 1000.0,
82 included_requests: 1_000_000,
83 overage_cost: 0.005,
84 },
85 rate_limit: None, sla_class: SlaClass::Platinum,
87 }
88 }
89}
90
91#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct TenantManagerConfig {
94 pub default_isolation: IsolationLevel,
96
97 pub billing_period: BillingPeriod,
99
100 pub strict_quotas: bool,
102
103 pub enable_access_control: bool,
105
106 pub enable_billing: bool,
108
109 pub auto_suspend_on_quota_exceeded: bool,
111}
112
113impl TenantManagerConfig {
114 pub fn default_config() -> Self {
116 Self {
117 default_isolation: IsolationLevel::Namespace,
118 billing_period: BillingPeriod::Monthly,
119 strict_quotas: true,
120 enable_access_control: true,
121 enable_billing: true,
122 auto_suspend_on_quota_exceeded: false,
123 }
124 }
125
126 pub fn production() -> Self {
128 Self {
129 default_isolation: IsolationLevel::SeparateIndex,
130 billing_period: BillingPeriod::Monthly,
131 strict_quotas: true,
132 enable_access_control: true,
133 enable_billing: true,
134 auto_suspend_on_quota_exceeded: true,
135 }
136 }
137}
138
139pub struct MultiTenantManager {
141 config: TenantManagerConfig,
143
144 tenants: Arc<RwLock<HashMap<TenantId, Tenant>>>,
146
147 statistics: Arc<RwLock<HashMap<TenantId, TenantStatistics>>>,
149
150 namespace_manager: Arc<NamespaceManager>,
152
153 quota_enforcer: Arc<QuotaEnforcer>,
155
156 rate_limiter: Arc<RateLimiter>,
158
159 access_control: Arc<AccessControl>,
161
162 billing_engine: Arc<BillingEngine>,
164}
165
166impl MultiTenantManager {
167 pub fn new(config: TenantManagerConfig) -> Self {
169 let isolation_strategy = IsolationStrategy::new(config.default_isolation);
170 let namespace_manager = Arc::new(NamespaceManager::new(isolation_strategy));
171 let quota_enforcer = Arc::new(QuotaEnforcer::new());
172 let rate_limiter = Arc::new(RateLimiter::new());
173 let access_control = Arc::new(AccessControl::new());
174 let billing_engine = Arc::new(BillingEngine::new(config.billing_period));
175
176 Self {
177 config,
178 tenants: Arc::new(RwLock::new(HashMap::new())),
179 statistics: Arc::new(RwLock::new(HashMap::new())),
180 namespace_manager,
181 quota_enforcer,
182 rate_limiter,
183 access_control,
184 billing_engine,
185 }
186 }
187
188 pub fn with_defaults() -> Self {
190 Self::new(TenantManagerConfig::default_config())
191 }
192
193 pub fn create_tenant(
195 &self,
196 tenant_id: impl Into<String>,
197 config: TenantConfig,
198 ) -> MultiTenancyResult<()> {
199 let tenant_id = tenant_id.into();
200 let tenant = Tenant::new(tenant_id.clone(), config.metadata);
201
202 self.namespace_manager.register_tenant(&tenant_id)?;
204
205 self.quota_enforcer.set_limits(config.quotas)?;
207
208 if let Some(rate) = config.rate_limit {
210 self.rate_limiter.set_rate(&tenant_id, rate)?;
211 }
212
213 self.billing_engine
215 .set_pricing(&tenant_id, config.pricing)?;
216
217 if self.config.enable_access_control {
219 self.access_control.create_default_policy(&tenant_id)?;
220 }
221
222 self.tenants
224 .write()
225 .map_err(|e| MultiTenancyError::InternalError {
226 message: format!("Lock error: {}", e),
227 })?
228 .insert(tenant_id.clone(), tenant);
229
230 self.statistics
232 .write()
233 .map_err(|e| MultiTenancyError::InternalError {
234 message: format!("Lock error: {}", e),
235 })?
236 .insert(tenant_id.clone(), TenantStatistics::new(tenant_id));
237
238 Ok(())
239 }
240
241 pub fn get_tenant(&self, tenant_id: &str) -> MultiTenancyResult<Tenant> {
243 self.tenants
244 .read()
245 .map_err(|e| MultiTenancyError::InternalError {
246 message: format!("Lock error: {}", e),
247 })?
248 .get(tenant_id)
249 .cloned()
250 .ok_or_else(|| MultiTenancyError::TenantNotFound {
251 tenant_id: tenant_id.to_string(),
252 })
253 }
254
255 pub fn update_tenant_status(
257 &self,
258 tenant_id: &str,
259 status: TenantStatus,
260 ) -> MultiTenancyResult<()> {
261 let mut tenants = self
262 .tenants
263 .write()
264 .map_err(|e| MultiTenancyError::InternalError {
265 message: format!("Lock error: {}", e),
266 })?;
267
268 let tenant =
269 tenants
270 .get_mut(tenant_id)
271 .ok_or_else(|| MultiTenancyError::TenantNotFound {
272 tenant_id: tenant_id.to_string(),
273 })?;
274
275 tenant.set_status(status);
276 Ok(())
277 }
278
279 pub fn delete_tenant(&self, tenant_id: &str) -> MultiTenancyResult<()> {
281 self.tenants
283 .write()
284 .map_err(|e| MultiTenancyError::InternalError {
285 message: format!("Lock error: {}", e),
286 })?
287 .remove(tenant_id);
288
289 self.namespace_manager.unregister_tenant(tenant_id)?;
291
292 self.statistics
294 .write()
295 .map_err(|e| MultiTenancyError::InternalError {
296 message: format!("Lock error: {}", e),
297 })?
298 .remove(tenant_id);
299
300 Ok(())
301 }
302
303 pub fn check_operation(
305 &self,
306 context: &TenantContext,
307 _operation: TenantOperation,
308 resource_delta: Option<(crate::multi_tenancy::quota::ResourceType, u64)>,
309 ) -> MultiTenancyResult<()> {
310 let tenant_id = &context.tenant_id;
311
312 let tenant = self.get_tenant(tenant_id)?;
314 if !tenant.is_operational() {
315 if tenant.status == TenantStatus::Suspended {
316 return Err(MultiTenancyError::TenantSuspended {
317 tenant_id: tenant_id.clone(),
318 });
319 }
320 return Err(MultiTenancyError::InternalError {
321 message: format!("Tenant {} is not operational", tenant_id),
322 });
323 }
324
325 if !self.rate_limiter.allow_request(tenant_id)? {
327 return Err(MultiTenancyError::RateLimitExceeded {
328 tenant_id: tenant_id.clone(),
329 });
330 }
331
332 if let Some((resource_type, amount)) = resource_delta {
334 if self.config.strict_quotas
335 && !self
336 .quota_enforcer
337 .check_quota(tenant_id, resource_type, amount)?
338 {
339 if self.config.auto_suspend_on_quota_exceeded {
340 self.update_tenant_status(tenant_id, TenantStatus::Suspended)?;
341 }
342 return Err(MultiTenancyError::QuotaExceeded {
343 tenant_id: tenant_id.clone(),
344 resource: resource_type.name(),
345 });
346 }
347 }
348
349 Ok(())
350 }
351
352 pub fn execute_operation<F, R>(
354 &self,
355 context: &TenantContext,
356 operation: TenantOperation,
357 func: F,
358 ) -> MultiTenancyResult<R>
359 where
360 F: FnOnce() -> MultiTenancyResult<R>,
361 {
362 self.check_operation(context, operation, None)?;
364
365 let start = chrono::Utc::now();
367 let result = func()?;
368 let latency_ms = (chrono::Utc::now() - start).num_milliseconds() as f64;
369
370 self.record_operation_completed(context, operation, latency_ms)?;
372
373 Ok(result)
374 }
375
376 fn record_operation_completed(
378 &self,
379 context: &TenantContext,
380 operation: TenantOperation,
381 latency_ms: f64,
382 ) -> MultiTenancyResult<()> {
383 let tenant_id = &context.tenant_id;
384
385 let mut stats = self
387 .statistics
388 .write()
389 .map_err(|e| MultiTenancyError::InternalError {
390 message: format!("Lock error: {}", e),
391 })?;
392
393 stats
394 .entry(tenant_id.clone())
395 .or_insert_with(|| TenantStatistics::new(tenant_id))
396 .record_operation(operation);
397
398 if operation == TenantOperation::VectorSearch {
399 stats
400 .get_mut(tenant_id)
401 .expect("tenant stats entry was just inserted via or_insert_with")
402 .record_query(latency_ms);
403 }
404
405 if self.config.enable_billing {
407 self.billing_engine.record_usage(tenant_id, operation, 1)?;
408 }
409
410 Ok(())
411 }
412
413 pub fn get_statistics(&self, tenant_id: &str) -> MultiTenancyResult<TenantStatistics> {
415 self.statistics
416 .read()
417 .map_err(|e| MultiTenancyError::InternalError {
418 message: format!("Lock error: {}", e),
419 })?
420 .get(tenant_id)
421 .cloned()
422 .ok_or_else(|| MultiTenancyError::TenantNotFound {
423 tenant_id: tenant_id.to_string(),
424 })
425 }
426
427 pub fn get_billing_metrics(&self, tenant_id: &str) -> MultiTenancyResult<BillingMetrics> {
429 self.billing_engine.get_metrics(tenant_id)
430 }
431
432 pub fn list_tenants(&self) -> MultiTenancyResult<Vec<Tenant>> {
434 Ok(self
435 .tenants
436 .read()
437 .map_err(|e| MultiTenancyError::InternalError {
438 message: format!("Lock error: {}", e),
439 })?
440 .values()
441 .cloned()
442 .collect())
443 }
444
445 pub fn namespace_manager(&self) -> &NamespaceManager {
447 &self.namespace_manager
448 }
449
450 pub fn quota_enforcer(&self) -> &QuotaEnforcer {
452 &self.quota_enforcer
453 }
454
455 pub fn access_control(&self) -> &AccessControl {
457 &self.access_control
458 }
459
460 pub fn billing_engine(&self) -> &BillingEngine {
462 &self.billing_engine
463 }
464}
465
466#[cfg(test)]
467mod tests {
468 type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
469 use super::*;
470
471 #[test]
472 fn test_tenant_config_tiers() {
473 let free = TenantConfig::free_tier("t1", "Free Tenant");
474 assert_eq!(free.metadata.tier, "free");
475 assert!(free.rate_limit.is_some());
476
477 let pro = TenantConfig::pro_tier("t2", "Pro Tenant");
478 assert_eq!(pro.metadata.tier, "pro");
479
480 let enterprise = TenantConfig::enterprise_tier("t3", "Enterprise");
481 assert_eq!(enterprise.metadata.tier, "enterprise");
482 assert!(enterprise.rate_limit.is_none()); }
484
485 #[test]
486 fn test_manager_creation() -> Result<()> {
487 let manager = MultiTenantManager::with_defaults();
488 assert_eq!(manager.list_tenants().expect("test value").len(), 0);
489 Ok(())
490 }
491
492 #[test]
493 fn test_create_and_get_tenant() -> Result<()> {
494 let manager = MultiTenantManager::with_defaults();
495 let config = TenantConfig::free_tier("tenant1", "Test Tenant");
496
497 manager.create_tenant("tenant1", config)?;
498
499 let tenant = manager.get_tenant("tenant1")?;
500 assert_eq!(tenant.id, "tenant1");
501 assert_eq!(tenant.metadata.name, "Test Tenant");
502 assert_eq!(tenant.status, TenantStatus::Active);
503 Ok(())
504 }
505
506 #[test]
507 fn test_tenant_operations() -> Result<()> {
508 let manager = MultiTenantManager::with_defaults();
509 let config = TenantConfig::free_tier("tenant1", "Test");
510 manager.create_tenant("tenant1", config)?;
511
512 let context = TenantContext::new("tenant1");
513
514 manager.check_operation(&context, TenantOperation::VectorSearch, None)?;
516
517 let result: MultiTenancyResult<i32> =
519 manager.execute_operation(&context, TenantOperation::VectorSearch, || Ok(42));
520 let __val = result?;
521 assert_eq!(__val, 42);
522
523 let stats = manager.get_statistics("tenant1")?;
525 assert_eq!(stats.total_queries, 1);
526 Ok(())
527 }
528
529 #[test]
530 fn test_tenant_status_changes() -> Result<()> {
531 let manager = MultiTenantManager::with_defaults();
532 let config = TenantConfig::free_tier("tenant1", "Test");
533 manager.create_tenant("tenant1", config)?;
534
535 manager.update_tenant_status("tenant1", TenantStatus::Suspended)?;
537
538 let tenant = manager.get_tenant("tenant1")?;
539 assert_eq!(tenant.status, TenantStatus::Suspended);
540
541 let context = TenantContext::new("tenant1");
543 assert!(manager
544 .check_operation(&context, TenantOperation::VectorSearch, None)
545 .is_err());
546 Ok(())
547 }
548
549 #[test]
550 fn test_delete_tenant() -> Result<()> {
551 let manager = MultiTenantManager::with_defaults();
552 let config = TenantConfig::free_tier("tenant1", "Test");
553 manager.create_tenant("tenant1", config)?;
554
555 assert!(manager.get_tenant("tenant1").is_ok());
556
557 manager.delete_tenant("tenant1")?;
558
559 assert!(manager.get_tenant("tenant1").is_err());
560 Ok(())
561 }
562
563 #[test]
564 fn test_list_tenants() -> Result<()> {
565 let manager = MultiTenantManager::with_defaults();
566
567 manager.create_tenant("tenant1", TenantConfig::free_tier("tenant1", "T1"))?;
568 manager.create_tenant("tenant2", TenantConfig::pro_tier("tenant2", "T2"))?;
569 manager.create_tenant("tenant3", TenantConfig::enterprise_tier("tenant3", "T3"))?;
570
571 let tenants = manager.list_tenants()?;
572 assert_eq!(tenants.len(), 3);
573 Ok(())
574 }
575
576 #[test]
577 fn test_billing_integration() -> Result<()> {
578 let manager = MultiTenantManager::with_defaults();
579 let config = TenantConfig::free_tier("tenant1", "Test");
580 manager.create_tenant("tenant1", config)?;
581
582 let context = TenantContext::new("tenant1");
583
584 for _ in 0..10 {
586 let _ = manager.execute_operation(&context, TenantOperation::VectorSearch, || Ok(()));
587 }
588
589 let metrics = manager.get_billing_metrics("tenant1")?;
591 assert_eq!(metrics.total_requests, 10);
592 assert!(metrics.total_cost > 0.0);
593 Ok(())
594 }
595
596 #[test]
597 fn test_manager_config() {
598 let config = TenantManagerConfig::default_config();
599 assert!(config.strict_quotas);
600 assert!(config.enable_access_control);
601 assert!(config.enable_billing);
602
603 let prod_config = TenantManagerConfig::production();
604 assert_eq!(prod_config.default_isolation, IsolationLevel::SeparateIndex);
605 assert!(prod_config.auto_suspend_on_quota_exceeded);
606 }
607}