Alt lookup speedup

Roughly makes looking up alts ~10x faster

Author: chatasma

src/database/cache.rs

@@ -111,7 +111,7 @@ pub async fn get_redis_pool(redis_host: &Option<String>) -> anyhow::Result<Redis
         None => Err(ConfigMissingFieldError {field_name: String::from("redis-host") }.into()),
         Some(redis_host) => {
             let redis_uri = format!("redis://{}", redis_host);
-            println!("Connecting to redis at {}", &redis_uri);
+            info!("Connecting to redis at {}", &redis_uri);
             let client = redis::Client::open(redis_uri)?;
             let manager = RedisConnectionManager::new(client);
             let pool = Pool::builder()

src/database/migrations/denormalize_ip_identities.rs

@@ -0,0 +1,106 @@
+use std::collections::HashMap;
+use std::ffi::c_int;
+use futures::{Stream, StreamExt, TryStreamExt};
+use futures::stream::FuturesUnordered;
+use mongodb::bson::{doc, Document};
+use mongodb::error::Error;
+use mongodb::options::FindOptions;
+use crate::database::Database;
+use crate::database::migrations::DatabaseMigration;
+use crate::database::models::ip_identity::IpIdentity;
+use crate::database::models::player::Player;
+
+pub struct DenormalizeIpIdentitiesMigration {}
+
+#[async_trait]
+impl DatabaseMigration for DenormalizeIpIdentitiesMigration {
+    fn get_id(&self) -> String {
+        String::from("denormalize_ip_identities")
+    }
+
+    async fn perform(&self, database: &Database) {
+        let mut find_options = FindOptions::default();
+        // batch-read into memory 50k records at a time
+        find_options.batch_size = Some(50_000);
+
+        let count = database.players.count_documents(doc! {}, None).await.unwrap_or(0);
+        info!("{} document(s) in the player collection to migrate", count);
+
+        let mut cursor = database.players.find(doc! {}, Some(find_options)).await.expect("find all players to succeed");
+        // aggregate 25k players' worth of IPs into memory at a time before flushing
+        let step_size = 25_000u32;
+        let mut total_accumulated = 0u32;
+        let mut accumulated = 0u32;
+        let mut error_count = 0u32;
+        let mut ip_map : HashMap<String, Vec<String>> = HashMap::new();
+        while let Ok(more) = cursor.advance().await {
+            if !more {
+                break
+            }
+            if accumulated >= step_size {
+                info!("Flushing batch of IPs to ip identities, accumulation progress: {}/{}", total_accumulated, count);
+                Self::flush_ips(database, ip_map).await;
+                ip_map = HashMap::new();
+                accumulated = 0;
+            }
+            let player = match cursor.deserialize_current() {
+                Ok(p) => p,
+                Err(e) => {
+                    let doc_result : mongodb::bson::raw::Result<Document> = cursor.current().try_into();
+                    let doc = match doc_result {
+                        Ok(doc) => doc,
+                        Err(e) => {
+                            warn!("Error to parse doc: {}", e);
+                            return
+                        }
+                    };
+                    warn!("document in question: {:?}", doc);
+                    warn!("Deserialization error: {}", e);
+                    error_count += 1;
+                    continue
+                }
+            };
+            for ip in &player.ips {
+                match ip_map.get_mut(ip) {
+                    Some(players_for_ip) => {
+                        players_for_ip.push(player.id.to_owned());
+                    }
+                    None => {
+                        ip_map.insert(ip.to_owned(), vec![player.id.to_owned()]);
+                    }
+                }
+            }
+            accumulated += 1;
+            total_accumulated += 1;
+        }
+        info!("Flushing any remaining IPs...");
+        Self::flush_ips(database, ip_map).await;
+        info!("Total flushed: {}", total_accumulated);
+        info!("Error count: {}", error_count);
+    }
+}
+
+impl DenormalizeIpIdentitiesMigration {
+    async fn flush_ips(database: &Database, ip_map : HashMap<String, Vec<String>>) {
+        let unordered_futures = FuturesUnordered::new();
+        for (ip, players) in ip_map.into_iter() {
+            let task = async {
+                let ip_identity_optional = IpIdentity::get_ip_identity_by_ip(&database.ip_identities, &ip).await;
+                match ip_identity_optional {
+                    Some(mut ip_identity) => {
+                        ip_identity.players.extend(players);
+                        database.save(&ip_identity).await;
+                    }
+                    None => {
+                        let ip_identity = IpIdentity {
+                            ip, players
+                        };
+                        database.save(&ip_identity).await
+                    }
+                }
+            };
+            unordered_futures.push(task);
+        }
+        unordered_futures.collect::<Vec<_>>().await;
+    }
+}

src/database/migrations/mod.rs

@@ -0,0 +1,33 @@
+use crate::database::Database;
+use crate::database::migrations::denormalize_ip_identities::DenormalizeIpIdentitiesMigration;
+
+pub mod denormalize_ip_identities;
+
+#[async_trait]
+pub trait DatabaseMigration {
+    fn get_id(&self) -> String;
+    async fn perform(&self, database: &Database);
+}
+
+pub struct MigrationExecutor {
+    migrations: Vec<Box<dyn DatabaseMigration>>
+}
+
+impl MigrationExecutor {
+    pub fn new() -> Self {
+        let denormalize_ip_identities_migration =
+            Box::new(DenormalizeIpIdentitiesMigration {});
+        Self {
+            migrations: vec![denormalize_ip_identities_migration]
+        }
+    }
+
+    pub async fn execute_migration_by_name(&self, database: &Database, name: String) {
+        match self.migrations.iter().find(|migration| migration.get_id() == name) {
+            Some(migration) => {
+                migration.perform(database).await
+            }
+            None => {}
+        }
+    }
+}

src/database/mod.rs

@@ -1,4 +1,5 @@
 use std::{str::FromStr, time::Duration};
+use std::collections::HashSet;
 
 use mars_api_rs_macro::IdentifiableDocument;
 use mongodb::{options::{ClientOptions, FindOneOptions, UpdateOptions}, Client, Collection, bson::{doc, oid::ObjectId}, Cursor, results::DeleteResult};
@@ -8,13 +9,17 @@ use rocket::serde::DeserializeOwned;
 use futures::StreamExt;
 use serde::Serialize;
 use anyhow::anyhow;
+use futures::stream::FuturesUnordered;
+use rocket::form::validate::Contains;
 
 use crate::{database::models::player::Player, util::r#macro::unwrap_helper};
+use crate::database::models::ip_identity::IpIdentity;
 use crate::util::validation::verbose_result_ok;
 
 use self::models::{achievement::Achievement, death::Death, level::Level, r#match::Match, punishment::Punishment, rank::Rank, session::Session};
 
 pub mod models;
+pub mod migrations;
 pub mod cache;
 
 pub trait CollectionOwner<T> {
@@ -32,7 +37,8 @@ pub struct Database {
     pub ranks: Collection<Rank>,
     pub matches: Collection<Match>,
     pub deaths: Collection<Death>,
-    pub levels: Collection<Level>
+    pub levels: Collection<Level>,
+    pub ip_identities: Collection<IpIdentity>
 }
 
 impl Database {
@@ -148,10 +154,26 @@ impl Database {
     }
 
     pub async fn get_alts_for_player(&self, player: &Player) -> Vec<Player> {
-        let cursor = unwrap_helper::result_return_default!(self.players.find(doc! {
-            "ips": {"$in": &player.ips}, "_id": {"$ne": &player.id}
-        }, None).await, Vec::new());
-        Database::consume_cursor_into_owning_vec(cursor).await
+        let unordered_futures = FuturesUnordered::new();
+        for ip in &player.ips {
+            unordered_futures.push(IpIdentity::find_players_for_ip(self, ip));
+        }
+        let players = {
+            let players_with_duplicates = unordered_futures.collect::<Vec<_>>().await
+                .into_iter().flatten().collect::<Vec<_>>();
+            let mut player_set : Vec<Player> = Vec::new();
+            for player in players_with_duplicates {
+                if !player_set.iter().map(|p| p.id.clone()).collect::<Vec<_>>().contains(&player.id) {
+                    player_set.push(player);
+                }
+            }
+            player_set
+        };
+        players
+        // let cursor = unwrap_helper::result_return_default!(self.players.find(doc! {
+        //     "ips": {"$in": &player.ips}, "_id": {"$ne": &player.id}
+        // }, None).await, Vec::new());
+        // Database::consume_cursor_into_owning_vec(cursor).await
     }
 
     pub async fn save<R>(&self, record: &R) where R: CollectionOwner<R> + Serialize + IdentifiableDocument {
@@ -211,10 +233,11 @@ pub async fn connect(db_url: &String, min_pool_size: Option<u32>, max_pool_size:
     let matches = db.collection::<Match>(Match::get_collection_name());
     let levels = db.collection::<Level>(Level::get_collection_name());
     let deaths = db.collection::<Death>(Death::get_collection_name());
+    let ip_identities = db.collection::<IpIdentity>(IpIdentity::get_collection_name());
 
     info!("Connected to database successfully.");
     Ok(Database { 
         mongo: db, tags, achievements, players, sessions, 
-        punishments, ranks, matches, levels, deaths 
+        punishments, ranks, matches, levels, deaths, ip_identities
     })
 }

src/database/models/ip_identity.rs

@@ -0,0 +1,65 @@
+use futures::stream::FuturesUnordered;
+use futures::StreamExt;
+use mars_api_rs_derive::IdentifiableDocument;
+use mars_api_rs_macro::IdentifiableDocument;
+use mongodb::Collection;
+use rocket::serde::{Deserialize, Serialize};
+use crate::database::{CollectionOwner, Database};
+use crate::database::models::player::Player;
+
+#[derive(Deserialize, Serialize, IdentifiableDocument, Clone)]
+pub struct IpIdentity {
+    #[id]
+    #[serde(rename = "_id")]
+    pub ip: String,
+    pub players: Vec<String>
+}
+
+impl IpIdentity {
+    pub async fn add_player_ip(database: &Database, ip: &String, player: &String) {
+        let collection = &database.ip_identities;
+        let ip_identity = match Self::get_ip_identity_by_ip(collection, ip).await {
+            Some(mut ip_identity) => {
+                if !ip_identity.players.contains(&player) {
+                    ip_identity.players.push(player.clone());
+                }
+                ip_identity
+            },
+            None => IpIdentity { ip: ip.clone(), players: vec![player.clone()] }
+        };
+        database.save(&ip_identity).await;
+    }
+
+    pub async fn find_players_for_ip(database: &Database, ip: &String) -> Vec<Player> {
+        let collection = &database.ip_identities;
+        let record = Self::get_ip_identity_by_ip(collection, ip).await;
+        match record {
+            Some(ip_identity) => {
+                let unordered_futures = FuturesUnordered::new();
+                for player in ip_identity.players.iter() {
+                    unordered_futures.push(
+                        Database::find_by_id(&database.players, player)
+                    );
+                }
+                let results : Vec<_> = unordered_futures.collect().await;
+                let players : Vec<_> = results.into_iter().filter_map(|r| r).collect();
+                players
+            }
+            None => Vec::new()
+        }
+    }
+
+    pub async fn get_ip_identity_by_ip(collection: &Collection<IpIdentity>, ip: &String) -> Option<IpIdentity> {
+        Database::find_by_id(collection, ip.as_str()).await
+    }
+}
+
+impl CollectionOwner<IpIdentity> for IpIdentity {
+    fn get_collection(database: &Database) -> &Collection<IpIdentity> {
+        &database.ip_identities
+    }
+
+    fn get_collection_name() -> &'static str {
+        "ip_identity"
+    }
+}

src/database/models/mod.rs

@@ -12,3 +12,4 @@ pub mod death;
 pub mod join_sound;
 pub mod server;
 pub mod achievement;
+pub mod ip_identity;

src/database/models/player.rs

@@ -85,7 +85,7 @@ pub struct PlayerStats {
     #[serde(default)]
     pub xp: u32,
     #[serde(default)]
-    pub server_playtime: u64,
+    pub server_playtime: i64,
     #[serde(default)]
     pub game_playtime: u64,
     #[serde(default)]

src/http/player/mod.rs

@@ -10,6 +10,7 @@ use sha2::{Sha256, Digest};
 
 use self::payloads::{PlayerPreLoginResponse, PlayerPreLoginResponder, PlayerLoginResponse, PlayerLogoutRequest, PlayerProfileResponder, PlayerProfileResponse, PlayerAltResponse};
 use std::{time::{SystemTime, UNIX_EPOCH}, collections::HashMap};
+use crate::database::models::ip_identity::IpIdentity;
 
 use super::punishment::payloads::PunishmentIssueRequest;
 
@@ -31,7 +32,8 @@ pub async fn prelogin(
     if let Some(mut returning_player) = player_optional {
         returning_player.name = data.player.name.clone();
         returning_player.name_lower = returning_player.name.to_lowercase();
-        if !returning_player.ips.contains(&ip) {
+        let new_ip = !returning_player.ips.contains(&ip);
+        if new_ip {
             returning_player.ips.push(ip.clone());
         };
 
@@ -60,6 +62,10 @@ pub async fn prelogin(
         };
         state.player_cache.set(&state.database, &returning_player.name, &returning_player, true).await;
         state.database.ensure_player_name_uniqueness(&data.player.name, &data.player.id).await;
+        // denormalize ip player relationship
+        if new_ip {
+            IpIdentity::add_player_ip(&state.database, &ip,  &returning_player.id).await;
+        }
 
         Ok(PlayerPreLoginResponder { 
             response: PlayerPreLoginResponse {
@@ -70,13 +76,13 @@ pub async fn prelogin(
             }
         })
     } else {
-        println!("Could not find player {} in database!", player_id);
+        debug!("Could not find player {} in database!", player_id);
         let time_millis : f64 = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis() as f64;
         let player = Player {
             id: data.player.id.clone(),
             name: data.player.name.clone(),
             name_lower: data.player.name.to_lowercase(),
-            ips: vec![ip],
+            ips: vec![ip.clone()],
             first_joined_at: time_millis,
             last_joined_at: time_millis,
             rank_ids: Vec::new(),
@@ -91,6 +97,7 @@ pub async fn prelogin(
 
         state.player_cache.set(&state.database, &player.name, &player, true).await;
         state.database.ensure_player_name_uniqueness(&data.player.name, &data.player.id).await;
+        IpIdentity::add_player_ip(&state.database, &ip, &data.player.id).await;
 
         Ok(PlayerPreLoginResponder {
             response: PlayerPreLoginResponse {
@@ -184,7 +191,7 @@ pub async fn logout(
 
     let time_millis : u64 = u64::try_from(SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis()).unwrap_or(u64::MAX);
     session.ended_at = Some(time_millis);
-    player.stats.server_playtime += data.playtime;
+    player.stats.server_playtime += (data.playtime as i64);
 
     state.leaderboards.server_playtime.increment(&player.id_name(), Some(u32::try_from(data.playtime).unwrap_or(u32::MAX))).await; // Will break in 2106
 
@@ -339,7 +346,10 @@ pub async fn lookup_player(
     let player_alts : Vec<PlayerAltResponse> = {
         let mut player_alts : Vec<PlayerAltResponse> = Vec::new();
         if alts {
+            let t1 = get_u64_time_millis();
             let fetched_alts = state.database.get_alts_for_player(&player).await;
+            let t2 = get_u64_time_millis();
+            debug!("Alt lookup for {} took {}ms", &player.name, (t2 - t1));
             let pun_tasks : Vec<_> = fetched_alts.iter().map(|alt| {
                 state.database.get_player_punishments(alt)
             }).collect();

src/http/server/mod.rs

@@ -47,7 +47,7 @@ async fn server_startup(
         sessions_to_write.push(hanging_session.to_owned());
 
         let mut cached_player = unwrap_helper::continue_default!(state.player_cache.get(&state.database, &hanging_session.player.name).await);
-        cached_player.stats.server_playtime += hanging_session.length().unwrap_or(0);
+        cached_player.stats.server_playtime += (hanging_session.length().unwrap_or(0) as i64);
         players_to_write.push(cached_player);
     }