From aa49b111edf3ee9eaccc43b689bacdb0207799af Mon Sep 17 00:00:00 2001
From: Charles Hall <charles@computer.surgery>
Date: Sun, 3 Nov 2024 11:34:15 -0800
Subject: [PATCH] allow selecting a subset of rooms

This can make debugging a specific room much easier/faster.
---
 src/cli.rs                           |  7 +++++++
 src/cli/get_room_states.rs           |  4 ++--
 src/cli/get_room_states/cache.rs     | 10 +++++++---
 src/cli/get_room_states/recompute.rs | 24 ++++++++++++++++--------
 4 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index 110359ae..874b2db3 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -6,6 +6,7 @@
 use std::path::PathBuf;
 
 use clap::{Parser, Subcommand};
+use ruma::OwnedRoomId;
 
 use crate::{
     config::{default_tracing_filter, EnvFilterClone, LogFormat},
@@ -107,6 +108,12 @@ pub(crate) struct GetRoomStatesArgs {
     #[clap(long, short)]
     pub(crate) recompute: bool,
 
+    /// Limit the output to a subset of rooms by their IDs
+    ///
+    /// This option can be specified multiple times.
+    #[clap(long, short)]
+    pub(crate) select: Vec<OwnedRoomId>,
+
     #[clap(flatten)]
     observability: ObservabilityArgs,
 }
diff --git a/src/cli/get_room_states.rs b/src/cli/get_room_states.rs
index e82a2f34..3a826bd8 100644
--- a/src/cli/get_room_states.rs
+++ b/src/cli/get_room_states.rs
@@ -65,9 +65,9 @@ pub(crate) async fn run(
     db.apply_migrations().await.map_err(Error::Database)?;
 
     let room_states = if args.recompute {
-        recompute::get_room_states(db)
+        recompute::get_room_states(db, &args.select)
     } else {
-        cache::get_room_states().await
+        cache::get_room_states(&args.select).await
     };
 
     serde_json::to_writer(std::io::stdout(), &room_states)?;
diff --git a/src/cli/get_room_states/cache.rs b/src/cli/get_room_states/cache.rs
index ba2f8ccf..23485834 100644
--- a/src/cli/get_room_states/cache.rs
+++ b/src/cli/get_room_states/cache.rs
@@ -8,12 +8,16 @@ use tracing as t;
 use super::StateEvent;
 use crate::{services, PduEvent};
 
-/// Get the state of all rooms
-#[t::instrument]
-pub(crate) async fn get_room_states() -> Vec<StateEvent> {
+/// Get the state of all rooms, or the selected subset
+#[t::instrument(skip(select))]
+pub(crate) async fn get_room_states(select: &[OwnedRoomId]) -> Vec<StateEvent> {
     let mut serializable_state = Vec::new();
 
     for room_id in services().rooms.metadata.iter_ids().filter_map(Result::ok) {
+        if !select.is_empty() && !select.contains(&room_id) {
+            continue;
+        }
+
         let Some(state) = get_room_state(room_id).await else {
             continue;
         };
diff --git a/src/cli/get_room_states/recompute.rs b/src/cli/get_room_states/recompute.rs
index c5f19b54..480aed89 100644
--- a/src/cli/get_room_states/recompute.rs
+++ b/src/cli/get_room_states/recompute.rs
@@ -18,7 +18,7 @@ use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use ruma::{
     events::{StateEventType, TimelineEventType},
     state_res::StateMap,
-    EventId, RoomId, RoomVersionId, UserId,
+    EventId, OwnedRoomId, RoomId, RoomVersionId, UserId,
 };
 use serde::Deserialize;
 use tracing as t;
@@ -64,9 +64,12 @@ struct StateResolutionEdges {
     auth_events: Vec<Arc<EventId>>,
 }
 
-/// Get the state of all rooms
-pub(crate) fn get_room_states(db: &KeyValueDatabase) -> Vec<StateEvent> {
-    let state_resolution_edges = get_state_resolution_edges(db);
+/// Get the state of all rooms, or the selected subset
+pub(crate) fn get_room_states(
+    db: &KeyValueDatabase,
+    select: &[OwnedRoomId],
+) -> Vec<StateEvent> {
+    let state_resolution_edges = get_state_resolution_edges(db, select);
     let graphs = get_state_event_graphs(db, state_resolution_edges);
 
     let states = resolve_room_states(db, graphs);
@@ -376,10 +379,11 @@ fn get_either_pdu_by_event_id(
 }
 
 /// Get the edges in the graph of state events for state resolution
-#[t::instrument(skip(db))]
-fn get_state_resolution_edges(
-    db: &KeyValueDatabase,
-) -> impl Iterator<Item = StateResolutionEdges> + '_ {
+#[t::instrument(skip(db, select))]
+fn get_state_resolution_edges<'a>(
+    db: &'a KeyValueDatabase,
+    select: &'a [OwnedRoomId],
+) -> impl Iterator<Item = StateResolutionEdges> + 'a {
     let filter_map = |key: Vec<u8>, value: Vec<u8>, map: &str| {
         let Ok(pdu) = serde_json::from_slice::<PduEvent>(&value) else {
             t::error!(
@@ -390,6 +394,10 @@ fn get_state_resolution_edges(
             return None;
         };
 
+        if !select.is_empty() && !select.contains(&pdu.room_id) {
+            return None;
+        }
+
         let Some(state_key) = pdu.state_key else {
             // Filter out non-state events
             return None;