use std::sync::Arc;
use tokio::sync::Mutex;
use tonic::transport::Channel;
use crate::common::AnyResult;
use crate::streaming::common::{
MetricsManager, PerformanceMetrics, StreamClientConfig,
};
use crate::protos::shredstream::shredstream_proxy_client::ShredstreamProxyClient;
#[derive(Clone)]
pub struct ShredStreamGrpc {
pub shredstream_client: Arc<ShredstreamProxyClient<Channel>>,
pub config: StreamClientConfig,
pub metrics: Arc<Mutex<PerformanceMetrics>>,
pub metrics_manager: MetricsManager,
}
impl ShredStreamGrpc {
pub async fn new(endpoint: String) -> AnyResult<Self> {
Self::new_with_config(endpoint, StreamClientConfig::default()).await
}
pub async fn new_with_config(endpoint: String, config: StreamClientConfig) -> AnyResult<Self> {
let shredstream_client = ShredstreamProxyClient::connect(endpoint.clone()).await?;
let metrics = Arc::new(Mutex::new(PerformanceMetrics::new()));
let config_arc = Arc::new(config.clone());
let metrics_manager = MetricsManager::new(
metrics.clone(),
config_arc,
"ShredStream".to_string()
);
Ok(Self {
shredstream_client: Arc::new(shredstream_client),
config,
metrics,
metrics_manager,
})
}
pub async fn new_high_performance(endpoint: String) -> AnyResult<Self> {
Self::new_with_config(endpoint, StreamClientConfig::high_performance()).await
}
pub async fn new_low_latency(endpoint: String) -> AnyResult<Self> {
Self::new_with_config(endpoint, StreamClientConfig::low_latency()).await
}
pub fn get_config(&self) -> &StreamClientConfig {
&self.config
}
pub fn update_config(&mut self, config: StreamClientConfig) {
self.config = config;
}
pub async fn get_metrics(&self) -> PerformanceMetrics {
self.metrics_manager.get_metrics().await
}
pub fn set_enable_metrics(&mut self, enabled: bool) {
self.config.enable_metrics = enabled;
}
pub async fn print_metrics(&self) {
self.metrics_manager.print_metrics().await;
}
pub async fn start_auto_metrics_monitoring(&self) {
self.metrics_manager.start_auto_monitoring().await;
}
}