use std::sync::Arc;
use super::config::HlsmiConfig;
use super::parser::GaudiMetricsData;
use super::process::ProcessManager;
use super::store::MetricsStore;
pub struct DataCollector {
process_manager: ProcessManager,
store: Arc<MetricsStore>,
}
impl DataCollector {
pub fn new(config: HlsmiConfig, store: Arc<MetricsStore>) -> Self {
let process_manager = ProcessManager::new(config, store.clone());
Self {
process_manager,
store,
}
}
pub fn start(&mut self) -> Result<(), Box<dyn std::error::Error>> {
self.process_manager.start()
}
pub fn stop(&mut self) {
self.process_manager.shutdown();
self.store.clear();
}
pub fn get_latest_data(&self) -> Result<GaudiMetricsData, Box<dyn std::error::Error>> {
self.store.get_latest_data()
}
#[cfg(test)]
pub(super) fn is_running(&self) -> bool {
self.process_manager.is_running()
}
#[cfg(test)]
pub(super) fn wait_for_initial_data(
&self,
timeout: std::time::Duration,
) -> Result<(), Box<dyn std::error::Error>> {
let start = std::time::Instant::now();
while start.elapsed() < timeout {
if self.get_latest_data().is_ok() {
return Ok(());
}
std::thread::sleep(std::time::Duration::from_millis(100));
}
Err("Timeout waiting for initial hl-smi data".into())
}
}
impl Drop for DataCollector {
fn drop(&mut self) {
self.stop();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_collector_creation() {
let config = HlsmiConfig::default();
let store = Arc::new(MetricsStore::new(120));
let collector = DataCollector::new(config, store);
assert!(!collector.is_running());
}
#[test]
fn test_wait_for_initial_data_timeout() {
let config = HlsmiConfig::default();
let store = Arc::new(MetricsStore::new(120));
let collector = DataCollector::new(config, store);
let result = collector.wait_for_initial_data(std::time::Duration::from_millis(100));
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("Timeout"));
}
}