streaming_telemetry/src/sensor_database/i2c_sensors/mod.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use tokio::sync::Mutex;

use crate::sensor_database::sensor_database::sensor_database::{AsyncReadSensor, Sensor};

use crate::sensor_database::sensor_configs::{
    find_sensor_file, FruInfo, PSUProperty, Sensor as SensorConfig,
};

use std::str::FromStr;
use std::thread;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio_uring::fs::File as UringFile;

const BUFFER_SIZE: usize = 128;

struct I2cSensorReader {
    _name: String,
    file_name: String,
    // This mutex is only for making this struct to have Send + Sync, as required by trait AsyncReadSensor
    // There should be no real contender for this lock in data path
    value_receiver: Arc<Mutex<mpsc::Receiver<Vec<(f64, SystemTime)>>>>,
    psu_property: PSUProperty,
    sampling_interval: Arc<AtomicU64>,
    reporting_interval: Arc<AtomicU64>,
    polling_tx: watch::Sender<bool>,
    polling_rx: watch::Receiver<bool>,
}

impl I2cSensorReader {
    fn new(
        name: &str,
        sysfs_path: &str,
        psu_property: PSUProperty,
        sampling_interval: Duration,
        reporting_interval: Duration,
    ) -> std::io::Result<Self> {
        assert!(
            reporting_interval >= sampling_interval,
            "reporting_interval must be greater than or equal to sampling_interval"
        );

        let (value_sender, value_receiver) = mpsc::channel(16);
        let sampling_interval = Arc::new(AtomicU64::new(sampling_interval.as_millis() as u64));
        let reporting_interval = Arc::new(AtomicU64::new(reporting_interval.as_millis() as u64));
        let (polling_tx, polling_rx) = watch::channel(false);

        let reader = Self {
            _name: name.to_owned(),
            file_name: sysfs_path.to_owned(),
            value_receiver: Arc::new(Mutex::new(value_receiver)),
            psu_property,
            sampling_interval: sampling_interval.clone(),
            reporting_interval: reporting_interval.clone(),
            polling_tx,
            polling_rx,
        };

        reader.start_io_uring_thread(value_sender, sampling_interval, reporting_interval);

        Ok(reader)
    }

    fn start_io_uring_thread(
        &self,
        value_sender: mpsc::Sender<Vec<(f64, SystemTime)>>,
        sampling_interval: Arc<AtomicU64>,
        reporting_interval: Arc<AtomicU64>,
    ) {
        let file_name = self.file_name.clone();
        let psu_property = self.psu_property.clone();
        let mut polling_rx = self.polling_rx.clone();

        thread::spawn(move || {
            tokio_uring::start(async move {
                let file = UringFile::open(&file_name).await.unwrap();
                let mut buf = vec![0u8; BUFFER_SIZE];
                let mut batch_values = Vec::new();
                let mut last_report_time = SystemTime::now();
                let factor = 10.0_f64.powi(psu_property.scale);

                loop {
                    let loop_start = Instant::now();

                    // Wait for the polling state to change to true
                    if !*polling_rx.borrow() {
                        let _ = polling_rx.changed().await;
                        continue;
                    }

                    let (res, new_buf) = file.read_at(buf, 0).await;
                    buf = new_buf; // Reuse the buffer

                    match res {
                        Ok(bytes_read) => {
                            if bytes_read == 0 {
                                continue;
                            }

                            let s = std::str::from_utf8(&buf[..bytes_read]).unwrap();
                            let raw_value = f64::from_str(s.trim()).unwrap();
                            let final_value = (raw_value / factor) + psu_property.offset as f64;

                            let now = SystemTime::now();
                            batch_values.push((final_value, now));

                            if now
                                .duration_since(last_report_time)
                                .unwrap_or(Duration::ZERO)
                                >= Duration::from_millis(reporting_interval.load(Ordering::Relaxed))
                            {
                                if value_sender.send(batch_values.clone()).await.is_err() {
                                    // The receiver has been dropped, exit the thread
                                    break;
                                }
                                batch_values.clear();
                                last_report_time = now;
                            }
                        }
                        Err(e) => {
                            eprintln!("Error reading file {}: {:?}", file_name, e);
                        }
                    }

                    let elapsed = loop_start.elapsed();
                    let intended_interval =
                        Duration::from_millis(sampling_interval.load(Ordering::Relaxed));

                    if elapsed < intended_interval {
                        let sleep_duration = intended_interval - elapsed;
                        tokio::time::sleep(sleep_duration).await;
                    }
                }
            });
        });
    }
}

#[async_trait::async_trait]
impl AsyncReadSensor for I2cSensorReader {
    async fn read_value(
        &self,
        sampling_interval: Duration,
        reporting_interval: Duration,
    ) -> Result<Vec<(f64, SystemTime)>, std::io::Error> {
        self.sampling_interval
            .store(sampling_interval.as_millis() as u64, Ordering::Relaxed);
        self.reporting_interval
            .store(reporting_interval.as_millis() as u64, Ordering::Relaxed);
        let mut value_receiver = self.value_receiver.lock().await;
        match value_receiver.recv().await {
            Some(values) => Ok(values),
            None => Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                "io_uring thread has terminated",
            )),
        }
    }

    fn start_polling(&self) {
        if let Err(e) = self.polling_tx.send(true) {
            eprintln!("Failed to start polling: {}", e);
        }
    }

    fn stop_polling(&self) {
        if let Err(e) = self.polling_tx.send(false) {
            eprintln!("Failed to stop polling: {}", e);
        }
    }
}

pub async fn create_i2c_sensor(
    sensor_config: &SensorConfig,
    all_sensor_frus: &HashMap<String, FruInfo>,
    sampling_interval: u64,
    reporting_interval: u64,
) -> Result<Sensor, Box<dyn std::error::Error>> {
    let sysfs_path = find_sensor_file(sensor_config)?;

    let psu_property = sensor_config.reading_type.0.clone();
    let sensor_reader = Arc::new(
        I2cSensorReader::new(
            &sensor_config.Name,
            &sysfs_path,
            psu_property,
            Duration::from_millis(sampling_interval),
            Duration::from_millis(reporting_interval),
        )
        .unwrap(),
    );

    let mut secondary_keys = Vec::new();

    // Add sensor's ReadingType as secondary key
    secondary_keys.push((
        "ReadingType".to_string(),
        format!("{:?}", sensor_config.reading_type.1),
    ));

    // Add sensor's type as secondary key
    secondary_keys.push(("Type".to_string(), "I2C".to_string()));

    // Add sensor's FruInfo if available
    if let Some(fru_info) = all_sensor_frus.get(&sensor_config.Name) {
        if let Some(location_code) = &fru_info.location_code {
            secondary_keys.push(("LocationCode".to_string(), location_code.clone()));
        }
        if let Some(service_label) = &fru_info.service_label {
            secondary_keys.push(("ServiceLabel".to_string(), service_label.clone()));
        }
    }

    let sensor = Sensor::new(
        sensor_config.Name.clone(),
        Some(secondary_keys),
        0.0, // Initial value
        sensor_reader,
        Duration::from_secs(1),
        Duration::from_millis(1),
        Duration::from_millis(reporting_interval),
    );

    Ok(sensor)
}