| 1 | use std::collections::{BTreeMap, BTreeSet}; |
| 2 | |
| 3 | use ruma::{OwnedRoomId, OwnedUserId}; |
| 4 | use serde_json::Value; |
| 5 | |
| 6 | use crate::{ |
| 7 | Error, ErrorKind, EventContext, MentionCaptureState, MentionCoverage, MentionEntry, MentionGap, |
| 8 | MentionGapReason, MentionKey, MentionsBoundary, MentionsPage, SavedMentionContext, |
| 9 | TimelineItem, normalise::MediaRegistration, |
| 10 | }; |
| 11 | |
| 12 | const CANDIDATE_LIFETIME_MS: u64 = 7 * 24 * 60 * 60 * 1000; |
| 13 | |
| 14 | #[derive(Clone)] |
| 15 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 16 | pub(super) struct SavedEvent { |
| 17 | pub item: TimelineItem, |
| 18 | pub media: Vec<MediaRegistration>, |
| 19 | pub raw: Option<Value>, |
| 20 | } |
| 21 | |
| 22 | impl PartialEq for SavedEvent { |
| 23 | fn eq(&self, other: &Self) -> bool { |
| 24 | self.item == other.item |
| 25 | && self.raw == other.raw |
| 26 | && self |
| 27 | .media |
| 28 | .iter() |
| 29 | .map(|media| &media.id) |
| 30 | .eq(other.media.iter().map(|media| &media.id)) |
| 31 | } |
| 32 | } |
| 33 | |
| 34 | #[derive(Clone, PartialEq)] |
| 35 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 36 | pub(super) struct Candidate { |
| 37 | pub room_id: ruma::OwnedRoomId, |
| 38 | pub raw: Value, |
| 39 | pub attempts: u32, |
| 40 | pub retry_at_ms: u64, |
| 41 | pub first_seen_ms: u64, |
| 42 | } |
| 43 | |
| 44 | impl Candidate { |
| 45 | pub fn new(room_id: ruma::OwnedRoomId, raw: Value, retry_at_ms: u64) -> Self { |
| 46 | Self { |
| 47 | room_id, |
| 48 | raw, |
| 49 | attempts: 0, |
| 50 | retry_at_ms, |
| 51 | first_seen_ms: now(), |
| 52 | } |
| 53 | } |
| 54 | |
| 55 | pub fn encrypted(&self) -> bool { |
| 56 | self.raw["type"] == "m.room.encrypted" |
| 57 | } |
| 58 | } |
| 59 | |
| 60 | #[derive(Clone, PartialEq)] |
| 61 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 62 | pub(super) struct Window { |
| 63 | pub before: Vec<String>, |
| 64 | pub target: Option<String>, |
| 65 | pub after: Vec<String>, |
| 66 | pub start: Option<String>, |
| 67 | pub end: Option<String>, |
| 68 | } |
| 69 | |
| 70 | #[derive(Clone, Default, PartialEq)] |
| 71 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 72 | pub(super) struct Data { |
| 73 | pub revision: u64, |
| 74 | pub entries: BTreeMap<String, MentionEntry>, |
| 75 | pub events: BTreeMap<String, SavedEvent>, |
| 76 | pub windows: BTreeMap<String, Window>, |
| 77 | pub capture_attempts: BTreeMap<String, u32>, |
| 78 | pub candidates: BTreeMap<String, Candidate>, |
| 79 | pub redacted: BTreeMap<String, u64>, |
| 80 | pub edits: BTreeMap<String, Value>, |
| 81 | pub coverage: MentionCoverage, |
| 82 | pub history_cursor: Option<String>, |
| 83 | pub history_exhausted: bool, |
| 84 | pub latest_cursor: Option<String>, |
| 85 | pub latest_anchor: Option<String>, |
| 86 | pub latest_head: Option<String>, |
| 87 | } |
| 88 | |
| 89 | #[derive(Clone, Default, PartialEq)] |
| 90 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 91 | pub(super) struct SyncState { |
| 92 | pub inbox_revision: u64, |
| 93 | pub checkpoints: BTreeMap<String, Checkpoint>, |
| 94 | } |
| 95 | |
| 96 | /// what a set of rooms holds in the inbox, compared before and after a change so an |
| 97 | /// unchanged inbox costs only the rooms it touched |
| 98 | #[derive(PartialEq)] |
| 99 | pub(super) struct Scope { |
| 100 | rooms: Vec<RoomScope>, |
| 101 | coverage: MentionCoverage, |
| 102 | cursors: ( |
| 103 | Option<String>, |
| 104 | bool, |
| 105 | Option<String>, |
| 106 | Option<String>, |
| 107 | Option<String>, |
| 108 | ), |
| 109 | } |
| 110 | |
| 111 | #[derive(PartialEq)] |
| 112 | struct RoomScope { |
| 113 | entries: Vec<MentionEntry>, |
| 114 | events: Vec<SavedEvent>, |
| 115 | windows: Vec<Window>, |
| 116 | capture_attempts: Vec<u32>, |
| 117 | candidates: Vec<Candidate>, |
| 118 | redacted: Vec<u64>, |
| 119 | edits: Vec<Value>, |
| 120 | keys: Vec<String>, |
| 121 | } |
| 122 | |
| 123 | fn in_room<V: Clone>( |
| 124 | map: &BTreeMap<String, V>, |
| 125 | room: &ruma::RoomId, |
| 126 | keys: &mut Vec<String>, |
| 127 | ) -> Vec<V> { |
| 128 | // no room id holds a newline, so a room's "room\nevent" keys sit below "room\x0b" |
| 129 | map.range(format!("{room}\n")..format!("{room}\u{b}")) |
| 130 | .map(|(key, value)| { |
| 131 | keys.push(key.clone()); |
| 132 | value.clone() |
| 133 | }) |
| 134 | .collect() |
| 135 | } |
| 136 | |
| 137 | #[derive(Clone, Debug, PartialEq)] |
| 138 | #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] |
| 139 | pub(super) struct Checkpoint { |
| 140 | pub pos: String, |
| 141 | pub previous: Option<String>, |
| 142 | } |
| 143 | |
| 144 | pub(super) fn id(key: &MentionKey) -> String { |
| 145 | format!("{}\n{}", key.room_id, key.event_id) |
| 146 | } |
| 147 | |
| 148 | pub(super) fn now() -> u64 { |
| 149 | matrix_sdk::ruma::MilliSecondsSinceUnixEpoch::now() |
| 150 | .get() |
| 151 | .into() |
| 152 | } |
| 153 | |
| 154 | pub(super) fn invalid(message: &'static str) -> Error { |
| 155 | Error::new(ErrorKind::Invalid, message) |
| 156 | } |
| 157 | |
| 158 | #[cfg(feature = "serde")] |
| 159 | pub(super) fn json_len(value: &impl serde::Serialize) -> usize { |
| 160 | struct Count(usize); |
| 161 | impl std::io::Write for Count { |
| 162 | fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> { |
| 163 | self.0 += bytes.len(); |
| 164 | Ok(bytes.len()) |
| 165 | } |
| 166 | fn flush(&mut self) -> std::io::Result<()> { |
| 167 | Ok(()) |
| 168 | } |
| 169 | } |
| 170 | let mut count = Count(0); |
| 171 | let _ = serde_json::to_writer(&mut count, value); |
| 172 | count.0 |
| 173 | } |
| 174 | |
| 175 | #[cfg(not(feature = "serde"))] |
| 176 | pub(super) fn json_len<T>(_: &T) -> usize { |
| 177 | 0 |
| 178 | } |
| 179 | |
| 180 | fn timestamp(value: &Value) -> u64 { |
| 181 | value["origin_server_ts"].as_u64().unwrap_or(0) |
| 182 | } |
| 183 | |
| 184 | impl Data { |
| 185 | pub fn scope<'a>(&self, rooms: impl IntoIterator<Item = &'a ruma::OwnedRoomId>) -> Scope { |
| 186 | Scope { |
| 187 | rooms: rooms |
| 188 | .into_iter() |
| 189 | .map(|room| { |
| 190 | let mut keys = Vec::new(); |
| 191 | RoomScope { |
| 192 | entries: in_room(&self.entries, room, &mut keys), |
| 193 | events: in_room(&self.events, room, &mut keys), |
| 194 | windows: in_room(&self.windows, room, &mut keys), |
| 195 | capture_attempts: in_room(&self.capture_attempts, room, &mut keys), |
| 196 | candidates: in_room(&self.candidates, room, &mut keys), |
| 197 | redacted: in_room(&self.redacted, room, &mut keys), |
| 198 | edits: in_room(&self.edits, room, &mut keys), |
| 199 | keys, |
| 200 | } |
| 201 | }) |
| 202 | .collect(), |
| 203 | coverage: self.coverage.clone(), |
| 204 | cursors: ( |
| 205 | self.history_cursor.clone(), |
| 206 | self.history_exhausted, |
| 207 | self.latest_cursor.clone(), |
| 208 | self.latest_anchor.clone(), |
| 209 | self.latest_head.clone(), |
| 210 | ), |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | /// folds an inbox built in memory while the store was unreachable into the one the |
| 215 | /// store held, keeping the newer side where both have a key |
| 216 | pub fn absorb(&mut self, newer: Data) { |
| 217 | self.revision = self.revision.max(newer.revision); |
| 218 | self.entries.extend(newer.entries); |
| 219 | self.events.extend(newer.events); |
| 220 | self.windows.extend(newer.windows); |
| 221 | self.capture_attempts.extend(newer.capture_attempts); |
| 222 | self.candidates.extend(newer.candidates); |
| 223 | self.redacted.extend(newer.redacted); |
| 224 | self.edits.extend(newer.edits); |
| 225 | for gap in newer.coverage.gaps { |
| 226 | self.gap(gap.room_id, gap.reason, gap.token); |
| 227 | } |
| 228 | self.coverage.earliest_timestamp_ms = self |
| 229 | .coverage |
| 230 | .earliest_timestamp_ms |
| 231 | .into_iter() |
| 232 | .chain(newer.coverage.earliest_timestamp_ms) |
| 233 | .min(); |
| 234 | self.coverage.truncated |= newer.coverage.truncated; |
| 235 | self.history_cursor = newer.history_cursor.or(self.history_cursor.take()); |
| 236 | self.latest_cursor = newer.latest_cursor.or(self.latest_cursor.take()); |
| 237 | self.latest_anchor = newer.latest_anchor.or(self.latest_anchor.take()); |
| 238 | self.latest_head = newer.latest_head.or(self.latest_head.take()); |
| 239 | } |
| 240 | |
| 241 | pub fn gap( |
| 242 | &mut self, |
| 243 | room_id: Option<ruma::OwnedRoomId>, |
| 244 | reason: MentionGapReason, |
| 245 | token: Option<String>, |
| 246 | ) { |
| 247 | if self |
| 248 | .coverage |
| 249 | .gaps |
| 250 | .iter() |
| 251 | .any(|gap| gap.room_id == room_id && gap.reason == reason && gap.token == token) |
| 252 | { |
| 253 | return; |
| 254 | } |
| 255 | self.coverage.gaps.push(MentionGap { |
| 256 | room_id, |
| 257 | reason, |
| 258 | token, |
| 259 | }); |
| 260 | if self.coverage.gaps.len() > 128 { |
| 261 | self.coverage.gaps.remove(0); |
| 262 | self.coverage.truncated = true; |
| 263 | } |
| 264 | } |
| 265 | |
| 266 | pub fn expire_candidates(&mut self, now_ms: u64) -> bool { |
| 267 | let expired: Vec<_> = self |
| 268 | .candidates |
| 269 | .iter() |
| 270 | .filter(|(_, candidate)| { |
| 271 | now_ms.saturating_sub(candidate.first_seen_ms) > CANDIDATE_LIFETIME_MS |
| 272 | }) |
| 273 | .map(|(key, candidate)| (key.clone(), candidate.room_id.clone())) |
| 274 | .collect(); |
| 275 | for (key, room_id) in &expired { |
| 276 | self.candidates.remove(key); |
| 277 | self.gap( |
| 278 | Some(room_id.clone()), |
| 279 | MentionGapReason::CandidateExpired, |
| 280 | None, |
| 281 | ); |
| 282 | } |
| 283 | !expired.is_empty() |
| 284 | } |
| 285 | |
| 286 | pub fn touch_contexts(&mut self, identity: &str) { |
| 287 | for (key, window) in &self.windows { |
| 288 | if window.target.as_deref() == Some(identity) |
| 289 | || window |
| 290 | .before |
| 291 | .iter() |
| 292 | .chain(&window.after) |
| 293 | .any(|key| key == identity) |
| 294 | { |
| 295 | if let Some(entry) = self.entries.get_mut(key) { |
| 296 | entry.revision = self.revision + 1; |
| 297 | } |
| 298 | } |
| 299 | } |
| 300 | } |
| 301 | |
| 302 | pub fn page( |
| 303 | &self, |
| 304 | account: &OwnedUserId, |
| 305 | cursor: Option<&str>, |
| 306 | limit: u16, |
| 307 | hidden: &BTreeSet<OwnedRoomId>, |
| 308 | ) -> Result<MentionsPage, Error> { |
| 309 | let (boundary, last) = if let Some(cursor) = cursor { |
| 310 | let value: Value = |
| 311 | serde_json::from_str(cursor).map_err(|_| invalid("invalid mentions cursor"))?; |
| 312 | if value["account"].as_str() != Some(account.as_str()) { |
| 313 | return Err(invalid("mentions cursor belongs to a different account")); |
| 314 | } |
| 315 | let boundary = value["revision"] |
| 316 | .as_u64() |
| 317 | .ok_or_else(|| invalid("invalid mentions cursor"))?; |
| 318 | let timestamp = value["timestamp"] |
| 319 | .as_u64() |
| 320 | .ok_or_else(|| invalid("invalid mentions cursor"))?; |
| 321 | let key = value["key"] |
| 322 | .as_str() |
| 323 | .ok_or_else(|| invalid("invalid mentions cursor"))? |
| 324 | .to_owned(); |
| 325 | if boundary > self.revision { |
| 326 | return Err(invalid("mentions cursor is newer than this store")); |
| 327 | } |
| 328 | (boundary, Some((timestamp, key))) |
| 329 | } else { |
| 330 | (self.revision, None) |
| 331 | }; |
| 332 | let mut entries: Vec<_> = self |
| 333 | .entries |
| 334 | .iter() |
| 335 | .filter(|(key, entry)| { |
| 336 | !hidden.contains(&entry.key.room_id) |
| 337 | && entry.discovery_revision <= boundary |
| 338 | && last.as_ref().is_none_or(|last| { |
| 339 | (entry.timestamp_ms, key.as_str()) < (last.0, last.1.as_str()) |
| 340 | }) |
| 341 | }) |
| 342 | .map(|(_, entry)| entry.clone()) |
| 343 | .collect(); |
| 344 | entries.sort_by(|a, b| (b.timestamp_ms, id(&b.key)).cmp(&(a.timestamp_ms, id(&a.key)))); |
| 345 | let limit = usize::from(limit.clamp(1, 50)); |
| 346 | let end_reached = entries.len() <= limit; |
| 347 | entries.truncate(limit); |
| 348 | let next_cursor = if end_reached { |
| 349 | None |
| 350 | } else { |
| 351 | entries.last().map(|entry| serde_json::json!({ |
| 352 | "account": account.as_str(), "revision": boundary, "timestamp": entry.timestamp_ms, "key": id(&entry.key) |
| 353 | }).to_string()) |
| 354 | }; |
| 355 | let mut coverage = self.coverage.clone(); |
| 356 | coverage.gaps.retain(|gap| { |
| 357 | gap.room_id |
| 358 | .as_ref() |
| 359 | .is_none_or(|room_id| !hidden.contains(room_id)) |
| 360 | }); |
| 361 | coverage.unresolved_candidates = self |
| 362 | .candidates |
| 363 | .values() |
| 364 | .filter(|candidate| !hidden.contains(&candidate.room_id)) |
| 365 | .count(); |
| 366 | Ok(MentionsPage { |
| 367 | entries, |
| 368 | next_cursor, |
| 369 | end_reached, |
| 370 | coverage, |
| 371 | unread_count: self |
| 372 | .entries |
| 373 | .values() |
| 374 | .filter(|entry| !entry.read && !hidden.contains(&entry.key.room_id)) |
| 375 | .count(), |
| 376 | revision: self.revision, |
| 377 | boundary: MentionsBoundary { |
| 378 | account: account.clone(), |
| 379 | revision: boundary, |
| 380 | }, |
| 381 | }) |
| 382 | } |
| 383 | |
| 384 | pub fn context( |
| 385 | &self, |
| 386 | key: &MentionKey, |
| 387 | ) -> Result<(SavedMentionContext, Vec<MediaRegistration>), Error> { |
| 388 | let entry = self |
| 389 | .entries |
| 390 | .get(&id(key)) |
| 391 | .ok_or_else(|| Error::new(ErrorKind::EventNotFound, "mention was not found"))?; |
| 392 | let mut media = Vec::new(); |
| 393 | let mut item = |key: &String| { |
| 394 | self.events.get(key).map(|event| { |
| 395 | media.extend(event.media.clone()); |
| 396 | event.item.clone() |
| 397 | }) |
| 398 | }; |
| 399 | let context = self.windows.get(&id(key)).and_then(|window| { |
| 400 | let event = window.target.as_ref().and_then(&mut item)?; |
| 401 | Some(EventContext { |
| 402 | events_before: window.before.iter().filter_map(&mut item).collect(), |
| 403 | event, |
| 404 | events_after: window.after.iter().filter_map(&mut item).collect(), |
| 405 | start_token: window.start.clone(), |
| 406 | end_token: window.end.clone(), |
| 407 | }) |
| 408 | }); |
| 409 | Ok(( |
| 410 | SavedMentionContext { |
| 411 | key: key.clone(), |
| 412 | state: entry.capture.clone(), |
| 413 | context, |
| 414 | revision: self.revision, |
| 415 | }, |
| 416 | media, |
| 417 | )) |
| 418 | } |
| 419 | |
| 420 | pub fn collect_events(&mut self) -> usize { |
| 421 | let mut used: BTreeSet<String> = self.entries.keys().cloned().collect(); |
| 422 | for window in self.windows.values() { |
| 423 | used.extend(window.before.iter().chain(&window.after).cloned()); |
| 424 | used.extend(window.target.iter().cloned()); |
| 425 | } |
| 426 | let mut freed = 0; |
| 427 | self.events.retain(|key, event| { |
| 428 | let keep = used.contains(key); |
| 429 | if !keep { |
| 430 | freed += key.len() + json_len(event); |
| 431 | } |
| 432 | keep |
| 433 | }); |
| 434 | self.capture_attempts |
| 435 | .retain(|key, _| self.entries.contains_key(key)); |
| 436 | freed |
| 437 | } |
| 438 | |
| 439 | /// returns roughly how many plaintext JSON bytes the eviction freed |
| 440 | pub fn evict_one(&mut self) -> Option<usize> { |
| 441 | let oldest = |keys: Vec<String>, entries: &BTreeMap<String, MentionEntry>| { |
| 442 | keys.into_iter() |
| 443 | .min_by_key(|key| entries.get(key).map_or(0, |entry| entry.timestamp_ms)) |
| 444 | }; |
| 445 | let freed = if let Some(key) = oldest( |
| 446 | self.windows |
| 447 | .keys() |
| 448 | .filter(|key| self.entries.get(*key).is_some_and(|entry| entry.read)) |
| 449 | .cloned() |
| 450 | .collect(), |
| 451 | &self.entries, |
| 452 | ) |
| 453 | .or_else(|| oldest(self.windows.keys().cloned().collect(), &self.entries)) |
| 454 | { |
| 455 | let window = self.windows.remove(&key); |
| 456 | if let Some(entry) = self.entries.get_mut(&key) { |
| 457 | entry.capture = MentionCaptureState::Unavailable; |
| 458 | entry.revision = self.revision; |
| 459 | } |
| 460 | key.len() + json_len(&window) |
| 461 | } else if let Some(key) = self |
| 462 | .entries |
| 463 | .iter() |
| 464 | .min_by_key(|(_, entry)| (!entry.read, entry.timestamp_ms)) |
| 465 | .map(|(key, _)| key.clone()) |
| 466 | { |
| 467 | key.len() + json_len(&self.entries.remove(&key)) |
| 468 | } else if let Some(key) = self |
| 469 | .candidates |
| 470 | .iter() |
| 471 | .min_by_key(|(_, candidate)| timestamp(&candidate.raw)) |
| 472 | .map(|(key, _)| key.clone()) |
| 473 | { |
| 474 | key.len() + json_len(&self.candidates.remove(&key)) |
| 475 | } else if let Some(key) = self.edits.keys().next().cloned() { |
| 476 | key.len() + json_len(&self.edits.remove(&key)) |
| 477 | } else { |
| 478 | let key = self.redacted.keys().next().cloned()?; |
| 479 | self.redacted.remove(&key); |
| 480 | key.len() + 24 |
| 481 | }; |
| 482 | self.coverage.truncated = true; |
| 483 | Some(freed + self.collect_events()) |
| 484 | } |
| 485 | |
| 486 | pub fn bound(&mut self, entries: usize, events: usize, candidates: usize) { |
| 487 | while self.entries.len() > entries { |
| 488 | if let Some(key) = self |
| 489 | .entries |
| 490 | .iter() |
| 491 | .min_by_key(|(_, entry)| (!entry.read, entry.timestamp_ms)) |
| 492 | .map(|(key, _)| key.clone()) |
| 493 | { |
| 494 | self.entries.remove(&key); |
| 495 | self.windows.remove(&key); |
| 496 | self.coverage.truncated = true; |
| 497 | } |
| 498 | } |
| 499 | self.collect_events(); |
| 500 | while self.events.len() > events && self.evict_one().is_some() {} |
| 501 | if self.candidates.len() > candidates { |
| 502 | let mut oldest: Vec<_> = self |
| 503 | .candidates |
| 504 | .iter() |
| 505 | .map(|(key, candidate)| (timestamp(&candidate.raw), key.clone())) |
| 506 | .collect(); |
| 507 | oldest.sort(); |
| 508 | for (_, key) in oldest.into_iter().take(self.candidates.len() - candidates) { |
| 509 | self.candidates.remove(&key); |
| 510 | } |
| 511 | self.coverage.truncated = true; |
| 512 | } |
| 513 | for map in [&mut self.redacted] { |
| 514 | while map.len() > entries + candidates { |
| 515 | let key = map |
| 516 | .iter() |
| 517 | .min_by_key(|(_, timestamp)| *timestamp) |
| 518 | .map(|(key, _)| key.clone()) |
| 519 | .unwrap(); |
| 520 | map.remove(&key); |
| 521 | self.coverage.truncated = true; |
| 522 | } |
| 523 | } |
| 524 | while self.edits.len() > candidates { |
| 525 | let key = self.edits.keys().next().cloned().unwrap(); |
| 526 | self.edits.remove(&key); |
| 527 | self.coverage.truncated = true; |
| 528 | } |
| 529 | self.coverage.unresolved_candidates = self.candidates.len(); |
| 530 | } |
| 531 | } |