refactor: create presence crate and add tests

Cargo.lock

@@ -210,9 +210,9 @@ dependencies = [
 
 [[package]]
 name = "async-std"
-version = "1.11.0"
+version = "1.12.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "52580991739c5cdb36cde8b2a516371c0a3b70dda36d916cc08b82372916808c"
+checksum = "62565bb4402e926b29953c785397c6dc0391b7b446e45008b0049eb43cec6f5d"
 dependencies = [
  "async-attributes",
  "async-channel",
@@ -229,7 +229,6 @@ dependencies = [
  "kv-log-macro",
  "log",
  "memchr",
- "num_cpus",
  "once_cell",
  "pin-project-lite 0.2.9",
  "pin-utils",
@@ -2808,6 +2807,7 @@ dependencies = [
  "log",
  "once_cell",
  "querystring",
+ "revolt-presence",
  "revolt-quark",
  "rmp-serde",
  "serde",
@@ -2879,6 +2879,17 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "revolt-presence"
+version = "0.0.1"
+dependencies = [
+ "async-std",
+ "log",
+ "once_cell",
+ "rand 0.8.5",
+ "redis-kiss",
+]
+
 [[package]]
 name = "revolt-quark"
 version = "0.5.19"
@@ -2911,6 +2922,7 @@ dependencies = [
  "redis-kiss",
  "regex",
  "reqwest",
+ "revolt-presence",
  "revolt-result",
  "revolt_okapi",
  "revolt_optional_struct",

crates/bonfire/Cargo.toml

@@ -26,3 +26,6 @@ serde = "1.0.136"
 futures = "0.3.21"
 async-tungstenite = { version = "0.17.0", features = ["async-std-runtime"] }
 async-std = { version = "1.8.0", features = ["tokio1", "tokio02", "attributes"] }
+
+# core
+revolt-presence = { path = "../core/presence" }

crates/bonfire/src/main.rs

@@ -1,7 +1,7 @@
 use std::env;
 
 use async_std::net::TcpListener;
-use revolt_quark::presence::presence_clear_region;
+use revolt_presence::clear_region;
 
 #[macro_use]
 extern crate log;
@@ -18,7 +18,7 @@ async fn main() {
     database::connect().await;
 
     // Clean up the current region information.
-    presence_clear_region(None).await;
+    clear_region(None).await;
 
     // Setup a TCP listener to accept WebSocket connections on.
     // By default, we bind to port 9000 on all interfaces.

crates/bonfire/src/websocket.rs

@@ -1,6 +1,7 @@
 use std::net::SocketAddr;
 
 use futures::{channel::oneshot, pin_mut, select, FutureExt, SinkExt, StreamExt, TryStreamExt};
+use revolt_presence::{create_session, delete_session};
 use revolt_quark::{
     events::{
         client::EventV1,
@@ -8,7 +9,6 @@ use revolt_quark::{
         state::{State, SubscriptionStateChange},
     },
     models::{user::UserHint, User},
-    presence::{presence_create_session, presence_delete_session},
     redis_kiss, Database,
 };
 
@@ -69,8 +69,7 @@ pub fn spawn_client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
                             let user_id = state.cache.user_id.clone();
 
                             // Create presence session.
-                            let (first_session, session_id) =
-                                presence_create_session(&user_id, 0).await;
+                            let (first_session, session_id) = create_session(&user_id, 0).await;
 
                             // Notify socket we have authenticated.
                             write
@@ -225,7 +224,7 @@ pub fn spawn_client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
                             }
 
                             // Clean up presence session.
-                            let last_session = presence_delete_session(&user_id, session_id).await;
+                            let last_session = delete_session(&user_id, session_id).await;
 
                             // If this was the last session, notify other users that we just went offline.
                             if last_session {

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

@@ -1,20 +1,31 @@
+#[macro_use]
+extern crate log;
+
+use once_cell::sync::Lazy;
+use rand::Rng;
+use redis_kiss::{get_connection, AsyncCommands};
 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};
+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 presence_create_session(user_id: &str, flags: u8) -> (bool, u32) {
+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 {
@@ -24,7 +35,7 @@ pub async fn presence_create_session(user_id: &str, flags: u8) -> (bool, u32) {
         // 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)
+            (rng.gen::<u32>() & !FLAG_BITS) | (flags as u32 & FLAG_BITS)
         };
 
         // Add session to user's sessions and to the region
@@ -41,16 +52,12 @@ pub async fn presence_create_session(user_id: &str, flags: u8) -> (bool, u32) {
 }
 
 /// 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
+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 presence_delete_session_internal(
-    user_id: &str,
-    session_id: u32,
-    skip_region: bool,
-) -> bool {
+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 {
@@ -78,7 +85,7 @@ async fn presence_delete_session_internal(
 }
 
 /// Check whether a given user ID is online
-pub async fn presence_is_online(user_id: &str) -> bool {
+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 {
@@ -87,7 +94,7 @@ pub async fn presence_is_online(user_id: &str) -> bool {
 }
 
 /// 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> {
+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() {
@@ -102,7 +109,7 @@ pub async fn presence_filter_online(user_ids: &'_ [String]) -> HashSet<String> {
     // 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 {
+        if is_online(&user_ids[0]).await {
             set.insert(user_ids[0].to_string());
         }
 
@@ -133,7 +140,7 @@ pub async fn presence_filter_online(user_ids: &'_ [String]) -> HashSet<String> {
 }
 
 /// Reset any stale presence data
-pub async fn presence_clear_region(region_id: Option<&str>) {
+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");
 
@@ -150,7 +157,7 @@ pub async fn presence_clear_region(region_id: Option<&str>) {
             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;
+                    delete_session_internal(user_id, session_id, true).await;
                 }
             }
         }
@@ -161,3 +168,59 @@ pub async fn presence_clear_region(region_id: Option<&str>) {
         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/quark/Cargo.toml

@@ -92,3 +92,4 @@ sentry = "0.25.0"
 
 # Core
 revolt-result = { path = "../core/result", features = [ "serde", "schemas" ] }
+revolt-presence = { path = "../core/presence" }

crates/quark/examples/test.rs

@@ -1,34 +0,0 @@
-use revolt_quark::models::user::PartialUser;
-use revolt_quark::presence::{
-    presence_create_session, presence_delete_session, presence_filter_online, presence_is_online,
-};
-use revolt_quark::*;
-
-#[async_std::main]
-async fn main() {
-    let db = DatabaseInfo::Dummy.connect().await.unwrap();
-
-    let sus = PartialUser {
-        username: Some("neat".into()),
-        ..Default::default()
-    };
-
-    db.update_user("user id", &sus, vec![]).await.unwrap();
-
-    dbg!(presence_create_session("entry", 0).await);
-    dbg!(presence_is_online("entry").await);
-    dbg!(presence_filter_online(&["a".into(), "b".into(), "entry".into()]).await);
-
-    dbg!(presence_delete_session("entry", 0).await);
-    dbg!(presence_is_online("entry").await);
-
-    dbg!(presence_create_session("entry", 0).await);
-    dbg!(presence_create_session("entry", 0).await);
-    dbg!(presence_delete_session("entry", 0).await);
-    dbg!(presence_is_online("entry").await);
-    dbg!(presence_delete_session("entry", 1).await);
-    dbg!(presence_is_online("entry").await);
-
-    // __set_key("dietz", vec![0xFF]).await;
-    // dbg!(presence_filter_online(&["dietz".into(), "nuts".into()]).await);
-}

crates/quark/src/events/impl.rs

@@ -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.

crates/quark/src/impl/generic/channels/message.rs

@@ -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>>();
                     }

crates/quark/src/impl/generic/users/user.rs

@@ -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)

crates/quark/src/lib.rs

@@ -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;

crates/quark/src/presence/entry.rs

@@ -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";

crates/quark/src/util/ref.rs

@@ -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)