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, } impl TransportPrefetchSnapshot { pub fn new(sync_preference: SyncPreference, requests: Vec) -> 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, pub repo_identity: TransportPrefetchRepoIdentity, pub mode: TransportPrefetchMode, pub tal_id: String, pub rir_id: String, pub priority: u8, pub requesters: Vec, } impl TransportPrefetchRequest { pub fn from_registered_request( identity: &RepoIdentity, requester: &RepoRequester, priority: u8, rsync_scope_uri: String, rsync_failure_scope_uri: Option, 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, 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, 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, request_order: Vec, } impl TransportPrefetchRecorder { pub fn record_registered_request( &mut self, identity: &RepoIdentity, requester: &RepoRequester, priority: u8, rsync_scope_uri: String, rsync_failure_scope_uri: Option, 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)); } }