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