use super::nostr_event_extensions::{MaybeConvertibleToRelayList, RelayList, TimestampedMuteList}; use crate::{ notification_manager::nostr_event_extensions::MaybeConvertibleToTimestampedMuteList, utils::time_delta::TimeDelta, }; use log; use nostr_sdk::prelude::*; use std::collections::HashMap; use std::str::FromStr; use tokio::time::{Duration, Instant}; struct CacheEntry { value: Option, // `None` means the event does not exist as far as we know (It does NOT mean expired) added_at: Instant, } impl CacheEntry { fn is_expired(&self, max_age: Duration) -> bool { self.added_at.elapsed() > max_age } pub fn new(value: T) -> Self { let added_at = Instant::now(); CacheEntry { value: Some(value), added_at, } } pub fn maybe(value: Option) -> Self { let added_at = Instant::now(); CacheEntry { value, added_at } } pub fn empty() -> Self { let added_at = Instant::now(); CacheEntry { value: None, added_at, } } pub fn value(&self) -> Option<&T> { self.value.as_ref() } } pub struct Cache { //entries: HashMap, mute_lists: HashMap>, contact_lists: HashMap>, relay_lists: HashMap>, max_age: Duration, } fn get_cache_entry( list: &mut HashMap>, pubkey: &PublicKey, max_age: Duration, name: &str, ) -> Result, CacheError> { let res = if let Some(entry) = list.get(pubkey) { if !entry.is_expired(max_age) { Ok(entry.value().cloned()) } else { log::debug!( "{} list for pubkey {} is expired, removing it from the cache", name, pubkey.to_hex() ); Err(CacheError::Expired) } } else { Err(CacheError::NotFound) }; if let Err(CacheError::Expired) = &res { list.remove(pubkey); } res } impl Cache { // MARK: - Initialization pub fn new(max_age: Duration) -> Self { Cache { //entries: HashMap::new(), mute_lists: HashMap::new(), contact_lists: HashMap::new(), relay_lists: HashMap::new(), max_age, } } // MARK: - Adding items to the cache pub fn add_optional_mute_list_with_author<'a>( &'a mut self, author: &PublicKey, mute_list: Option<&Event>, ) { if let Some(mute_list) = mute_list { self.add_event(mute_list); } else { self.mute_lists .insert(author.to_owned(), CacheEntry::empty()); } } pub fn add_optional_relay_list_with_author<'a>( &'a mut self, author: &PublicKey, relay_list_event: Option<&Event>, ) { if let Some(relay_list_event) = relay_list_event { self.add_event(relay_list_event); } else { self.relay_lists .insert(author.to_owned(), CacheEntry::empty()); } } pub fn add_optional_contact_list_with_author<'a>( &'a mut self, author: &PublicKey, contact_list: Option<&Event>, ) { if let Some(contact_list) = contact_list { self.add_event(contact_list); } else { self.contact_lists .insert(author.to_owned(), CacheEntry::empty()); } } pub fn add_event(&mut self, event: &Event) { match event.kind { Kind::MuteList => { self.mute_lists.insert( event.pubkey.clone(), CacheEntry::maybe(event.to_timestamped_mute_list()), ); log::debug!( "Added mute list to the cache. Event ID: {}", event.id.to_hex() ); } Kind::ContactList => { log::debug!( "Added contact list to the cache. Event ID: {}", event.id.to_hex() ); self.contact_lists .insert(event.pubkey.clone(), CacheEntry::new(event.to_owned())); } Kind::RelayList => { log::debug!( "Added relay list to the cache. Event ID: {}", event.id.to_hex() ); self.relay_lists.insert( event.pubkey.clone(), CacheEntry::maybe(event.to_relay_list()), ); } _ => { log::debug!( "Unknown event kind, not adding to any cache. Event ID: {}", event.id.to_hex() ); } } } // MARK: - Fetching items from the cache pub fn get_mute_list( &mut self, pubkey: &PublicKey, ) -> Result, CacheError> { get_cache_entry(&mut self.mute_lists, pubkey, self.max_age, "Mute") } pub fn get_relay_list(&mut self, pubkey: &PublicKey) -> Result, CacheError> { get_cache_entry(&mut self.relay_lists, pubkey, self.max_age, "Relay") } pub fn get_contact_list(&mut self, pubkey: &PublicKey) -> Result, CacheError> { get_cache_entry(&mut self.contact_lists, pubkey, self.max_age, "Contact") } } // Error type #[derive(Debug, Eq, PartialEq)] pub enum CacheError { NotFound, Expired, } #[cfg(test)] mod tests { use super::*; use nostr_sdk::prelude::*; use std::time::Duration; use tokio::time::sleep; // Helper function to create a dummy event of a given kind for testing. fn create_dummy_event(pubkey: PublicKey, kind: Kind) -> Event { // In a real test, you might generate keys or events more dynamically. let id = EventId::from_hex("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa") .unwrap(); let created_at = Timestamp::now(); let content = ""; let sig_str = "8e1a61523765a6e577e3ca0c87afe3694ed518719aea067701c35262dd2a3c7e3ca0946fe98463a3af706dd333695ceec6cb3b29254c557c8630d3db1171ea3d"; let sig = Signature::from_str(sig_str).unwrap(); Event::new(id, pubkey, created_at, kind, [], content, sig) } // Helper function to create a dummy public key for testing. fn create_dummy_pubkey() -> PublicKey { // In a real project, you'd generate a key. For the sake of tests, just parse a known hex. PublicKey::from_hex("32e1827635450ebb3c5a7d12c1f8e7b2b514439ac10a67eef3d9fd9c5c68e245") .unwrap() } #[tokio::test] async fn test_add_and_retrieve_contact_list() { let pubkey = create_dummy_pubkey(); let max_age = Duration::from_secs(60); let mut cache = Cache::new(max_age); // Initially, no contact list should be found. assert!(matches!( cache.get_contact_list(&pubkey), Err(CacheError::NotFound) )); // Add a contact list event. let event = create_dummy_event(pubkey, Kind::ContactList); cache.add_event(&event); // Now we should be able to retrieve it. let retrieved = cache.get_contact_list(&pubkey).unwrap(); assert!(retrieved.is_some()); let retrieved_event = retrieved.unwrap(); assert_eq!(retrieved_event.id, event.id); } #[tokio::test] async fn test_add_and_retrieve_mute_list() { let pubkey = create_dummy_pubkey(); let max_age = Duration::from_secs(60); let mut cache = Cache::new(max_age); // No mute list initially assert!(matches!( cache.get_mute_list(&pubkey), Err(CacheError::NotFound) )); // Add a mute list event. let mutelist_event = { let event = create_dummy_event(pubkey, Kind::MuteList); event }; cache.add_event(&mutelist_event); let retrieved = cache.get_mute_list(&pubkey).unwrap(); assert!(retrieved.is_some()); // Should have a Some(TimestampedMuteList) now } #[tokio::test] async fn test_add_and_retrieve_relay_list() { let pubkey = create_dummy_pubkey(); let max_age = Duration::from_secs(60); let mut cache = Cache::new(max_age); // No relay list initially assert!(matches!( cache.get_relay_list(&pubkey), Err(CacheError::NotFound) )); // Add a relay list event. let relaylist_event = create_dummy_event(pubkey, Kind::RelayList); cache.add_event(&relaylist_event); let retrieved = cache.get_relay_list(&pubkey).unwrap(); assert!(retrieved.is_some()); } #[tokio::test] async fn test_expired_entries() { // Very short max_age to test expiration logic quickly. let max_age = Duration::from_millis(100); let pubkey = create_dummy_pubkey(); let mut cache = Cache::new(max_age); // Add a contact list event that will expire soon. let event = create_dummy_event(pubkey, Kind::ContactList); cache.add_event(&event); // Initially, we can retrieve it. let retrieved = cache.get_contact_list(&pubkey).unwrap(); assert!(retrieved.is_some()); // Wait for it to expire. sleep(Duration::from_millis(200)).await; // Now it should be expired and removed. let result = cache.get_contact_list(&pubkey); assert_eq!(result, Err(CacheError::Expired)); } #[tokio::test] async fn test_empty_entries() { let pubkey = create_dummy_pubkey(); let max_age = Duration::from_secs(60); let mut cache = Cache::new(max_age); // Add empty mute list cache.add_optional_mute_list_with_author(&pubkey, None); // We should now find a mute list entry, but it's None. let result = cache.get_mute_list(&pubkey).unwrap(); assert!(result.is_none()); // Add empty contact list cache.add_optional_contact_list_with_author(&pubkey, None); let result = cache.get_contact_list(&pubkey).unwrap(); assert!(result.is_none()); // Add empty relay list cache.add_optional_relay_list_with_author(&pubkey, None); let result = cache.get_relay_list(&pubkey).unwrap(); assert!(result.is_none()); } #[tokio::test] async fn test_re_insertion() { let pubkey = create_dummy_pubkey(); let max_age = Duration::from_secs(60); let mut cache = Cache::new(max_age); // Insert empty first cache.add_optional_contact_list_with_author(&pubkey, None); assert!(cache.get_contact_list(&pubkey).unwrap().is_none()); // Now insert a real event let event = create_dummy_event(pubkey, Kind::ContactList); cache.add_event(&event); // It should now return the actual event let retrieved = cache.get_contact_list(&pubkey).unwrap().unwrap(); assert_eq!(retrieved.id, event.id); } }