macltok

public

Star0

Watch0

3 commits 3 branches 0 tags 0 releases3.0 MiB

use std::collections::{HashMap, HashSet, VecDeque};

use matrix_sdk::{
    Client, Room,
    ruma::{
        Int,
        api::client::room::{
            Visibility as SdkVisibility,
            create_room::v3::{Request as CreateRoomRequest, RoomPreset as SdkPreset},
        },
        events::{
            AnyGlobalAccountDataEventContent, GlobalAccountDataEventType, MessageLikeEventType,
            StateEventType, room::power_levels::RoomPowerLevelsEventContent,
        },
        serde::Raw,
    },
};
use ruma::{OwnedRoomId, OwnedUserId};
use tokio::{sync::mpsc, task::JoinSet};

use super::commands::{Job, guarded};
use crate::{
    CommandOutput, CreateRoomOptions, Error, ErrorKind, PowerLevelChanges, RoomCommand, RoomPreset,
    RoomVisibility,
};

/// state changes that must not race each other. same room, same alias, same dm
/// partner or same account data type run in order, everything else runs at once
pub(super) enum Mutation {
    Room(RoomCommand),
    GlobalAccountData {
        event_type: GlobalAccountDataEventType,
        content: Raw<AnyGlobalAccountDataEventContent>,
    },
}

pub(super) type MutationJob = Job<Mutation>;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
enum Key {
    Room(OwnedRoomId),
    Alias(String),
    Direct(OwnedUserId),
    AccountData(String),
    Independent(u64),
}

pub(super) async fn run_scheduler(
    client: Client,
    mut jobs: mpsc::Receiver<MutationJob>,
    queue_limit: usize,
) {
    let mut waiting = HashMap::<Key, VecDeque<MutationJob>>::new();
    let mut running = HashSet::<Key>::new();
    let mut tasks = JoinSet::<Key>::new();
    let mut queued = 0_usize;
    let mut independent = 0_u64;
    let mut receiving = true;
    while receiving || !tasks.is_empty() {
        tokio::select! {
            biased;
            Some(finished) = tasks.join_next(), if !tasks.is_empty() => {
                // a job's own panic is caught inside it, so this is only ever a key
                let Ok(key) = finished else { continue };
                let next = waiting.get_mut(&key).and_then(VecDeque::pop_front);
                if waiting.get(&key).is_some_and(VecDeque::is_empty) {
                    waiting.remove(&key);
                }
                match next {
                    Some(job) => {
                        queued -= 1;
                        spawn(&mut tasks, &client, key, job);
                    }
                    None => {
                        running.remove(&key);
                    }
                }
            }
            job = jobs.recv(), if receiving => {
                let Some(job) = job else {
                    receiving = false;
                    continue;
                };
                let key = key(&job.command).unwrap_or_else(|| {
                    independent += 1;
                    Key::Independent(independent)
                });
                if running.insert(key.clone()) {
                    spawn(&mut tasks, &client, key, job);
                } else if queued < queue_limit {
                    waiting.entry(key).or_default().push_back(job);
                    queued += 1;
                } else {
                    job.reply.done(Err(Error::new(
                        ErrorKind::QueueFull,
                        "room command queue is full",
                    )));
                }
            }
        }
    }
}

fn key(mutation: &Mutation) -> Option<Key> {
    let room_id = match mutation {
        Mutation::GlobalAccountData { event_type, .. } => {
            return Some(Key::AccountData(event_type.to_string()));
        }
        Mutation::Room(RoomCommand::Join { room, .. }) => {
            return Some(match OwnedRoomId::try_from(room.clone()) {
                Ok(room_id) => Key::Room(room_id),
                Err(alias) => Key::Alias(alias.to_string()),
            });
        }
        Mutation::Room(RoomCommand::CreateDirect(user_id)) => {
            return Some(Key::Direct(user_id.clone()));
        }
        Mutation::Room(
            RoomCommand::Leave(room_id)
            | RoomCommand::Forget(room_id)
            | RoomCommand::AcceptInvite(room_id)
            | RoomCommand::RejectInvite(room_id)
            | RoomCommand::InviteUser { room_id, .. }
            | RoomCommand::KickUser { room_id, .. }
            | RoomCommand::BanUser { room_id, .. }
            | RoomCommand::UnbanUser { room_id, .. }
            | RoomCommand::SetPowerLevels { room_id, .. }
            | RoomCommand::SetUserPowerLevel { room_id, .. }
            | RoomCommand::SetAccountData { room_id, .. },
        ) => room_id,
        Mutation::Room(_) => return None,
    };
    Some(Key::Room(room_id.clone()))
}

fn spawn(tasks: &mut JoinSet<Key>, client: &Client, key: Key, job: MutationJob) {
    let client = client.clone();
    tasks.spawn(async move {
        let Job { command, reply } = job;
        reply.send(guarded("room command", execute(&client, command)).await);
        key
    });
}

