telemetry_visualization/server/src/telemetry/management_service.rs

use crate::telemetry::data::TelemetryData;
use crate::telemetry::data_item::TelemetryDataItem;
use crate::telemetry::definition::TelemetryDefinition;
use crate::telemetry::history::{TelemetryHistory, TelemetryHistoryService};
use anyhow::bail;
use api::messages::telemetry_definition::{
    TelemetryDefinitionRequest, TelemetryDefinitionResponse,
};
use api::messages::telemetry_entry::TelemetryEntry;
use chrono::SecondsFormat;
use log::{error, info, warn};
use papaya::{HashMap, HashMapRef, LocalGuard};
use std::fs;
use std::fs::File;
use std::hash::RandomState;
use std::io::{Read, Write};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::sleep;
use uuid::Uuid;

const RELEASED_ATTEMPTS: usize = 5;

pub struct TelemetryManagementService {
    uuid_index: HashMap<String, Uuid>,
    tlm_data: HashMap<Uuid, Arc<TelemetryHistory>>,
    telemetry_history_service: Arc<TelemetryHistoryService>,
    metadata_file: Arc<Mutex<File>>,
}

impl TelemetryManagementService {
    pub fn new(telemetry_history_service: TelemetryHistoryService) -> anyhow::Result<Self> {
        let metadata_file = telemetry_history_service.get_metadata_file();

        let uuid_index = HashMap::new();
        let tlm_data = HashMap::new();

        // TODO: Load metadata from file
        match File::open(&metadata_file) {
            Ok(mut metadata_file) => {
                let uuid_index = uuid_index.pin();
                let tlm_data = tlm_data.pin();

                // Read all data from the file
                let mut data = String::new();
                metadata_file.read_to_string(&mut data)?;
                drop(metadata_file);

                // Each line is a separate entry
                for line in data.split("\n") {
                    if line.is_empty() {
                        // Skip empty lines
                        continue;
                    }
                    // Skip invalid entries
                    match serde_json::from_str::<TelemetryDefinition>(line) {
                        Ok(tlm_def) => {
                            let _ = uuid_index.insert(tlm_def.name.clone(), tlm_def.uuid);
                            let _ = tlm_data.insert(tlm_def.uuid, Arc::new(tlm_def.into()));
                        }
                        Err(err) => {
                            error!("Failed to parse metadata entry {err}");
                        }
                    }
                }
            }
            Err(err) => {
                warn!("Failed to open metadata file {err}. Continuing");
            }
        }

        Ok(Self {
            uuid_index,
            tlm_data,
            telemetry_history_service: Arc::new(telemetry_history_service),
            metadata_file: Arc::new(Mutex::new(
                fs::OpenOptions::new()
                    .create(true)
                    .append(true)
                    .open(metadata_file)?,
            )),
        })
    }

    pub fn register(
        &self,
        telemetry_definition_request: TelemetryDefinitionRequest,
    ) -> anyhow::Result<TelemetryDefinitionResponse> {
        let uuid_index = self.uuid_index.pin();
        let tlm_data = self.tlm_data.pin();

        let uuid =
            *uuid_index.get_or_insert_with(telemetry_definition_request.name.clone(), Uuid::new_v4);

        let inserted = tlm_data.try_insert(
            uuid,
            Arc::new(
                TelemetryDefinition {
                    uuid,
                    name: telemetry_definition_request.name.clone(),
                    data_type: telemetry_definition_request.data_type,
                }
                .into(),
            ),
        );

        if let Ok(newly_inserted) = inserted {
            // This data also needs to be written to disk
            let file = self.metadata_file.clone();
            let newly_inserted = newly_inserted.data.definition.clone();
            tokio::spawn(async move {
                let mut file = file.lock_owned().await;
                let newly_inserted = serde_json::to_string(&newly_inserted);
                match newly_inserted {
                    Ok(newly_inserted) => {
                        if let Err(err) = file.write(newly_inserted.as_bytes()) {
                            error!("Failed to write TelemetryDefinition to file {err}");
                            return;
                        }
                        if let Err(err) = file.write("\n".as_bytes()) {
                            error!("Failed to write newline to file {err}");
                            return;
                        }
                        if let Err(err) = file.flush() {
                            error!("Failed to flush file {err}");
                        }
                    }
                    Err(err) => {
                        error!("Failed to serialize TelemetryDefinition {err}");
                    }
                }
            });
        }

        Ok(TelemetryDefinitionResponse { uuid })
    }

