| 1 | use std::collections::{BTreeMap, BTreeSet}; |
| 2 | |
| 3 | use ruma::OwnedUserId; |
| 4 | use serde_json::Value; |
| 5 | |
| 6 | use crate::{ |
| 7 | Error, ErrorKind, EventContext, MentionCaptureState, MentionCoverage, MentionEntry, MentionKey, |
| 8 | MentionsBoundary, MentionsPage, SavedMentionContext, TimelineItem, |
| 9 | normalise::MediaRegistration, |
| 10 | }; |
| 11 | |
| 12 | #[derive(Clone)] |
| 13 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 14 | pub(super) struct SavedEvent { |
| 15 | pub item: TimelineItem, |
| 16 | pub media: Vec<MediaRegistration>, |
| 17 | pub raw: Option<Value>, |
| 18 | } |
| 19 | |
| 20 | #[derive(Clone)] |
| 21 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 22 | pub(super) struct Candidate { |
| 23 | pub room_id: ruma::OwnedRoomId, |
| 24 | pub raw: Value, |
| 25 | pub attempts: u32, |
| 26 | pub retry_at_ms: u64, |
| 27 | } |
| 28 | |
| 29 | #[derive(Clone)] |
| 30 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 31 | pub(super) struct Window { |
| 32 | pub before: Vec<String>, |
| 33 | pub target: Option<String>, |
| 34 | pub after: Vec<String>, |
| 35 | pub start: Option<String>, |
| 36 | pub end: Option<String>, |
| 37 | } |
| 38 | |
| 39 | #[derive(Clone, Default)] |
| 40 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 41 | pub(super) struct Data { |
| 42 | pub revision: u64, |
| 43 | pub entries: BTreeMap<String, MentionEntry>, |
| 44 | pub events: BTreeMap<String, SavedEvent>, |
| 45 | pub windows: BTreeMap<String, Window>, |
| 46 | pub capture_attempts: BTreeMap<String, u32>, |
| 47 | pub candidates: BTreeMap<String, Candidate>, |
| 48 | pub redacted: BTreeMap<String, u64>, |
| 49 | pub edits: BTreeMap<String, Value>, |
| 50 | pub coverage: MentionCoverage, |
| 51 | pub checkpoints: BTreeMap<String, String>, |
| 52 | pub observed_rooms: BTreeSet<ruma::OwnedRoomId>, |
| 53 | pub history_cursor: Option<String>, |
| 54 | pub history_exhausted: bool, |
| 55 | pub latest_cursor: Option<String>, |
| 56 | #[cfg_attr(feature = "serde", serde(default))] |
| 57 | pub latest_anchor: Option<String>, |
| 58 | #[cfg_attr(feature = "serde", serde(default))] |
| 59 | pub latest_head: Option<String>, |
| 60 | } |
| 61 | |
| 62 | pub(super) fn id(key: &MentionKey) -> String { |
| 63 | format!("{}\n{}", key.room_id, key.event_id) |
| 64 | } |
| 65 | |
| 66 | pub(super) fn now() -> u64 { |
| 67 | matrix_sdk::ruma::MilliSecondsSinceUnixEpoch::now() |
| 68 | .get() |
| 69 | .into() |
| 70 | } |
| 71 | |
| 72 | pub(super) fn invalid(message: &'static str) -> Error { |
| 73 | Error::new(ErrorKind::Invalid, message) |
| 74 | } |
| 75 | |
| 76 | impl Data { |
| 77 | pub fn touch_contexts(&mut self, identity: &str) { |
| 78 | for (key, window) in &self.windows { |
| 79 | if window.target.as_deref() == Some(identity) |
| 80 | || window |
| 81 | .before |
| 82 | .iter() |
| 83 | .chain(&window.after) |
| 84 | .any(|key| key == identity) |
| 85 | { |
| 86 | if let Some(entry) = self.entries.get_mut(key) { |
| 87 | entry.revision = self.revision + 1; |
| 88 | } |
| 89 | } |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | pub fn page( |
| 94 | &self, |
| 95 | account: &OwnedUserId, |
| 96 | cursor: Option<&str>, |
| 97 | limit: u16, |
| 98 | ) -> Result<MentionsPage, Error> { |
| 99 | let (boundary, last) = if let Some(cursor) = cursor { |
| 100 | let value: Value = |
| 101 | serde_json::from_str(cursor).map_err(|_| invalid("invalid mentions cursor"))?; |
| 102 | if value["account"].as_str() != Some(account.as_str()) { |
| 103 | return Err(invalid("mentions cursor belongs to a different account")); |
| 104 | } |
| 105 | let boundary = value["revision"] |
| 106 | .as_u64() |
| 107 | .ok_or_else(|| invalid("invalid mentions cursor"))?; |
| 108 | let timestamp = value["timestamp"] |
| 109 | .as_u64() |
| 110 | .ok_or_else(|| invalid("invalid mentions cursor"))?; |
| 111 | let key = value["key"] |
| 112 | .as_str() |
| 113 | .ok_or_else(|| invalid("invalid mentions cursor"))? |
| 114 | .to_owned(); |
| 115 | if boundary > self.revision { |
| 116 | return Err(invalid("mentions cursor is newer than this store")); |
| 117 | } |
| 118 | (boundary, Some((timestamp, key))) |
| 119 | } else { |
| 120 | (self.revision, None) |
| 121 | }; |
| 122 | let mut entries: Vec<_> = self |
| 123 | .entries |
| 124 | .iter() |
| 125 | .filter(|(key, entry)| { |
| 126 | entry.discovery_revision <= boundary |
| 127 | && last.as_ref().is_none_or(|last| { |
| 128 | (entry.timestamp_ms, key.as_str()) < (last.0, last.1.as_str()) |
| 129 | }) |
| 130 | }) |
| 131 | .map(|(_, entry)| entry.clone()) |
| 132 | .collect(); |
| 133 | entries.sort_by(|a, b| (b.timestamp_ms, id(&b.key)).cmp(&(a.timestamp_ms, id(&a.key)))); |
| 134 | let limit = usize::from(limit.clamp(1, 50)); |
| 135 | let end_reached = entries.len() <= limit; |
| 136 | entries.truncate(limit); |
| 137 | let next_cursor = if end_reached { |
| 138 | None |
| 139 | } else { |
| 140 | entries.last().map(|entry| serde_json::json!({ |
| 141 | "account": account.as_str(), "revision": boundary, "timestamp": entry.timestamp_ms, "key": id(&entry.key) |
| 142 | }).to_string()) |
| 143 | }; |
| 144 | let mut coverage = self.coverage.clone(); |
| 145 | coverage.unresolved_candidates = self.candidates.len(); |
| 146 | Ok(MentionsPage { |
| 147 | entries, |
| 148 | next_cursor, |
| 149 | end_reached, |
| 150 | coverage, |
| 151 | unread_count: self.entries.values().filter(|entry| !entry.read).count(), |
| 152 | revision: self.revision, |
| 153 | boundary: MentionsBoundary { |
| 154 | account: account.clone(), |
| 155 | revision: boundary, |
| 156 | }, |
| 157 | }) |
| 158 | } |
| 159 | |
| 160 | pub fn context( |
| 161 | &self, |
| 162 | key: &MentionKey, |
| 163 | ) -> Result<(SavedMentionContext, Vec<MediaRegistration>), Error> { |
| 164 | let entry = self |
| 165 | .entries |
| 166 | .get(&id(key)) |
| 167 | .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "mention was not found"))?; |
| 168 | let mut media = Vec::new(); |
| 169 | let mut item = |key: &String| { |
| 170 | self.events.get(key).map(|event| { |
| 171 | media.extend(event.media.clone()); |
| 172 | event.item.clone() |
| 173 | }) |
| 174 | }; |
| 175 | let context = self.windows.get(&id(key)).map(|window| EventContext { |
| 176 | events_before: window.before.iter().filter_map(&mut item).collect(), |
| 177 | event: window.target.as_ref().and_then(&mut item), |
| 178 | events_after: window.after.iter().filter_map(&mut item).collect(), |
| 179 | start_token: window.start.clone(), |
| 180 | end_token: window.end.clone(), |
| 181 | }); |
| 182 | Ok(( |
| 183 | SavedMentionContext { |
| 184 | key: key.clone(), |
| 185 | state: entry.capture.clone(), |
| 186 | context, |
| 187 | revision: self.revision, |
| 188 | }, |
| 189 | media, |
| 190 | )) |
| 191 | } |
| 192 | |
| 193 | pub fn collect_events(&mut self) { |
| 194 | let mut used: BTreeSet<String> = self.entries.keys().cloned().collect(); |
| 195 | for window in self.windows.values() { |
| 196 | used.extend(window.before.iter().chain(&window.after).cloned()); |
| 197 | used.extend(window.target.iter().cloned()); |
| 198 | } |
| 199 | self.events.retain(|key, _| used.contains(key)); |
| 200 | self.capture_attempts |
| 201 | .retain(|key, _| self.entries.contains_key(key)); |
| 202 | } |
| 203 | |
| 204 | pub fn evict_one(&mut self) -> bool { |
| 205 | let oldest = |keys: Vec<String>, entries: &BTreeMap<String, MentionEntry>| { |
| 206 | keys.into_iter() |
| 207 | .min_by_key(|key| entries.get(key).map_or(0, |entry| entry.timestamp_ms)) |
| 208 | }; |
| 209 | if let Some(key) = oldest( |
| 210 | self.windows |
| 211 | .keys() |
| 212 | .filter(|key| self.entries.get(*key).is_some_and(|entry| entry.read)) |
| 213 | .cloned() |
| 214 | .collect(), |
| 215 | &self.entries, |
| 216 | ) |
| 217 | .or_else(|| oldest(self.windows.keys().cloned().collect(), &self.entries)) |
| 218 | { |
| 219 | self.windows.remove(&key); |
| 220 | if let Some(entry) = self.entries.get_mut(&key) { |
| 221 | entry.capture = MentionCaptureState::Unavailable; |
| 222 | entry.revision = self.revision; |
| 223 | } |
| 224 | } else if let Some(key) = self |
| 225 | .entries |
| 226 | .iter() |
| 227 | .min_by_key(|(_, entry)| (!entry.read, entry.timestamp_ms)) |
| 228 | .map(|(key, _)| key.clone()) |
| 229 | { |
| 230 | self.entries.remove(&key); |
| 231 | } else if let Some(key) = self.candidates.keys().next().cloned() { |
| 232 | self.candidates.remove(&key); |
| 233 | } else if let Some(key) = self.edits.keys().next().cloned() { |
| 234 | self.edits.remove(&key); |
| 235 | } else if let Some(key) = self.redacted.keys().next().cloned() { |
| 236 | self.redacted.remove(&key); |
| 237 | } else { |
| 238 | return false; |
| 239 | } |
| 240 | self.coverage.truncated = true; |
| 241 | self.collect_events(); |
| 242 | true |
| 243 | } |
| 244 | |
| 245 | pub fn bound(&mut self, entries: usize, events: usize, candidates: usize) { |
| 246 | while self.entries.len() > entries { |
| 247 | if let Some(key) = self |
| 248 | .entries |
| 249 | .iter() |
| 250 | .min_by_key(|(_, entry)| (!entry.read, entry.timestamp_ms)) |
| 251 | .map(|(key, _)| key.clone()) |
| 252 | { |
| 253 | self.entries.remove(&key); |
| 254 | self.windows.remove(&key); |
| 255 | self.coverage.truncated = true; |
| 256 | } |
| 257 | } |
| 258 | self.collect_events(); |
| 259 | while self.events.len() > events && self.evict_one() {} |
| 260 | while self.candidates.len() > candidates { |
| 261 | let key = self.candidates.keys().next().cloned().unwrap(); |
| 262 | self.candidates.remove(&key); |
| 263 | self.coverage.truncated = true; |
| 264 | } |
| 265 | for map in [&mut self.redacted] { |
| 266 | while map.len() > entries + candidates { |
| 267 | let key = map |
| 268 | .iter() |
| 269 | .min_by_key(|(_, timestamp)| *timestamp) |
| 270 | .map(|(key, _)| key.clone()) |
| 271 | .unwrap(); |
| 272 | map.remove(&key); |
| 273 | self.coverage.truncated = true; |
| 274 | } |
| 275 | } |
| 276 | while self.edits.len() > candidates { |
| 277 | let key = self.edits.keys().next().cloned().unwrap(); |
| 278 | self.edits.remove(&key); |
| 279 | self.coverage.truncated = true; |
| 280 | } |
| 281 | self.coverage.unresolved_candidates = self.candidates.len(); |
| 282 | } |
| 283 | } |