async fn execute(client: &Client, mutation: Mutation) -> Result<CommandOutput, Error> {
    let command = match mutation {
        Mutation::GlobalAccountData {
            event_type,
            content,
        } => {
            client
                .account()
                .set_account_data_raw(event_type, content)
                .await
                .map_err(|source| Error::from_sdk("writing account data failed", source))?;
            return Ok(CommandOutput::Done);
        }
        Mutation::Room(command) => command,
    };
    let room = |room_id: &OwnedRoomId| client.get_room(room_id).ok_or_else(Error::room_unknown);
    let joined = |room: Room| CommandOutput::Room(room.room_id().to_owned());
    let sdk = |message: &'static str| move |source| Error::from_sdk(message, source);
    match command {
        RoomCommand::Join { room, via } => client
            .join_room_by_id_or_alias(&room, &via)
            .await
            .map(joined)
            .map_err(sdk("joining room failed")),
        RoomCommand::Leave(room_id) => room(&room_id)?
            .leave()
            .await
            .map(done)
            .map_err(sdk("leaving room failed")),
        RoomCommand::Forget(room_id) => room(&room_id)?
            .forget()
            .await
            .map(done)
            .map_err(sdk("forgetting room failed")),
        RoomCommand::AcceptInvite(room_id) => room(&room_id)?
            .join()
            .await
            .map(done)
            .map_err(sdk("accepting invite failed")),
        RoomCommand::RejectInvite(room_id) => room(&room_id)?
            .leave()
            .await
            .map(done)
            .map_err(sdk("rejecting invite failed")),
        RoomCommand::InviteUser { room_id, user_id } => room(&room_id)?
            .invite_user_by_id(&user_id)
            .await
            .map(done)
            .map_err(sdk("inviting user failed")),
        RoomCommand::KickUser {
            room_id,
            user_id,
            reason,
        } => room(&room_id)?
            .kick_user(&user_id, reason.as_deref())
            .await
            .map(done)
            .map_err(sdk("kicking user failed")),
        RoomCommand::BanUser {
            room_id,
            user_id,
            reason,
        } => room(&room_id)?
            .ban_user(&user_id, reason.as_deref())
            .await
            .map(done)
            .map_err(sdk("banning user failed")),
        RoomCommand::UnbanUser {
            room_id,
            user_id,
            reason,
        } => room(&room_id)?
            .unban_user(&user_id, reason.as_deref())
            .await
            .map(done)
            .map_err(sdk("unbanning user failed")),
        RoomCommand::SetPowerLevels { room_id, changes } => {
            apply_power_level_changes(&room(&room_id)?, changes)
                .await
                .map(done)
        }
        RoomCommand::SetUserPowerLevel {
            room_id,
            user_id,
            level,
        } => room(&room_id)?
            .update_power_levels(vec![(&user_id, power_level(level)?)])
            .await
            .map(done)
            .map_err(sdk("updating user power level failed")),
        RoomCommand::SetAccountData {
            room_id,
            event_type,
            content,
        } => room(&room_id)?
            .set_account_data_raw(event_type, content)
            .await
            .map(done)
            .map_err(sdk("writing room account data failed")),
        RoomCommand::Create(options) => client
            .create_room(create_room_request(options))
            .await
            .map(joined)
            .map_err(sdk("creating room failed")),
        RoomCommand::CreateDirect(user_id) => match client.get_dm_room(&user_id) {
            Some(existing) => Ok(joined(existing)),
            None => client
                .create_dm(&user_id)
                .await
                .map(joined)
                .map_err(sdk("creating direct room failed")),
        },
        _ => Err(Error::new(ErrorKind::Other, "room command was misrouted")),
    }
}

/// one state event carrying every requested change, since the sdk's own patcher is private
async fn apply_power_level_changes(room: &Room, changes: PowerLevelChanges) -> Result<(), Error> {
    let mut levels = room
        .power_levels()
        .await
        .map_err(|source| Error::from_sdk("reading power levels failed", source.into()))?;
    for (target, value) in [
        (&mut levels.ban, changes.ban),
        (&mut levels.invite, changes.invite),
        (&mut levels.kick, changes.kick),
        (&mut levels.redact, changes.redact),
        (&mut levels.events_default, changes.events_default),
        (&mut levels.state_default, changes.state_default),
        (&mut levels.users_default, changes.users_default),
        (&mut levels.notifications.room, changes.room_notification),
    ] {
        if let Some(value) = value {
            *target = power_level(value)?;
        }
    }
    for (event_type, value) in [
        (StateEventType::RoomName.into(), changes.room_name),
        (StateEventType::RoomAvatar.into(), changes.room_avatar),
        (StateEventType::RoomTopic.into(), changes.room_topic),
        (StateEventType::SpaceChild.into(), changes.space_child),
        (MessageLikeEventType::Beacon.into(), changes.beacon),
        (StateEventType::BeaconInfo.into(), changes.beacon_info),
    ] {
        if let Some(value) = value {
            levels.events.insert(event_type, power_level(value)?);
        }
    }
    let content = RoomPowerLevelsEventContent::try_from(levels).map_err(|source| {
        Error::with_source(
            ErrorKind::Invalid,
            "normalising power levels failed",
            source,
        )
    })?;
    room.send_state_event(content)
        .await
        .map(drop)
        .map_err(|source| Error::from_sdk("updating power levels failed", source))
}

fn done<T>(_: T) -> CommandOutput {
    CommandOutput::Done
}

fn power_level(value: i64) -> Result<Int, Error> {
    value.try_into().map_err(|source| {
        Error::with_source(
            ErrorKind::Invalid,
            "power level is outside the Matrix integer range",
            source,
        )
    })
}

fn create_room_request(options: CreateRoomOptions) -> CreateRoomRequest {
    let mut request = CreateRoomRequest::new();
    request.name = options.name;
    request.topic = options.topic;
    request.room_alias_name = options.alias_localpart;
    request.invite = options.invite;
    request.is_direct = options.direct;
    request.visibility = match options.visibility {
        RoomVisibility::Private => SdkVisibility::Private,
        RoomVisibility::Public => SdkVisibility::Public,
    };
    request.preset = options.preset.map(|preset| match preset {
        RoomPreset::PrivateChat => SdkPreset::PrivateChat,
        RoomPreset::TrustedPrivateChat => SdkPreset::TrustedPrivateChat,
        RoomPreset::PublicChat => SdkPreset::PublicChat,
    });
    request
}

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