refactor: create presence crate and add tests

crates/core/presence/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "revolt-presence"
+version = "0.0.1"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dev-dependencies]
+# Async
+async-std = { version = "1.8.0", features = ["attributes"] }
+
+[dependencies]
+# Utility
+log = "0.4.17"
+rand = "0.8.5"
+once_cell = "1.17.1"
+
+# Redis
+redis-kiss = "0.1.4"

crates/core/presence/src/lib.rs

@@ -0,0 +1,226 @@
+#[macro_use]
+extern crate log;
+
+use once_cell::sync::Lazy;
+use rand::Rng;
+use redis_kiss::{get_connection, AsyncCommands};
+use std::collections::HashSet;
+
+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,
+};
+
+pub static REGION_ID: Lazy<u16> = Lazy::new(|| {
+    std::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";
+
+pub static FLAG_BITS: u32 = 0b1;
+
+/// Create a new presence session, returns the ID of this session
+pub async fn 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>() & !FLAG_BITS) | (flags as u32 & FLAG_BITS)
+        };
+
+        // 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 delete_session(user_id: &str, session_id: u32) -> bool {
+    delete_session_internal(user_id, session_id, false).await
+}
+
+/// Delete existing presence session (but also choose whether to skip region)
+async fn 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 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 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 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 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() {
+                    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.");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{clear_region, create_session, delete_session, filter_online, is_online};
+    use rand::Rng;
+
+    #[async_std::test]
+    async fn it_works() {
+        // Clear the region before we start the tests:
+        clear_region(None).await;
+
+        // Generate some data we'll use:
+        let user_id = rand::thread_rng().gen::<u32>().to_string();
+        let other_id = rand::thread_rng().gen::<u32>().to_string();
+        let flags = 1;
+
+        // Create a session
+        let (first_session, session_id) = create_session(&user_id, flags).await;
+        assert!(first_session);
+        assert_ne!(session_id, 0);
+        assert_eq!(session_id as u8 & flags, flags);
+
+        // Check if the user is online
+        assert!(is_online(&user_id).await);
+
+        let user_ids = filter_online(&[user_id.to_string()]).await;
+        assert_eq!(user_ids.len(), 1);
+        assert!(user_ids.contains(&user_id));
+
+        // Create a few more sessions
+        let (first_session, second_session_id) = create_session(&user_id, 0).await;
+        assert!(!first_session);
+        dbg!(second_session_id);
+        assert_eq!(second_session_id as u8 & 1, 0);
+
+        let (first_session, other_session_id) = create_session(&other_id, 0).await;
+        assert!(first_session);
+
+        let user_ids = filter_online(&[user_id.to_string(), other_id.to_string()]).await;
+        assert_eq!(user_ids.len(), 2);
+        assert!(user_ids.contains(&user_id));
+        assert!(user_ids.contains(&other_id));
+
+        // Remove sessions
+        delete_session(&user_id, session_id).await;
+        delete_session(&other_id, other_session_id).await;
+        assert!(!is_online(&other_id).await);
+
+        // Check if we can wipe everything too
+        clear_region(None).await;
+        assert!(!is_online(&user_id).await);
+
+        let user_ids = filter_online(&[user_id.to_string(), other_id.to_string()]).await;
+        assert!(user_ids.is_empty())
+    }
+}

crates/core/presence/src/operations.rs

@@ -0,0 +1,42 @@
+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");
+}