1use std::{borrow::Borrow, collections::HashMap, mem::take, sync::Arc};
2
3use futures::{
4 FutureExt, StreamExt, TryFutureExt,
5 future::{join, try_join},
6};
7use ruma::{
8 EventId, OwnedEventId, OwnedRoomId, RoomId, RoomVersionId,
9 events::{StateEventType, TimelineEventType},
10 room_version_rules::RoomVersionRules,
11};
12use tracing::Span;
13use tuwunel_core::{
14 Result, debug, debug_warn, defer, err, implement,
15 matrix::{
16 Event, PduEvent, StateKey,
17 pdu::PrevEvents,
18 room_version::{self, from_create_event},
19 },
20 trace,
21 utils::stream::{BroadbandExt, IterStream, ReadyExt, WidebandExt},
22 warn,
23};
24
25use crate::rooms::{
26 short::{ShortStateHash, ShortStateKey},
27 state_compressor::CompressedState,
28 state_res::auth_check,
29};
30
31type StateIds = HashMap<ShortStateKey, OwnedEventId>;
33
34#[derive(Debug)]
36pub struct LocalBuildReport {
37 pub state_len: Option<usize>,
38 pub visited: usize,
39 pub forks: usize,
40 pub gate_drops: usize,
41 pub memo_hits: usize,
42 pub fallback: Option<String>,
43}
44
45#[derive(Clone, Copy, Eq, PartialEq)]
48pub(super) enum WalkMode {
49 Active,
50 Shadow,
51}
52
53struct Walk<'a> {
55 room_id: &'a RoomId,
56 room_version: &'a RoomVersionId,
57 room_rules: RoomVersionRules,
58 create_event_id: &'a EventId,
59 mode: WalkMode,
60 max_nodes: usize,
61 top_prevs: PrevEvents,
62 class: HashMap<OwnedEventId, Class>,
63 nodes: Vec<Node>,
64 order: Vec<usize>,
65 frontier: HashMap<OwnedEventId, usize>,
66 resolved: HashMap<OwnedEventId, Arc<StateIds>>,
67 live_entries: usize,
68 peak_entries: usize,
69 forks: usize,
70 gate_drops: usize,
71 memo_hits: usize,
72 fallback: Option<Fallback>,
73}
74
75struct Node {
77 pdu: PduEvent,
78 consumers: usize,
79}
80
81#[derive(Clone, Copy)]
83enum Class {
84 Committed(ShortStateHash),
86
87 Memoized,
89
90 Held(usize),
92}
93
94#[derive(Clone, Copy)]
96enum Fallback {
97 Absent,
98 Ceiling,
99 AuthMissing,
100 AllCommitted,
101 Entries,
102 Canary,
103 CreateMismatch,
104 Error,
105}
106
107const MAX_LIVE_ENTRIES: usize = 1 << 19;
112
113const DIVERGENCE_SAMPLE: usize = 16;
115
116#[implement(super::Service)]
123pub(super) async fn state_at_incoming_local<Pdu>(
124 &self,
125 room_id: &RoomId,
126 incoming_pdu: &Pdu,
127 room_version: &RoomVersionId,
128 create_event_id: &EventId,
129 mode: WalkMode,
130) -> Result<Option<StateIds>>
131where
132 Pdu: Event,
133{
134 let top_prevs = incoming_pdu
135 .prev_events()
136 .map(ToOwned::to_owned)
137 .collect();
138
139 let services = self.services.clone();
140 let room_id = room_id.to_owned();
141 let room_version = room_version.clone();
142 let create_event_id = create_event_id.to_owned();
143 let parent = Span::current();
144
145 let task = self.services.server.runtime().spawn(async move {
146 services
147 .event_handler
148 .walk_task(room_id, room_version, create_event_id, mode, top_prevs, parent)
149 .await
150 });
151
152 let abort = task.abort_handle();
154 defer! {{ abort.abort(); }};
155
156 task.await.unwrap_or_else(|error| {
157 debug_warn!(
158 %error,
159 "Local state build task failed; falling back to federation fetch.",
160 );
161
162 Ok(None)
163 })
164}
165
166#[implement(super::Service)]
170#[tracing::instrument(name = "local", level = "debug", parent = &parent, skip_all)]
171async fn walk_task(
172 &self,
173 room_id: OwnedRoomId,
174 room_version: RoomVersionId,
175 create_event_id: OwnedEventId,
176 mode: WalkMode,
177 top_prevs: PrevEvents,
178 parent: Span,
179) -> Result<Option<StateIds>> {
180 let max_nodes = self
181 .services
182 .server
183 .config
184 .resolve_state_locally_max;
185
186 let mut walk =
187 Walk::new(&room_id, &room_version, &create_event_id, mode, max_nodes, top_prevs)?;
188
189 let state = self.walk_state(&mut walk).await?;
190
191 debug!(
192 visited = walk.nodes.len(),
193 forks = walk.forks,
194 gate_drops = walk.gate_drops,
195 memo_hits = walk.memo_hits,
196 live_entries_peak = walk.peak_entries,
197 outcome = walk.fallback.map_or("resolved", Fallback::name),
198 "Local state build finished.",
199 );
200
201 if let Some(fallback) = walk.fallback {
202 debug_warn!(
203 reason = fallback.name(),
204 "Local state build falling back to federation fetch.",
205 );
206 }
207
208 Ok(state)
209}
210
211#[implement(super::Service)]
214pub async fn local_state_report(&self, event_id: &EventId) -> Result<LocalBuildReport> {
215 let pdu = self.services.timeline.get_pdu(event_id).await?;
216
217 let create_event = self
218 .services
219 .state_accessor
220 .room_state_get(pdu.room_id(), &StateEventType::RoomCreate, "")
221 .await?;
222
223 let room_version = from_create_event(&create_event)?;
224 let max_nodes = self
225 .services
226 .server
227 .config
228 .resolve_state_locally_max;
229
230 let top_prevs = pdu.prev_events().map(ToOwned::to_owned).collect();
231
232 let mut walk = Walk::new(
233 pdu.room_id(),
234 &room_version,
235 create_event.event_id(),
236 WalkMode::Shadow,
237 max_nodes,
238 top_prevs,
239 )?;
240
241 let state = self.walk_state(&mut walk).await?;
242
243 Ok(LocalBuildReport {
244 state_len: state.map(|state| state.len()),
245 visited: walk.nodes.len(),
246 forks: walk.forks,
247 gate_drops: walk.gate_drops,
248 memo_hits: walk.memo_hits,
249 fallback: walk
250 .fallback
251 .map(|fallback| fallback.name().to_owned()),
252 })
253}
254
255pub(super) fn compare_shadow(
258 room_id: &RoomId,
259 event_id: &EventId,
260 local: &StateIds,
261 fetched: &StateIds,
262) {
263 let only_local: Vec<ShortStateKey> = diverging(local, fetched).collect();
264 let only_fetch: Vec<ShortStateKey> = diverging(fetched, local).collect();
265
266 if only_local.is_empty() && only_fetch.is_empty() {
267 debug!(%room_id, %event_id, "Shadow local state build matches fetched state.");
268 return;
269 }
270
271 warn!(
272 %room_id,
273 %event_id,
274 only_local = only_local.len(),
275 only_fetch = only_fetch.len(),
276 "Shadow local state build diverges from fetched state.",
277 );
278
279 let sample: Vec<_> = only_local
280 .iter()
281 .chain(only_fetch.iter())
282 .copied()
283 .take(DIVERGENCE_SAMPLE)
284 .collect();
285
286 debug!(?sample, "Diverging shortstatekeys.");
287}
288
289fn diverging<'a>(a: &'a StateIds, b: &'a StateIds) -> impl Iterator<Item = ShortStateKey> + 'a {
291 a.iter()
292 .filter(|&(shortstatekey, event_id)| b.get(shortstatekey) != Some(event_id))
293 .map(|(&shortstatekey, _)| shortstatekey)
294}
295
296#[implement(super::Service)]
299async fn walk_state(&self, walk: &mut Walk<'_>) -> Result<Option<StateIds>> {
300 self.walk_discover(walk).await?;
301
302 if walk.fallback.is_some() {
303 return Ok(None);
304 }
305
306 self.walk_build(walk).await
307}
308
309#[implement(super::Service)]
313async fn walk_discover(&self, walk: &mut Walk<'_>) -> Result {
314 let mut stack: Vec<(OwnedEventId, bool)> = walk
315 .top_prevs
316 .iter()
317 .map(|prev| (prev.clone(), false))
318 .collect();
319
320 while let Some((event_id, expanded)) = stack.pop() {
321 self.services.server.check_running()?;
322
323 if expanded {
324 let Some(Class::Held(index)) = walk.class.get(&event_id).copied() else {
326 debug_assert!(false, "expanded stack entries are held nodes");
327 walk.fallback = Some(Fallback::Error);
328 return Ok(());
329 };
330
331 walk.order.push(index);
332 continue;
333 }
334
335 if walk.class.contains_key(&event_id) {
336 continue;
337 }
338
339 if let Ok(shortstatehash) = self
340 .services
341 .state
342 .pdu_shortstatehash(&event_id)
343 .await
344 {
345 walk.class
346 .insert(event_id, Class::Committed(shortstatehash));
347
348 continue;
349 }
350
351 if self
352 .db
353 .eventid_resolvedstate
354 .exists(&event_id)
355 .await
356 .is_ok()
357 {
358 walk.class.insert(event_id, Class::Memoized);
359 continue;
360 }
361
362 let Ok(pdu) = self.services.timeline.get_pdu(&event_id).await else {
363 trace!(%event_id, "Ancestor is not held locally.");
364 walk.fallback = Some(Fallback::Absent);
365 return Ok(());
366 };
367
368 if walk.nodes.len() >= walk.max_nodes {
369 walk.fallback = Some(Fallback::Ceiling);
370 return Ok(());
371 }
372
373 if pdu.prev_events().next().is_none() {
374 debug_warn!(%event_id, "Held uncommitted ancestor has no prev events.");
375 walk.fallback = Some(Fallback::Error);
376 return Ok(());
377 }
378
379 if !self.walk_auth_present(walk, &pdu).await {
380 walk.fallback = Some(Fallback::AuthMissing);
381 return Ok(());
382 }
383
384 walk.class
385 .insert(event_id.clone(), Class::Held(walk.nodes.len()));
386
387 stack.push((event_id, true));
388 stack.extend(
389 pdu.prev_events()
390 .map(|prev| (prev.to_owned(), false)),
391 );
392 walk.nodes.push(Node { pdu, consumers: 0 });
393 }
394
395 if walk.nodes.is_empty() {
396 walk.fallback = Some(Fallback::AllCommitted);
399 return Ok(());
400 }
401
402 walk.count_consumers();
403
404 Ok(())
405}
406
407#[implement(super::Service)]
411async fn walk_auth_present(&self, walk: &Walk<'_>, pdu: &PduEvent) -> bool {
412 let is_hydra = !walk
413 .room_rules
414 .event_format
415 .allow_room_create_in_auth_events;
416
417 let not_create = *pdu.kind() != TimelineEventType::RoomCreate;
418 let hydra_create_id = (not_create && is_hydra)
419 .then(|| pdu.room_id().as_event_id().ok())
420 .flatten();
421
422 pdu.auth_events()
423 .chain(hydra_create_id.as_deref())
424 .stream()
425 .all(|auth_id| self.services.timeline.pdu_exists(auth_id))
426 .await
427}
428
429#[implement(super::Service)]
432async fn walk_build(&self, walk: &mut Walk<'_>) -> Result<Option<StateIds>> {
433 let order = take(&mut walk.order);
434 for index in order {
435 self.services.server.check_running()?;
436
437 if !self.walk_node(walk, index).await {
438 return Ok(None);
439 }
440 }
441
442 let top_prevs = take(&mut walk.top_prevs);
443 let state = match top_prevs.as_slice() {
444 | [prev] => self.state_after(walk, prev).await,
445 | _ => self.fork_resolve(walk, &top_prevs, None).await,
446 };
447
448 let Some(state) = state else {
449 return Ok(None);
450 };
451
452 let create_entry = self
455 .services
456 .short
457 .get_shortstatekey(&StateEventType::RoomCreate, "")
458 .await
459 .ok()
460 .and_then(|shortstatekey| state.get(&shortstatekey))
461 .map(AsRef::as_ref);
462
463 if state.is_empty() || create_entry != Some(walk.create_event_id) {
464 walk.fallback = Some(Fallback::CreateMismatch);
465 return Ok(None);
466 }
467
468 walk.resolved.clear();
469
470 let state = Arc::try_unwrap(state).unwrap_or_else(|state| (*state).clone());
471
472 Ok(Some(state))
473}
474
475#[implement(super::Service)]
478async fn walk_node(&self, walk: &mut Walk<'_>, index: usize) -> bool {
479 let node = &walk.nodes[index];
480 let event_id = node.pdu.event_id().to_owned();
481 let prevs: PrevEvents = node
482 .pdu
483 .prev_events()
484 .map(ToOwned::to_owned)
485 .collect();
486
487 let before = match prevs.as_slice() {
488 | [prev] => self.state_after(walk, prev).await,
489 | _ =>
490 self.fork_resolve(walk, &prevs, Some(&event_id))
491 .await,
492 };
493
494 let Some(before) = before else {
495 return false;
496 };
497
498 let after = match walk.nodes[index].pdu.state_key() {
499 | None => before,
500 | Some(_) =>
501 self.gated_fold(
502 &walk.room_rules,
503 &mut walk.gate_drops,
504 &walk.nodes[index].pdu,
505 &before,
506 )
507 .await,
508 };
509
510 if !walk.retain(event_id, after) {
511 return false;
512 }
513
514 walk.release(&prevs);
515
516 true
517}
518
519#[implement(super::Service)]
522async fn state_after(&self, walk: &mut Walk<'_>, event_id: &EventId) -> Option<Arc<StateIds>> {
523 if let Some(state) = walk.resolved.get(event_id) {
524 return Some(state.clone());
525 }
526
527 let state = match walk.class.get(event_id).copied() {
528 | Some(Class::Committed(shortstatehash)) =>
529 self.committed_state_after(walk, event_id, shortstatehash)
530 .await,
531 | Some(Class::Memoized) => self.memoized_state_after(walk, event_id).await,
532 | Some(Class::Held(_)) | None => {
533 debug_assert!(false, "held nodes resolve before their consumers");
534 walk.fallback = Some(Fallback::Error);
535 None
536 },
537 }?;
538
539 walk.retain(event_id.to_owned(), state.clone())
540 .then_some(state)
541}
542
543#[implement(super::Service)]
547async fn committed_state_after(
548 &self,
549 walk: &mut Walk<'_>,
550 event_id: &EventId,
551 shortstatehash: ShortStateHash,
552) -> Option<Arc<StateIds>> {
553 let pdu = self.services.timeline.get_pdu(event_id);
554
555 let state = self
556 .services
557 .state_accessor
558 .state_full_ids(shortstatehash)
559 .collect::<StateIds>()
560 .map(Ok);
561
562 let Ok((pdu, mut state)) = try_join(pdu, state)
563 .inspect_err(|e| debug_warn!(%event_id, %e, "Failed loading committed state."))
564 .await
565 else {
566 walk.fallback = Some(Fallback::Error);
567 return None;
568 };
569
570 if let Some(state_key) = pdu.state_key() {
571 let event_type = pdu.event_type().to_cow_str().into();
572 let shortstatekey = self
573 .services
574 .short
575 .get_or_create_shortstatekey(&event_type, state_key)
576 .await;
577
578 state.insert(shortstatekey, event_id.to_owned());
579 }
580
581 Some(Arc::new(state))
582}
583
584#[implement(super::Service)]
587async fn memoized_state_after(
588 &self,
589 walk: &mut Walk<'_>,
590 event_id: &EventId,
591) -> Option<Arc<StateIds>> {
592 walk.memo_hits = walk.memo_hits.saturating_add(1);
593
594 let state = self.cached_resolved_state(event_id);
595
596 let pdu = self
597 .services
598 .timeline
599 .get_pdu(event_id)
600 .inspect_err(|e| debug_warn!(%event_id, %e, "Failed loading memoized event."));
601
602 let (state, pdu) = join(state, pdu).await;
603
604 let Some(state) = state else {
605 walk.fallback = Some(Fallback::Canary);
606 return None;
607 };
608
609 let Ok(pdu) = pdu else {
610 walk.fallback = Some(Fallback::Error);
611 return None;
612 };
613
614 let before = Arc::new(state);
615 if pdu.state_key().is_none() {
616 return Some(before);
617 }
618
619 let after = self
620 .gated_fold(&walk.room_rules, &mut walk.gate_drops, &pdu, &before)
621 .await;
622
623 Some(after)
624}
625
626#[implement(super::Service)]
631async fn gated_fold(
632 &self,
633 room_rules: &RoomVersionRules,
634 gate_drops: &mut usize,
635 pdu: &PduEvent,
636 before: &Arc<StateIds>,
637) -> Arc<StateIds> {
638 let state_fetch = async |k: StateEventType, s: StateKey| {
639 let shortstatekey = self
640 .services
641 .short
642 .get_shortstatekey(&k, s.as_str())
643 .await?;
644
645 let event_id = before
646 .get(&shortstatekey)
647 .ok_or_else(|| err!(Request(NotFound("Not in state before event."))))?;
648
649 self.services.timeline.get_pdu(event_id).await
650 };
651
652 let event_fetch = async |event_id: OwnedEventId| self.event_fetch(&event_id).await;
653
654 if let Err(e) = auth_check(room_rules, pdu, &event_fetch, &state_fetch).await {
655 debug!(event_id = %pdu.event_id(), %e, "Auth gate rejected fold.");
656 *gate_drops = gate_drops.saturating_add(1);
657 return before.clone();
658 }
659
660 let state_key = pdu.state_key().expect("only state events fold");
661
662 let event_type = pdu.event_type().to_cow_str().into();
663 let shortstatekey = self
664 .services
665 .short
666 .get_or_create_shortstatekey(&event_type, state_key)
667 .await;
668
669 let mut state = StateIds::clone(before);
670 state.insert(shortstatekey, pdu.event_id().to_owned());
671
672 Arc::new(state)
673}
674
675#[implement(super::Service)]
679async fn fork_resolve(
680 &self,
681 walk: &mut Walk<'_>,
682 prevs: &[OwnedEventId],
683 memo_event_id: Option<&EventId>,
684) -> Option<Arc<StateIds>> {
685 walk.forks = walk.forks.saturating_add(1);
686
687 let mut afters = Vec::with_capacity(prevs.len());
689 for prev in prevs {
690 afters.push(self.state_after(walk, prev).await?);
691 }
692
693 let (room_id, room_version) = (walk.room_id, walk.room_version);
694 let fork_states = afters.iter().stream().wide_then(|after| {
695 let state = after
696 .iter()
697 .map(|(shortstatekey, event_id)| (*shortstatekey, event_id));
698
699 self.fork_state(state)
700 });
701
702 let auth_chains = prevs
703 .iter()
704 .zip(&afters)
705 .stream()
706 .wide_then(|(prev_event, after)| {
707 self.fork_chain(room_id, room_version, after.values().map(Borrow::borrow))
708 .inspect_err(move |e| {
709 debug_warn!(%prev_event, %e, "Skipping failed fork auth chain.");
710 })
711 })
712 .ready_filter_map(Result::ok);
713
714 let Ok(resolved) = self
715 .state_resolution(room_id, room_version, fork_states, auth_chains)
716 .await
717 else {
718 walk.fallback = Some(Fallback::Error);
719 return None;
720 };
721
722 let state: StateIds = resolved
723 .into_iter()
724 .stream()
725 .broad_then(async |((event_type, state_key), event_id)| {
726 self.services
727 .short
728 .get_or_create_shortstatekey(&event_type, &state_key)
729 .map(move |shortstatekey| (shortstatekey, event_id))
730 .await
731 })
732 .collect()
733 .await;
734
735 if let Some(event_id) = memo_event_id.filter(|_| walk.mode == WalkMode::Active) {
736 let compressed: Arc<CompressedState> = self
737 .services
738 .state_compressor
739 .compress_state_events(
740 state
741 .iter()
742 .map(|(shortstatekey, event_id)| (shortstatekey, event_id.borrow())),
743 )
744 .collect()
745 .map(Arc::new)
746 .await;
747
748 self.cache_resolved_state(walk.room_id, event_id, compressed)
749 .await;
750 }
751
752 Some(Arc::new(state))
753}
754
755impl<'a> Walk<'a> {
756 fn new(
757 room_id: &'a RoomId,
758 room_version: &'a RoomVersionId,
759 create_event_id: &'a EventId,
760 mode: WalkMode,
761 max_nodes: usize,
762 top_prevs: PrevEvents,
763 ) -> Result<Self> {
764 Ok(Self {
765 room_id,
766 room_version,
767 room_rules: room_version::rules(room_version)?,
768 create_event_id,
769 mode,
770 max_nodes,
771 top_prevs,
772 class: HashMap::new(),
773 nodes: Vec::new(),
774 order: Vec::new(),
775 frontier: HashMap::new(),
776 resolved: HashMap::new(),
777 live_entries: 0,
778 peak_entries: 0,
779 forks: 0,
780 gate_drops: 0,
781 memo_hits: 0,
782 fallback: None,
783 })
784 }
785
786 fn count_consumers(&mut self) {
789 let mut held = vec![0_usize; self.nodes.len()];
790
791 let edges = self
792 .nodes
793 .iter()
794 .flat_map(|node| node.pdu.prev_events())
795 .chain(self.top_prevs.iter().map(AsRef::as_ref));
796
797 for prev in edges {
798 match self.class.get(prev).copied() {
799 | Some(Class::Held(index)) => held[index] = held[index].saturating_add(1),
800 | Some(_) => {
801 let consumers = self.frontier.entry(prev.to_owned()).or_default();
802
803 *consumers = consumers.saturating_add(1);
804 },
805 | None => debug_assert!(false, "every walk edge is classified"),
806 }
807 }
808
809 for (node, consumers) in self.nodes.iter_mut().zip(held) {
810 node.consumers = consumers;
811 }
812 }
813
814 fn retain(&mut self, event_id: OwnedEventId, state: Arc<StateIds>) -> bool {
818 let live_entries = self.live_entries.saturating_add(state.len());
819 if live_entries > MAX_LIVE_ENTRIES {
820 self.fallback = Some(Fallback::Entries);
821 return false;
822 }
823
824 self.live_entries = live_entries;
825 self.peak_entries = self.peak_entries.max(live_entries);
826 self.resolved.insert(event_id, state);
827
828 true
829 }
830
831 fn release(&mut self, prevs: &[OwnedEventId]) {
834 for prev in prevs {
835 let remaining = match self.class.get(prev).copied() {
836 | Some(Class::Held(index)) => {
837 let node = &mut self.nodes[index];
838 node.consumers = node.consumers.saturating_sub(1);
839 node.consumers
840 },
841 | _ => {
842 let Some(consumers) = self.frontier.get_mut(prev) else {
843 continue;
844 };
845
846 *consumers = consumers.saturating_sub(1);
847 *consumers
848 },
849 };
850
851 if remaining == 0
852 && let Some(state) = self.resolved.remove(prev)
853 {
854 self.live_entries = self.live_entries.saturating_sub(state.len());
855 }
856 }
857 }
858}
859
860impl Fallback {
861 fn name(self) -> &'static str {
862 match self {
863 | Self::Absent => "absent",
864 | Self::Ceiling => "ceiling",
865 | Self::AuthMissing => "auth_missing",
866 | Self::AllCommitted => "all_committed",
867 | Self::Entries => "entries",
868 | Self::Canary => "canary",
869 | Self::CreateMismatch => "create_mismatch",
870 | Self::Error => "error",
871 }
872 }
873}