| 1 | use std::{ |
| 2 | collections::{HashMap, HashSet, VecDeque}, |
| 3 | sync::Arc, |
| 4 | }; |
| 5 | |
| 6 | use matrix_sdk::{ |
| 7 | Client, Room, |
| 8 | ruma::{ |
| 9 | Int, |
| 10 | api::client::room::{ |
| 11 | Visibility as SdkVisibility, |
| 12 | create_room::v3::{Request as CreateRoomRequest, RoomPreset as SdkPreset}, |
| 13 | }, |
| 14 | events::{ |
| 15 | AnyGlobalAccountDataEventContent, GlobalAccountDataEventType, MessageLikeEventType, |
| 16 | StateEventType, room::power_levels::RoomPowerLevelsEventContent, |
| 17 | }, |
| 18 | serde::Raw, |
| 19 | }, |
| 20 | }; |
| 21 | use ruma::{OwnedRoomId, OwnedUserId}; |
| 22 | use tokio::sync::mpsc; |
| 23 | |
| 24 | use super::commands::{Job, guarded}; |
| 25 | use super::executor::TaskSet; |
| 26 | use crate::{ |
| 27 | CommandOutput, CreateRoomOptions, Error, ErrorKind, PowerLevelChanges, RoomCommand, RoomPreset, |
| 28 | RoomVisibility, |
| 29 | }; |
| 30 | |
| 31 | /// state changes that must not race each other. same room, same alias, same dm |
| 32 | /// partner or same account data type run in order, everything else runs at once |
| 33 | pub(super) enum Mutation { |
| 34 | Room(RoomCommand), |
| 35 | GlobalAccountData { |
| 36 | event_type: GlobalAccountDataEventType, |
| 37 | content: Raw<AnyGlobalAccountDataEventContent>, |
| 38 | }, |
| 39 | } |
| 40 | |
| 41 | pub(super) type MutationJob = Job<Mutation>; |
| 42 | |
| 43 | #[derive(Clone, Debug, Eq, Hash, PartialEq)] |
| 44 | enum Key { |
| 45 | Room(OwnedRoomId), |
| 46 | Alias(String), |
| 47 | Direct(OwnedUserId), |
| 48 | AccountData(String), |
| 49 | Independent(u64), |
| 50 | } |
| 51 | |
| 52 | pub(super) async fn run_scheduler( |
| 53 | client: Client, |
| 54 | mut jobs: mpsc::Receiver<MutationJob>, |
| 55 | queue_limit: usize, |
| 56 | ) { |
| 57 | let fresh_dms = FreshDms::default(); |
| 58 | let mut waiting = HashMap::<Key, VecDeque<MutationJob>>::new(); |
| 59 | let mut running = HashSet::<Key>::new(); |
| 60 | let mut tasks = TaskSet::<Key>::new(); |
| 61 | let mut queued = 0_usize; |
| 62 | let mut independent = 0_u64; |
| 63 | let mut receiving = true; |
| 64 | while receiving || !tasks.is_empty() { |
| 65 | tokio::select! { |
| 66 | biased; |
| 67 | Some(finished) = tasks.join_next(), if !tasks.is_empty() => { |
| 68 | // a job's own panic is caught inside it, so this is only ever a key |
| 69 | let Ok(key) = finished else { continue }; |
| 70 | let next = waiting.get_mut(&key).and_then(VecDeque::pop_front); |
| 71 | if waiting.get(&key).is_some_and(VecDeque::is_empty) { |
| 72 | waiting.remove(&key); |
| 73 | } |
| 74 | match next { |
| 75 | Some(job) => { |
| 76 | queued -= 1; |
| 77 | spawn(&mut tasks, &client, &fresh_dms, key, job); |
| 78 | } |
| 79 | None => { |
| 80 | running.remove(&key); |
| 81 | } |
| 82 | } |
| 83 | } |
| 84 | job = jobs.recv(), if receiving => { |
| 85 | let Some(job) = job else { |
| 86 | receiving = false; |
| 87 | continue; |
| 88 | }; |
| 89 | let key = key(&job.command, client.user_id().map(ruma::UserId::server_name)).unwrap_or_else(|| { |
| 90 | independent += 1; |
| 91 | Key::Independent(independent) |
| 92 | }); |
| 93 | if running.insert(key.clone()) { |
| 94 | spawn(&mut tasks, &client, &fresh_dms, key, job); |
| 95 | } else if queued < queue_limit { |
| 96 | waiting.entry(key).or_default().push_back(job); |
| 97 | queued += 1; |
| 98 | } else { |
| 99 | job.reply.done(Err(Error::new( |
| 100 | ErrorKind::QueueFull, |
| 101 | "room command queue is full", |
| 102 | ))); |
| 103 | } |
| 104 | } |
| 105 | } |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | /// `server` is ours, since a new alias lives there |
| 110 | fn key(mutation: &Mutation, server: Option<&ruma::ServerName>) -> Option<Key> { |
| 111 | let room_id = match mutation { |
| 112 | Mutation::GlobalAccountData { event_type, .. } => { |
| 113 | return Some(Key::AccountData(event_type.to_string())); |
| 114 | } |
| 115 | Mutation::Room(RoomCommand::Join { room, .. }) => { |
| 116 | return Some(match OwnedRoomId::try_from(room.clone()) { |
| 117 | Ok(room_id) => Key::Room(room_id), |
| 118 | Err(alias) => Key::Alias(alias.to_string()), |
| 119 | }); |
| 120 | } |
| 121 | Mutation::Room(RoomCommand::CreateDirect(user_id)) => { |
| 122 | return Some(Key::Direct(user_id.clone())); |
| 123 | } |
| 124 | Mutation::Room(RoomCommand::Create(options)) => { |
| 125 | return match (&options.alias_localpart, options.invite.as_slice()) { |
| 126 | (Some(localpart), _) => Some(Key::Alias(format!("#{localpart}:{}", server?))), |
| 127 | (None, [user_id]) if options.direct => Some(Key::Direct(user_id.clone())), |
| 128 | _ => None, |
| 129 | }; |
| 130 | } |
| 131 | Mutation::Room( |
| 132 | RoomCommand::Leave(room_id) |
| 133 | | RoomCommand::Forget(room_id) |
| 134 | | RoomCommand::AcceptInvite(room_id) |
| 135 | | RoomCommand::RejectInvite(room_id) |
| 136 | | RoomCommand::InviteUser { room_id, .. } |
| 137 | | RoomCommand::KickUser { room_id, .. } |
| 138 | | RoomCommand::BanUser { room_id, .. } |
| 139 | | RoomCommand::UnbanUser { room_id, .. } |
| 140 | | RoomCommand::SetPowerLevels { room_id, .. } |
| 141 | | RoomCommand::SetUserPowerLevel { room_id, .. } |
| 142 | | RoomCommand::SetAccountData { room_id, .. }, |
| 143 | ) => room_id, |
| 144 | Mutation::Room(_) => return None, |
| 145 | }; |
| 146 | Some(Key::Room(room_id.clone())) |
| 147 | } |
| 148 | |
| 149 | /// dms made here that the next sync hasn't told the sdk about yet, since `get_dm_room` |
| 150 | /// only knows `m.direct` once it comes back round |
| 151 | type FreshDms = Arc<std::sync::Mutex<HashMap<OwnedUserId, OwnedRoomId>>>; |
| 152 | |
| 153 | fn spawn( |
| 154 | tasks: &mut TaskSet<Key>, |
| 155 | client: &Client, |
| 156 | fresh_dms: &FreshDms, |
| 157 | key: Key, |
| 158 | job: MutationJob, |
| 159 | ) { |
| 160 | let client = client.clone(); |
| 161 | let fresh_dms = fresh_dms.clone(); |
| 162 | tasks.spawn(async move { |
| 163 | let Job { command, reply } = job; |
| 164 | reply.send(guarded("room command", execute(&client, &fresh_dms, command)).await); |
| 165 | key |
| 166 | }); |
| 167 | } |
| 168 | |
| 169 | async fn execute( |
| 170 | client: &Client, |
| 171 | fresh_dms: &FreshDms, |
| 172 | mutation: Mutation, |
| 173 | ) -> Result<CommandOutput, Error> { |
| 174 | let command = match mutation { |
| 175 | Mutation::GlobalAccountData { |
| 176 | event_type, |
| 177 | content, |
| 178 | } => { |
| 179 | client |
| 180 | .account() |
| 181 | .set_account_data_raw(event_type, content) |
| 182 | .await |
| 183 | .map_err(|source| Error::from_sdk("writing account data failed", source))?; |
| 184 | return Ok(CommandOutput::Done); |
| 185 | } |
| 186 | Mutation::Room(command) => command, |
| 187 | }; |
| 188 | let room = |room_id: &OwnedRoomId| client.get_room(room_id).ok_or_else(Error::room_unknown); |
| 189 | let joined = |room: Room| CommandOutput::Room(room.room_id().to_owned()); |
| 190 | let sdk = |message: &'static str| move |source| Error::from_sdk(message, source); |
| 191 | match command { |
| 192 | RoomCommand::Join { room, via } => client |
| 193 | .join_room_by_id_or_alias(&room, &via) |
| 194 | .await |
| 195 | .map(joined) |
| 196 | .map_err(sdk("joining room failed")), |
| 197 | RoomCommand::Leave(room_id) => room(&room_id)? |
| 198 | .leave() |
| 199 | .await |
| 200 | .map(done) |
| 201 | .map_err(sdk("leaving room failed")), |
| 202 | RoomCommand::Forget(room_id) => room(&room_id)? |
| 203 | .forget() |
| 204 | .await |
| 205 | .map(done) |
| 206 | .map_err(sdk("forgetting room failed")), |
| 207 | RoomCommand::AcceptInvite(room_id) => room(&room_id)? |
| 208 | .join() |
| 209 | .await |
| 210 | .map(done) |
| 211 | .map_err(sdk("accepting invite failed")), |
| 212 | RoomCommand::RejectInvite(room_id) => room(&room_id)? |
| 213 | .leave() |
| 214 | .await |
| 215 | .map(done) |
| 216 | .map_err(sdk("rejecting invite failed")), |
| 217 | RoomCommand::InviteUser { room_id, user_id } => room(&room_id)? |
| 218 | .invite_user_by_id(&user_id) |
| 219 | .await |
| 220 | .map(done) |
| 221 | .map_err(sdk("inviting user failed")), |
| 222 | RoomCommand::KickUser { |
| 223 | room_id, |
| 224 | user_id, |
| 225 | reason, |
| 226 | } => room(&room_id)? |
| 227 | .kick_user(&user_id, reason.as_deref()) |
| 228 | .await |
| 229 | .map(done) |
| 230 | .map_err(sdk("kicking user failed")), |
| 231 | RoomCommand::BanUser { |
| 232 | room_id, |
| 233 | user_id, |
| 234 | reason, |
| 235 | } => room(&room_id)? |
| 236 | .ban_user(&user_id, reason.as_deref()) |
| 237 | .await |
| 238 | .map(done) |
| 239 | .map_err(sdk("banning user failed")), |
| 240 | RoomCommand::UnbanUser { |
| 241 | room_id, |
| 242 | user_id, |
| 243 | reason, |
| 244 | } => room(&room_id)? |
| 245 | .unban_user(&user_id, reason.as_deref()) |
| 246 | .await |
| 247 | .map(done) |
| 248 | .map_err(sdk("unbanning user failed")), |
| 249 | RoomCommand::SetPowerLevels { room_id, changes } => { |
| 250 | apply_power_level_changes(&room(&room_id)?, changes) |
| 251 | .await |
| 252 | .map(done) |
| 253 | } |
| 254 | RoomCommand::SetUserPowerLevel { |
| 255 | room_id, |
| 256 | user_id, |
| 257 | level, |
| 258 | } => room(&room_id)? |
| 259 | .update_power_levels(vec![(&user_id, power_level(level)?)]) |
| 260 | .await |
| 261 | .map(done) |
| 262 | .map_err(sdk("updating user power level failed")), |
| 263 | RoomCommand::SetAccountData { |
| 264 | room_id, |
| 265 | event_type, |
| 266 | content, |
| 267 | } => room(&room_id)? |
| 268 | .set_account_data_raw(event_type, content) |
| 269 | .await |
| 270 | .map(done) |
| 271 | .map_err(sdk("writing room account data failed")), |
| 272 | RoomCommand::Create(options) => client |
| 273 | .create_room(create_room_request(options)) |
| 274 | .await |
| 275 | .map(joined) |
| 276 | .map_err(sdk("creating room failed")), |
| 277 | RoomCommand::CreateDirect(user_id) => { |
| 278 | if let Some(existing) = client.get_dm_room(&user_id) { |
| 279 | return Ok(joined(existing)); |
| 280 | } |
| 281 | let fresh = fresh_dms |
| 282 | .lock() |
| 283 | .expect("nothing panics holding this") |
| 284 | .get(&user_id) |
| 285 | .cloned(); |
| 286 | if let Some(room_id) = fresh.filter(|room_id| client.get_room(room_id).is_some()) { |
| 287 | return Ok(CommandOutput::Room(room_id)); |
| 288 | } |
| 289 | let room = client |
| 290 | .create_dm(&user_id) |
| 291 | .await |
| 292 | .map_err(sdk("creating direct room failed"))?; |
| 293 | fresh_dms |
| 294 | .lock() |
| 295 | .expect("nothing panics holding this") |
| 296 | .insert(user_id, room.room_id().to_owned()); |
| 297 | Ok(joined(room)) |
| 298 | } |
| 299 | _ => Err(Error::new(ErrorKind::Other, "room command was misrouted")), |
| 300 | } |
| 301 | } |
| 302 | |
| 303 | /// one state event carrying every requested change, since the sdk's own patcher is private |
| 304 | async fn apply_power_level_changes(room: &Room, changes: PowerLevelChanges) -> Result<(), Error> { |
| 305 | let mut levels = room |
| 306 | .power_levels() |
| 307 | .await |
| 308 | .map_err(|source| Error::from_sdk("reading power levels failed", source.into()))?; |
| 309 | for (target, value) in [ |
| 310 | (&mut levels.ban, changes.ban), |
| 311 | (&mut levels.invite, changes.invite), |
| 312 | (&mut levels.kick, changes.kick), |
| 313 | (&mut levels.redact, changes.redact), |
| 314 | (&mut levels.events_default, changes.events_default), |
| 315 | (&mut levels.state_default, changes.state_default), |
| 316 | (&mut levels.users_default, changes.users_default), |
| 317 | (&mut levels.notifications.room, changes.room_notification), |
| 318 | ] { |
| 319 | if let Some(value) = value { |
| 320 | *target = power_level(value)?; |
| 321 | } |
| 322 | } |
| 323 | for (event_type, value) in [ |
| 324 | (StateEventType::RoomName.into(), changes.room_name), |
| 325 | (StateEventType::RoomAvatar.into(), changes.room_avatar), |
| 326 | (StateEventType::RoomTopic.into(), changes.room_topic), |
| 327 | (StateEventType::SpaceChild.into(), changes.space_child), |
| 328 | (MessageLikeEventType::Beacon.into(), changes.beacon), |
| 329 | (StateEventType::BeaconInfo.into(), changes.beacon_info), |
| 330 | ] { |
| 331 | if let Some(value) = value { |
| 332 | levels.events.insert(event_type, power_level(value)?); |
| 333 | } |
| 334 | } |
| 335 | let content = RoomPowerLevelsEventContent::try_from(levels).map_err(|source| { |
| 336 | Error::with_source( |
| 337 | ErrorKind::Invalid, |
| 338 | "normalising power levels failed", |
| 339 | source, |
| 340 | ) |
| 341 | })?; |
| 342 | room.send_state_event(content) |
| 343 | .await |
| 344 | .map(drop) |
| 345 | .map_err(|source| Error::from_sdk("updating power levels failed", source)) |
| 346 | } |
| 347 | |
| 348 | fn done<T>(_: T) -> CommandOutput { |
| 349 | CommandOutput::Done |
| 350 | } |
| 351 | |
| 352 | fn power_level(value: i64) -> Result<Int, Error> { |
| 353 | value.try_into().map_err(|source| { |
| 354 | Error::with_source( |
| 355 | ErrorKind::Invalid, |
| 356 | "power level is outside the Matrix integer range", |
| 357 | source, |
| 358 | ) |
| 359 | }) |
| 360 | } |
| 361 | |
| 362 | fn create_room_request(options: CreateRoomOptions) -> CreateRoomRequest { |
| 363 | let mut request = CreateRoomRequest::new(); |
| 364 | request.name = options.name; |
| 365 | request.topic = options.topic; |
| 366 | request.room_alias_name = options.alias_localpart; |
| 367 | request.invite = options.invite; |
| 368 | request.is_direct = options.direct; |
| 369 | request.visibility = match options.visibility { |
| 370 | RoomVisibility::Private => SdkVisibility::Private, |
| 371 | RoomVisibility::Public => SdkVisibility::Public, |
| 372 | }; |
| 373 | request.preset = options.preset.map(|preset| match preset { |
| 374 | RoomPreset::PrivateChat => SdkPreset::PrivateChat, |
| 375 | RoomPreset::TrustedPrivateChat => SdkPreset::TrustedPrivateChat, |
| 376 | RoomPreset::PublicChat => SdkPreset::PublicChat, |
| 377 | }); |
| 378 | request |
| 379 | } |