rpki/src/parallel/transport_prefetch.rs

use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::parallel::types::{
    RepoDedupKey, RepoIdentity, RepoRequester, RepoTransportMode, RepoTransportTask,
};
use crate::policy::SyncPreference;

pub const TRANSPORT_PREFETCH_SCHEMA_VERSION: u32 = 1;

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchSnapshot {
    pub schema_version: u32,
    pub generated_at_unix: i64,
    pub sync_preference: SyncPreference,
    pub requests: Vec<TransportPrefetchRequest>,
}

impl TransportPrefetchSnapshot {
    pub fn new(sync_preference: SyncPreference, requests: Vec<TransportPrefetchRequest>) -> Self {
        Self {
            schema_version: TRANSPORT_PREFETCH_SCHEMA_VERSION,
            generated_at_unix: time::OffsetDateTime::now_utc().unix_timestamp(),
            sync_preference,
            requests,
        }
    }

    pub fn is_compatible_with(&self, sync_preference: SyncPreference) -> bool {
        self.schema_version == TRANSPORT_PREFETCH_SCHEMA_VERSION
            && self.sync_preference == sync_preference
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchRequest {
    pub dedup_key: TransportPrefetchDedupKey,
    pub rsync_scope_uri: String,
    pub rsync_failure_scope_uri: Option<String>,
    pub repo_identity: TransportPrefetchRepoIdentity,
    pub mode: TransportPrefetchMode,
    pub tal_id: String,
    pub rir_id: String,
    pub priority: u8,
    pub requesters: Vec<TransportPrefetchRequester>,
}

impl TransportPrefetchRequest {
    pub fn from_registered_request(
        identity: &RepoIdentity,
        requester: &RepoRequester,
        priority: u8,
        rsync_scope_uri: String,
        rsync_failure_scope_uri: Option<String>,
        sync_preference: SyncPreference,
    ) -> Self {
        let (dedup_key, mode) = if sync_preference == SyncPreference::RrdpThenRsync {
            if let Some(notification_uri) = identity.notification_uri.clone() {
                (
                    TransportPrefetchDedupKey::RrdpNotify { notification_uri },
                    TransportPrefetchMode::Rrdp,
                )
            } else {
                (
                    TransportPrefetchDedupKey::RsyncScope {
                        rsync_scope_uri: rsync_scope_uri.clone(),
                    },
                    TransportPrefetchMode::Rsync,
                )
            }
        } else {
            (
                TransportPrefetchDedupKey::RsyncScope {
                    rsync_scope_uri: rsync_scope_uri.clone(),
                },
                TransportPrefetchMode::Rsync,
            )
        };

        Self {
            dedup_key,
            rsync_scope_uri,
            rsync_failure_scope_uri,
            repo_identity: TransportPrefetchRepoIdentity::from_identity(identity),
            mode,
            tal_id: requester.tal_id.clone(),
            rir_id: requester.rir_id.clone(),
            priority,
            requesters: vec![TransportPrefetchRequester::from_requester(requester)],
        }
    }

    pub fn from_task(task: &RepoTransportTask, rsync_scope_uri: String) -> Self {
        Self {
            dedup_key: TransportPrefetchDedupKey::from_repo_key(&task.dedup_key),
            rsync_scope_uri,
            rsync_failure_scope_uri: task.rsync_failure_scope_uri.clone(),
            repo_identity: TransportPrefetchRepoIdentity::from_identity(&task.repo_identity),
            mode: TransportPrefetchMode::from_mode(task.mode),
            tal_id: task.tal_id.clone(),
            rir_id: task.rir_id.clone(),
            priority: task.priority,
            requesters: task
                .requesters
                .iter()
                .map(TransportPrefetchRequester::from_requester)
                .collect(),
        }
    }

    pub fn to_identity(&self) -> RepoIdentity {
        self.repo_identity.to_identity()
    }

    pub fn to_requester(&self) -> RepoRequester {
        self.requesters
            .first()
            .map(TransportPrefetchRequester::to_requester)
            .unwrap_or_else(|| RepoRequester {
                tal_id: self.tal_id.clone(),
                rir_id: self.rir_id.clone(),
                parent_node_id: None,
                ca_instance_handle_id: format!(
                    "{}:{}",
                    self.tal_id, self.repo_identity.rsync_base_uri
                ),
                publication_point_rsync_uri: self.repo_identity.rsync_base_uri.clone(),
                manifest_rsync_uri: format!("{}prefetch.mft", self.repo_identity.rsync_base_uri),
            })
    }

    fn recorder_key(&self) -> String {
        self.dedup_key.stable_key()
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TransportPrefetchDedupKey {
    RrdpNotify { notification_uri: String },
    RsyncScope { rsync_scope_uri: String },
}

impl TransportPrefetchDedupKey {
    fn from_repo_key(key: &RepoDedupKey) -> Self {
        match key {
            RepoDedupKey::RrdpNotify { notification_uri } => Self::RrdpNotify {
                notification_uri: notification_uri.clone(),
            },
            RepoDedupKey::RsyncScope { rsync_scope_uri } => Self::RsyncScope {
                rsync_scope_uri: rsync_scope_uri.clone(),
            },
        }
    }

    fn stable_key(&self) -> String {
        match self {
            Self::RrdpNotify { notification_uri } => format!("rrdp:{notification_uri}"),
            Self::RsyncScope { rsync_scope_uri } => format!("rsync:{rsync_scope_uri}"),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchRepoIdentity {
    pub notification_uri: Option<String>,
    pub rsync_base_uri: String,
}

impl TransportPrefetchRepoIdentity {
    fn from_identity(identity: &RepoIdentity) -> Self {
        Self {
            notification_uri: identity.notification_uri.clone(),
            rsync_base_uri: identity.rsync_base_uri.clone(),
        }
    }

    fn to_identity(&self) -> RepoIdentity {
        RepoIdentity::new(self.notification_uri.clone(), self.rsync_base_uri.clone())
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TransportPrefetchMode {
    Rrdp,
    Rsync,
}

impl TransportPrefetchMode {
    fn from_mode(mode: RepoTransportMode) -> Self {
        match mode {
            RepoTransportMode::Rrdp => Self::Rrdp,
            RepoTransportMode::Rsync => Self::Rsync,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchRequester {
    pub tal_id: String,
    pub rir_id: String,
    pub parent_node_id: Option<u64>,
    pub ca_instance_handle_id: String,
    pub publication_point_rsync_uri: String,
    pub manifest_rsync_uri: String,
}

impl TransportPrefetchRequester {
    fn from_requester(requester: &RepoRequester) -> Self {
        Self {
            tal_id: requester.tal_id.clone(),
            rir_id: requester.rir_id.clone(),
            parent_node_id: requester.parent_node_id,
            ca_instance_handle_id: requester.ca_instance_handle_id.clone(),
            publication_point_rsync_uri: requester.publication_point_rsync_uri.clone(),
            manifest_rsync_uri: requester.manifest_rsync_uri.clone(),
        }
    }

    fn to_requester(&self) -> RepoRequester {
        RepoRequester {
            tal_id: self.tal_id.clone(),
            rir_id: self.rir_id.clone(),
            parent_node_id: self.parent_node_id,
            ca_instance_handle_id: self.ca_instance_handle_id.clone(),
            publication_point_rsync_uri: self.publication_point_rsync_uri.clone(),
            manifest_rsync_uri: self.manifest_rsync_uri.clone(),
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TransportPrefetchRecorder {
    requests_by_key: BTreeMap<String, TransportPrefetchRequest>,
    request_order: Vec<String>,
}

impl TransportPrefetchRecorder {
    pub fn record_registered_request(
        &mut self,
        identity: &RepoIdentity,
        requester: &RepoRequester,
        priority: u8,
        rsync_scope_uri: String,
        rsync_failure_scope_uri: Option<String>,
        sync_preference: SyncPreference,
    ) {
        let request = TransportPrefetchRequest::from_registered_request(
            identity,
            requester,
            priority,
            rsync_scope_uri,
            rsync_failure_scope_uri,
            sync_preference,
        );
        self.record_request(request);
    }

    pub fn record_task(&mut self, task: &RepoTransportTask, rsync_scope_uri: String) {
        let request = TransportPrefetchRequest::from_task(task, rsync_scope_uri);
        self.record_request(request);
    }

    pub fn snapshot(&self, sync_preference: SyncPreference) -> TransportPrefetchSnapshot {
        TransportPrefetchSnapshot::new(
            sync_preference,
            self.request_order
                .iter()
                .filter_map(|key| self.requests_by_key.get(key))
                .cloned()
                .collect(),
        )
    }

    pub fn len(&self) -> usize {
        self.requests_by_key.len()
    }

    fn record_request(&mut self, request: TransportPrefetchRequest) {
        let key = request.recorder_key();
        if !self.requests_by_key.contains_key(&key) {
            self.request_order.push(key.clone());
            self.requests_by_key.insert(key, request);
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TransportPrefetchDispatchStats {
    pub loaded_requests: u64,
    pub enqueued_tasks: u64,
    pub waiting_requests: u64,
    pub reused_results: u64,
    pub skipped_incompatible: u64,
}

#[cfg(test)]
mod tests {
    use super::*;

    fn requester(uri: &str) -> RepoRequester {
        RepoRequester {
            tal_id: "apnic".to_string(),
            rir_id: "apnic".to_string(),
            parent_node_id: None,
            ca_instance_handle_id: format!("apnic:{uri}"),
            publication_point_rsync_uri: "rsync://example.test/repo/".to_string(),
            manifest_rsync_uri: uri.to_string(),
        }
    }

    fn task(notification_uri: &str, manifest_uri: &str) -> RepoTransportTask {
        RepoTransportTask {
            dedup_key: RepoDedupKey::RrdpNotify {
                notification_uri: notification_uri.to_string(),
            },
            rsync_failure_scope_uri: Some("rsync://example.test/".to_string()),
            repo_identity: RepoIdentity::new(
                Some(notification_uri.to_string()),
                "rsync://example.test/repo/",
            ),
            mode: RepoTransportMode::Rrdp,
            tal_id: "apnic".to_string(),
            rir_id: "apnic".to_string(),
            validation_time: time::OffsetDateTime::UNIX_EPOCH,
            priority: 0,
            requesters: vec![requester(manifest_uri)],
        }
    }

    #[test]
    fn recorder_deduplicates_by_transport_key() {
        let mut recorder = TransportPrefetchRecorder::default();
        recorder.record_task(
            &task(
                "https://example.test/notification.xml",
                "rsync://example.test/repo/a.mft",
            ),
            "rsync://example.test/repo/".to_string(),
        );
        recorder.record_task(
            &task(
                "https://example.test/notification.xml",
                "rsync://example.test/repo/b.mft",
            ),
            "rsync://example.test/repo/".to_string(),
        );

        let snapshot = recorder.snapshot(SyncPreference::RrdpThenRsync);
        assert_eq!(snapshot.requests.len(), 1);
        assert_eq!(
            snapshot.requests[0].requesters[0].manifest_rsync_uri,
            "rsync://example.test/repo/a.mft"
        );
    }

    #[test]
    fn recorder_preserves_first_discovery_order() {
        let mut recorder = TransportPrefetchRecorder::default();
        recorder.record_task(
            &task(
                "https://z.example.test/notification.xml",
                "rsync://z.example.test/repo/root.mft",
            ),
            "rsync://z.example.test/repo/".to_string(),
        );
        recorder.record_task(
            &task(
                "https://a.example.test/notification.xml",
                "rsync://a.example.test/repo/root.mft",
            ),
            "rsync://a.example.test/repo/".to_string(),
        );

        let snapshot = recorder.snapshot(SyncPreference::RrdpThenRsync);
        assert_eq!(
            snapshot.requests[0]
                .repo_identity
                .notification_uri
                .as_deref(),
            Some("https://z.example.test/notification.xml")
        );
        assert_eq!(
            snapshot.requests[1]
                .repo_identity
                .notification_uri
                .as_deref(),
            Some("https://a.example.test/notification.xml")
        );
    }

    #[test]
    fn snapshot_checks_schema_and_sync_preference() {
        let snapshot = TransportPrefetchSnapshot::new(SyncPreference::RrdpThenRsync, Vec::new());
        assert!(snapshot.is_compatible_with(SyncPreference::RrdpThenRsync));
        assert!(!snapshot.is_compatible_with(SyncPreference::RsyncOnly));
    }
}