use crate::dynamic::ServiceRegistry;
use crate::reflection::{
cache::DescriptorCache,
config::ProxyConfig,
connection_pool::ConnectionPool,
smart_mock_generator::{SmartMockConfig, SmartMockGenerator},
};
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc, Mutex,
};
use std::time::Duration;
use tonic::Status;
use tracing::debug;
pub struct ConnectionGuard {
counter: Arc<AtomicUsize>,
}
impl Drop for ConnectionGuard {
fn drop(&mut self) {
self.counter.fetch_sub(1, Ordering::Relaxed);
}
}
pub struct MockReflectionProxy {
pub(crate) cache: DescriptorCache,
pub(crate) config: ProxyConfig,
pub(crate) timeout_duration: Duration,
#[allow(dead_code)]
pub(crate) connection_pool: ConnectionPool,
pub(crate) service_registry: Arc<ServiceRegistry>,
pub(crate) smart_generator: Arc<Mutex<SmartMockGenerator>>,
pub(crate) active_connections: Arc<AtomicUsize>,
pub(crate) total_requests: Arc<AtomicUsize>,
}
impl MockReflectionProxy {
pub async fn new(
config: ProxyConfig,
service_registry: Arc<ServiceRegistry>,
) -> Result<Self, Status> {
debug!(
"Creating mock reflection proxy with {} services",
service_registry.service_names().len()
);
let cache = DescriptorCache::new();
cache.populate_from_pool(Some(service_registry.descriptor_pool())).await;
let connection_pool = ConnectionPool::new();
let timeout_duration = Duration::from_secs(config.request_timeout_seconds);
let smart_generator = Arc::new(Mutex::new(SmartMockGenerator::new(SmartMockConfig {
field_name_inference: true,
use_faker: true,
field_overrides: std::collections::HashMap::new(),
service_profiles: std::collections::HashMap::new(),
max_depth: 3,
seed: config.mock_seed,
deterministic: false,
})));
Ok(Self {
cache,
config,
timeout_duration,
connection_pool,
service_registry,
smart_generator,
active_connections: Arc::new(AtomicUsize::new(0)),
total_requests: Arc::new(AtomicUsize::new(0)),
})
}
pub fn config(&self) -> &ProxyConfig {
&self.config
}
pub fn cache(&self) -> &DescriptorCache {
&self.cache
}
pub fn service_registry(&self) -> &Arc<ServiceRegistry> {
&self.service_registry
}
pub fn service_names(&self) -> Vec<String> {
self.service_registry.service_names()
}
pub fn smart_generator(&self) -> &Arc<Mutex<SmartMockGenerator>> {
&self.smart_generator
}
pub fn should_mock_service_method(&self, service_name: &str, method_name: &str) -> bool {
if let Some(service) = self.service_registry.get(service_name) {
service.methods().iter().any(|m| m.name == method_name)
} else {
false
}
}
pub fn timeout_duration(&self) -> Duration {
self.timeout_duration
}
pub fn update_config(&mut self, config: ProxyConfig) {
self.config = config;
self.timeout_duration = Duration::from_secs(self.config.request_timeout_seconds);
}
pub fn track_connection(&self) -> ConnectionGuard {
self.active_connections.fetch_add(1, Ordering::Relaxed);
ConnectionGuard {
counter: self.active_connections.clone(),
}
}
pub async fn get_stats(&self) -> ProxyStats {
ProxyStats {
cached_services: self.cache.service_count().await,
cached_methods: self.cache.method_count().await,
registered_services: self.service_registry.service_names().len(),
total_requests: self.total_requests.load(Ordering::Relaxed) as u64,
active_connections: self.active_connections.load(Ordering::Relaxed),
}
}
}
#[derive(Debug, Clone)]
pub struct ProxyStats {
pub cached_services: usize,
pub cached_methods: usize,
pub registered_services: usize,
pub total_requests: u64,
pub active_connections: usize,
}
impl Clone for MockReflectionProxy {
fn clone(&self) -> Self {
Self {
cache: self.cache.clone(),
config: self.config.clone(),
timeout_duration: self.timeout_duration,
connection_pool: self.connection_pool.clone(),
service_registry: self.service_registry.clone(),
smart_generator: self.smart_generator.clone(),
active_connections: self.active_connections.clone(),
total_requests: self.total_requests.clone(),
}
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_module_compiles() {
}
}