jmug/stoatchat
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

refactor: create presence crate and add tests

Browse Source
This commit is contained in:
Paul Makles
2023-04-22 20:01:23 +01:00
parent c817c2dd40
commit 12d963d2bd
15 changed files with 136 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
crates/quark/src/events/impl.rs
View File

@@ -7,11 +7,11 @@ use crate::{
user::{PartialUser, Presence, RelationshipStatus},
Channel, Member, User,
},
perms,
presence::presence_filter_online,
Database, Permission, Result,
perms, Database, Permission, Result,
};
use revolt_presence::filter_online;
use super::{
client::EventV1,
state::{Cache, State},
@@ -135,8 +135,7 @@ impl State {
}
// Fetch presence data for known users.
let online_ids =
presence_filter_online(&user_ids.iter().cloned().collect::<Vec<String>>()).await;
let online_ids = filter_online(&user_ids.iter().cloned().collect::<Vec<String>>()).await;
user.online = Some(true);
// Fetch user data.

4
crates/quark/src/impl/generic/channels/message.rs
View File

@@ -1,5 +1,6 @@
use std::collections::HashSet;
use revolt_presence::filter_online;
use serde_json::json;
use ulid::Ulid;
use validator::Validate;
@@ -14,7 +15,6 @@ use crate::{
Channel, Emoji, Message, User,
},
permissions::PermissionCalculator,
presence::presence_filter_online,
tasks::ack::AckEvent,
types::{
january::{Embed, Text},
@@ -79,7 +79,7 @@ impl Message {
Channel::DirectMessage { recipients, .. }
| Channel::Group { recipients, .. } => {
target_ids = (&recipients.iter().cloned().collect::<HashSet<String>>()
- &presence_filter_online(recipients).await)
- &filter_online(recipients).await)
.into_iter()
.collect::<Vec<String>>();
}

4
crates/quark/src/impl/generic/users/user.rs
View File

@@ -4,11 +4,11 @@ use crate::models::user::{
};
use crate::permissions::defn::UserPerms;
use crate::permissions::r#impl::user::get_relationship;
use crate::presence::presence_filter_online;
use crate::{perms, Database, Error, Result};
use futures::try_join;
use impl_ops::impl_op_ex_commutative;
use revolt_presence::filter_online;
use std::ops;
impl_op_ex_commutative!(+ |a: &i32, b: &Badges| -> i32 { *a | *b as i32 });
@@ -98,7 +98,7 @@ impl User {
/// Fetch foreign users by a list of IDs
pub async fn fetch_foreign_users(db: &Database, user_ids: &[String]) -> Result<Vec<User>> {
let online_ids = presence_filter_online(user_ids).await;
let online_ids = filter_online(user_ids).await;
Ok(db
.fetch_users(user_ids)

1
crates/quark/src/lib.rs
View File

@@ -22,7 +22,6 @@ pub use redis_kiss;
pub mod events;
pub mod r#impl;
pub mod models;
pub mod presence;
pub mod tasks;
pub mod types;
pub mod util;

13
crates/quark/src/presence/entry.rs
View File

@@ -1,13 +0,0 @@
use std::env;
use once_cell::sync::Lazy;
pub static REGION_ID: Lazy<u16> = Lazy::new(|| {
env::var("REGION_ID")
.unwrap_or_else(|_| "0".to_string())
.parse()
.unwrap()
});
pub static REGION_KEY: Lazy<String> = Lazy::new(|| format!("region{}", &*REGION_ID));
pub static ONLINE_SET: &str = "online";

163
crates/quark/src/presence/mod.rs
View File

@@ -1,163 +0,0 @@
use std::collections::HashSet;
use redis_kiss::{get_connection, AsyncCommands};
use rand::Rng;
mod entry;
mod operations;
use operations::{
__add_to_set_string, __add_to_set_u32, __delete_key, __get_set_members_as_string,
__get_set_size, __remove_from_set_string, __remove_from_set_u32,
};
use self::entry::{ONLINE_SET, REGION_KEY};
/// Create a new presence session, returns the ID of this session
pub async fn presence_create_session(user_id: &str, flags: u8) -> (bool, u32) {
info!("Creating a presence session for {user_id} with flags {flags}");
if let Ok(mut conn) = get_connection().await {
// Check whether this is the first session
let was_empty = __get_set_size(&mut conn, user_id).await == 0;
// A session ID is comprised of random data and any flags ORed to the end
let session_id = {
let mut rng = rand::thread_rng();
(rng.gen::<u32>() ^ 1) | (flags as u32 & 1)
};
// Add session to user's sessions and to the region
__add_to_set_u32(&mut conn, user_id, session_id).await;
__add_to_set_string(&mut conn, ONLINE_SET, user_id).await;
__add_to_set_string(&mut conn, &REGION_KEY, &format!("{user_id}:{session_id}")).await;
info!("Created session for {user_id}, assigned them a session ID of {session_id}.");
(was_empty, session_id)
} else {
// Fail through
(false, 0)
}
}
/// Delete existing presence session
pub async fn presence_delete_session(user_id: &str, session_id: u32) -> bool {
presence_delete_session_internal(user_id, session_id, false).await
}
/// Delete existing presence session (but also choose whether to skip region)
async fn presence_delete_session_internal(
user_id: &str,
session_id: u32,
skip_region: bool,
) -> bool {
info!("Deleting presence session for {user_id} with id {session_id}");
if let Ok(mut conn) = get_connection().await {
// Remove the session
__remove_from_set_u32(&mut conn, user_id, session_id).await;
// Remove from the region
if !skip_region {
__remove_from_set_string(&mut conn, &REGION_KEY, &format!("{user_id}:{session_id}"))
.await;
}
// Return whether this was the last session
let is_empty = __get_set_size(&mut conn, user_id).await == 0;
if is_empty {
__remove_from_set_string(&mut conn, ONLINE_SET, user_id).await;
info!("User ID {} just went offline.", &user_id);
}
is_empty
} else {
// Fail through
false
}
}
/// Check whether a given user ID is online
pub async fn presence_is_online(user_id: &str) -> bool {
if let Ok(mut conn) = get_connection().await {
conn.exists(user_id).await.unwrap_or(false)
} else {
false
}
}
/// Check whether a set of users is online, returns a set of the online user IDs
pub async fn presence_filter_online(user_ids: &'_ [String]) -> HashSet<String> {
// Ignore empty list immediately, to save time.
let mut set = HashSet::new();
if user_ids.is_empty() {
return set;
}
// NOTE: at the point that we need mobile indicators
// you can interpret the data here and return a new data
// structure like HashMap<String /* id */, u8 /* flags */>
// We need to handle a special case where only one is present
// as for some reason or another, Redis does not like us sending
// a list of just one ID to the server.
if user_ids.len() == 1 {
if presence_is_online(&user_ids[0]).await {
set.insert(user_ids[0].to_string());
}
return set;
}
// Otherwise, go ahead as normal.
if let Ok(mut conn) = get_connection().await {
// Ok so, if this breaks, that means we've lost the Redis patch which adds SMISMEMBER
// Currently it's patched in through a forked repository, investigate what happen to it
let data: Vec<bool> = conn
.sismember("online", user_ids)
.await
.expect("this shouldn't happen, please read this code! presence/mod.rs");
if data.is_empty() {
return set;
}
// We filter known values to figure out who is online.
for i in 0..user_ids.len() {
if data[i] {
set.insert(user_ids[i].to_string());
}
}
}
set
}
/// Reset any stale presence data
pub async fn presence_clear_region(region_id: Option<&str>) {
let region_id = region_id.unwrap_or(&*REGION_KEY);
let mut conn = get_connection().await.expect("Redis connection");
let sessions = __get_set_members_as_string(&mut conn, region_id).await;
if !sessions.is_empty() {
info!(
"Cleaning up {} sessions, this may take a while...",
sessions.len()
);
// Iterate and delete each session, this will
// also send out any relevant events.
for session in sessions {
let parts = session.split(':').collect::<Vec<&str>>();
if let (Some(user_id), Some(session_id)) = (parts.first(), parts.get(1)) {
if let Ok(session_id) = session_id.parse() {
presence_delete_session_internal(user_id, session_id, true).await;
}
}
}
// Then clear the set in Redis.
__delete_key(&mut conn, region_id).await;
info!("Clean up complete.");
}
}

42
crates/quark/src/presence/operations.rs
View File

@@ -1,42 +0,0 @@
use redis_kiss::{AsyncCommands, Conn};
/// Add to set (string)
pub async fn __add_to_set_string(conn: &mut Conn, key: &str, value: &str) {
let _: Option<()> = conn.sadd(key, value).await.ok();
}
/// Add to set (u32)
pub async fn __add_to_set_u32(conn: &mut Conn, key: &str, value: u32) {
let _: Option<()> = conn.sadd(key, value).await.ok();
}
/// Remove from set (string)
pub async fn __remove_from_set_string(conn: &mut Conn, key: &str, value: &str) {
let _: Option<()> = conn.srem(key, value).await.ok();
}
/// Remove from set (u32)
pub async fn __remove_from_set_u32(conn: &mut Conn, key: &str, value: u32) {
let _: Option<()> = conn.srem(key, value).await.ok();
}
/// Get set members as string
pub async fn __get_set_members_as_string(conn: &mut Conn, key: &str) -> Vec<String> {
conn.smembers::<_, Vec<String>>(key)
.await
.expect("could not get set members as string")
}
/// Get set size
pub async fn __get_set_size(conn: &mut Conn, id: &str) -> u32 {
conn.scard::<_, u32>(id)
.await
.expect("could not get set size")
}
/// Delete key by id
pub async fn __delete_key(conn: &mut Conn, id: &str) {
conn.del::<_, ()>(id)
.await
.expect("could not delete key by id");
}

4
crates/quark/src/util/ref.rs
View File

@@ -1,4 +1,5 @@
use futures::future::join;
use revolt_presence::is_online;
use rocket::request::FromParam;
use schemars::schema::{InstanceType, Schema, SchemaObject, SingleOrVec};
use schemars::JsonSchema;
@@ -7,7 +8,6 @@ use serde::{Deserialize, Serialize};
use crate::models::{
Bot, Channel, Emoji, Invite, Member, Message, Report, Server, ServerBan, User,
};
use crate::presence::presence_is_online;
use crate::{Database, Error, Result};
/// Reference to some object in the database
@@ -25,7 +25,7 @@ impl Ref {
/// Fetch user from Ref
pub async fn as_user(&self, db: &Database) -> Result<User> {
let (user, online) = join(db.fetch_user(&self.id), presence_is_online(&self.id)).await;
let (user, online) = join(db.fetch_user(&self.id), is_online(&self.id)).await;
let mut user = user?;
user.online = Some(online);
Ok(user)