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
391 .get_mut(tenant_id)
392 .expect("tenant stats entry was just inserted via or_insert_with")
393 .record_query(latency_ms);
394 }
395
396 if self.config.enable_billing {
398 self.billing_engine.record_usage(tenant_id, operation, 1)?;
399 }
400
401 Ok(())
402 }
403
404 pub fn get_statistics(&self, tenant_id: &str) -> MultiTenancyResult<TenantStatistics> {
406 self.statistics
407 .read()
408 .map_err(|e| MultiTenancyError::InternalError {
409 message: format!("Lock error: {}", e),
410 })?
411 .get(tenant_id)
412 .cloned()
413 .ok_or_else(|| MultiTenancyError::TenantNotFound {
414 tenant_id: tenant_id.to_string(),
415 })
416 }
417
418 pub fn get_billing_metrics(&self, tenant_id: &str) -> MultiTenancyResult<BillingMetrics> {
420 self.billing_engine.get_metrics(tenant_id)
421 }
422
423 pub fn list_tenants(&self) -> MultiTenancyResult<Vec<Tenant>> {
425 Ok(self
426 .tenants
427 .read()
428 .map_err(|e| MultiTenancyError::InternalError {
429 message: format!("Lock error: {}", e),
430 })?
431 .values()
432 .cloned()
433 .collect())
434 }
435
436 pub fn namespace_manager(&self) -> &NamespaceManager {
438 &self.namespace_manager
439 }
440
441 pub fn quota_enforcer(&self) -> &QuotaEnforcer {
443 &self.quota_enforcer
444 }
445
446 pub fn access_control(&self) -> &AccessControl {
448 &self.access_control
449 }
450
451 pub fn billing_engine(&self) -> &BillingEngine {
453 &self.billing_engine
454 }
455}
456
457#[cfg(test)]
458mod tests {
459 type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
460 use super::*;
461
462 #[test]
463 fn test_tenant_config_tiers() {
464 let free = TenantConfig::free_tier("t1", "Free Tenant");
465 assert_eq!(free.metadata.tier, "free");
466 assert!(free.rate_limit.is_some());
467
468 let pro = TenantConfig::pro_tier("t2", "Pro Tenant");
469 assert_eq!(pro.metadata.tier, "pro");
470
471 let enterprise = TenantConfig::enterprise_tier("t3", "Enterprise");
472 assert_eq!(enterprise.metadata.tier, "enterprise");
473 assert!(enterprise.rate_limit.is_none()); }
475
476 #[test]
477 fn test_manager_creation() -> Result<()> {
478 let manager = MultiTenantManager::with_defaults();
479 assert_eq!(manager.list_tenants().expect("test value").len(), 0);
480 Ok(())
481 }
482
483 #[test]
484 fn test_create_and_get_tenant() -> Result<()> {
485 let manager = MultiTenantManager::with_defaults();
486 let config = TenantConfig::free_tier("tenant1", "Test Tenant");
487
488 manager.create_tenant("tenant1", config)?;
489
490 let tenant = manager.get_tenant("tenant1")?;
491 assert_eq!(tenant.id, "tenant1");
492 assert_eq!(tenant.metadata.name, "Test Tenant");
493 assert_eq!(tenant.status, TenantStatus::Active);
494 Ok(())
495 }
496
497 #[test]
498 fn test_tenant_operations() -> Result<()> {
499 let manager = MultiTenantManager::with_defaults();
500 let config = TenantConfig::free_tier("tenant1", "Test");
501 manager.create_tenant("tenant1", config)?;
502
503 let context = TenantContext::new("tenant1");
504
505 manager.check_operation(&context, TenantOperation::VectorSearch, None)?;
507
508 let result: MultiTenancyResult<i32> =
510 manager.execute_operation(&context, TenantOperation::VectorSearch, || Ok(42));
511 let __val = result?;
512 assert_eq!(__val, 42);
513
514 let stats = manager.get_statistics("tenant1")?;
516 assert_eq!(stats.total_queries, 1);
517 Ok(())
518 }
519
520 #[test]
521 fn test_tenant_status_changes() -> Result<()> {
522 let manager = MultiTenantManager::with_defaults();
523 let config = TenantConfig::free_tier("tenant1", "Test");
524 manager.create_tenant("tenant1", config)?;
525
526 manager.update_tenant_status("tenant1", TenantStatus::Suspended)?;
528
529 let tenant = manager.get_tenant("tenant1")?;
530 assert_eq!(tenant.status, TenantStatus::Suspended);
531
532 let context = TenantContext::new("tenant1");
534 assert!(manager
535 .check_operation(&context, TenantOperation::VectorSearch, None)
536 .is_err());
537 Ok(())
538 }
539
540 #[test]
541 fn test_delete_tenant() -> Result<()> {
542 let manager = MultiTenantManager::with_defaults();
543 let config = TenantConfig::free_tier("tenant1", "Test");
544 manager.create_tenant("tenant1", config)?;
545
546 assert!(manager.get_tenant("tenant1").is_ok());
547
548 manager.delete_tenant("tenant1")?;
549
550 assert!(manager.get_tenant("tenant1").is_err());
551 Ok(())
552 }
553
554 #[test]
555 fn test_list_tenants() -> Result<()> {
556 let manager = MultiTenantManager::with_defaults();
557
558 manager.create_tenant("tenant1", TenantConfig::free_tier("tenant1", "T1"))?;
559 manager.create_tenant("tenant2", TenantConfig::pro_tier("tenant2", "T2"))?;
560 manager.create_tenant("tenant3", TenantConfig::enterprise_tier("tenant3", "T3"))?;
561
562 let tenants = manager.list_tenants()?;
563 assert_eq!(tenants.len(), 3);
564 Ok(())
565 }
566
567 #[test]
568 fn test_billing_integration() -> Result<()> {
569 let manager = MultiTenantManager::with_defaults();
570 let config = TenantConfig::free_tier("tenant1", "Test");
571 manager.create_tenant("tenant1", config)?;
572
573 let context = TenantContext::new("tenant1");
574
575 for _ in 0..10 {
577 let _ = manager.execute_operation(&context, TenantOperation::VectorSearch, || Ok(()));
578 }
579
580 let metrics = manager.get_billing_metrics("tenant1")?;
582 assert_eq!(metrics.total_requests, 10);
583 assert!(metrics.total_cost > 0.0);
584 Ok(())
585 }
586
587 #[test]
588 fn test_manager_config() {
589 let config = TenantManagerConfig::default_config();
590 assert!(config.strict_quotas);
591 assert!(config.enable_access_control);
592 assert!(config.enable_billing);
593
594 let prod_config = TenantManagerConfig::production();
595 assert_eq!(prod_config.default_isolation, IsolationLevel::SeparateIndex);
596 assert!(prod_config.auto_suspend_on_quota_exceeded);
597 }
598}