    pub fn add_tlm_item(&self, tlm_item: TelemetryEntry) -> anyhow::Result<()> {
        let tlm_management_pin = self.pin();
        let Some(tlm_data) = tlm_management_pin.get_by_uuid(&tlm_item.uuid) else {
            bail!("Telemetry Item Not Found");
        };

        let expected_type = tlm_item.value.to_data_type();
        if expected_type != tlm_data.data.definition.data_type {
            bail!("Data Type Mismatch");
        };

        let _ = tlm_data.data.data.send_replace(Some(TelemetryDataItem {
            value: tlm_item.value,
            timestamp: tlm_item
                .timestamp
                .to_rfc3339_opts(SecondsFormat::Millis, true),
        }));
        TelemetryHistory::insert_sync(
            tlm_data.clone(),
            self.history_service(),
            tlm_item.value,
            tlm_item.timestamp,
        );

        Ok(())
    }

    pub fn get_by_name(&self, name: &String) -> Option<TelemetryData> {
        let uuid_index = self.uuid_index.pin();
        let uuid = uuid_index.get(name)?;
        self.get_by_uuid(uuid)
    }

    pub fn get_by_uuid(&self, uuid: &Uuid) -> Option<TelemetryData> {
        let tlm_data = self.tlm_data.pin();
        tlm_data
            .get(uuid)
            .map(|data_history| &data_history.data)
            .cloned()
    }

    pub fn get_all_definitions(&self) -> Vec<TelemetryDefinition> {
        let tlm_data = self.tlm_data.pin();
        tlm_data
            .values()
            .map(|x| x.data.definition.clone())
            .collect()
    }

    pub fn pin(&self) -> TelemetryManagementServicePin<'_> {
        TelemetryManagementServicePin {
            tlm_data: self.tlm_data.pin(),
        }
    }

    pub fn history_service(&self) -> Arc<TelemetryHistoryService> {
        self.telemetry_history_service.clone()
    }

    pub async fn cleanup(self) -> anyhow::Result<()> {
        info!("Saving Telemetry");
        let history_service = self.telemetry_history_service;
        let metadata_file = self.metadata_file;
        let tlm_data = self.tlm_data.pin_owned();
        let mut tasks = vec![];

        for _ in 0..RELEASED_ATTEMPTS {
            if Arc::strong_count(&history_service) != 1 {
                sleep(Duration::from_secs(1)).await;
            }
        }
        if Arc::strong_count(&history_service) != 1 {
            warn!("History Service Has not been released after {RELEASED_ATTEMPTS} attempts!");
        }

        for data_history in tlm_data.values() {
            tasks.push(data_history.cleanup(&history_service));
        }
        {
            if let Some(tlm) = Arc::into_inner(metadata_file) {
                tlm.into_inner().sync_all()?;
            } else {
                error!("Could not close metadata file: still in use")
            }
        }

        for task in tasks {
            task.await?;
        }

        Ok(())
    }
}

pub struct TelemetryManagementServicePin<'a> {
    tlm_data: HashMapRef<'a, Uuid, Arc<TelemetryHistory>, RandomState, LocalGuard<'a>>,
}

impl<'a> TelemetryManagementServicePin<'a> {
    pub fn get_by_uuid(&'a self, uuid: &Uuid) -> Option<&'a Arc<TelemetryHistory>> {
        self.tlm_data.get(uuid)
    }
}