chore: refactor geoip, update deps

Signed-off-by: Henry Gressmann <[email protected]>

Author: explodingcamera

Cargo.lock

@@ -16,11 +16,10 @@ path="src/main.rs"
 
 [dependencies]
 # async/concurrency
+arc-swap="1.7"
 tokio={version="1.48", default-features=false, features=["macros", "rt-multi-thread", "signal"]}
 tokio-util={version="0.7", features=["io"]}
 futures-lite={version="2.6", default-features=false, features=["alloc"]}
-crossbeam-utils="0.8"
-crossbeam-channel="0.5"
 quick_cache={version="0.6", features=["parking_lot", "ahash"]}
 
 # encoding
@@ -63,7 +62,8 @@ reqwest={version="0.12", default-features=false, features=[
 ]}
 rustls={version="0.23", features=["aws_lc_rs"]}
 
-duckdb={version="1.4", features=["buildtime_bindgen", "bundled", "chrono", "r2d2"]} # database
+# database
+duckdb={version="1.4", features=["buildtime_bindgen", "bundled", "chrono", "r2d2"]}
 rusqlite={version="0.37", features=["bundled", "modern_sqlite", "chrono"]}
 r2d2={version="0.8", default-features=false}
 r2d2_sqlite="0.31"

Cargo.toml

@@ -1,9 +1,9 @@
 use std::sync::Arc;
 
+use arc_swap::ArcSwap;
 use chrono::{DateTime, Utc};
-use crossbeam_channel::Receiver;
-use crossbeam_utils::sync::ShardedLock;
 use eyre::{Result, bail};
+use std::sync::mpsc::Receiver;
 
 use crate::app::models::{Event, event_params};
 use crate::app::{DuckDBPool, EVENT_BATCH_INTERVAL, SqlitePool};
@@ -13,7 +13,7 @@ use crate::utils::hash::generate_salt;
 pub struct LiwanEvents {
     duckdb: DuckDBPool,
     sqlite: SqlitePool,
-    daily_salt: Arc<ShardedLock<(String, DateTime<Utc>)>>,
+    daily_salt: Arc<ArcSwap<(String, DateTime<Utc>)>>,
 }
 
 impl LiwanEvents {
@@ -24,15 +24,12 @@ impl LiwanEvents {
                 Ok((row.get(0)?, row.get(1)?))
             })?
         };
-        Ok(Self { duckdb, sqlite, daily_salt: ShardedLock::new(daily_salt).into() })
+        Ok(Self { duckdb, sqlite, daily_salt: ArcSwap::new(daily_salt.into()).into() })
     }
 
     /// Get the daily salt, generating a new one if the current one is older than 24 hours
     pub fn get_salt(&self) -> Result<String> {
-        let (salt, updated_at) = {
-            let salt = self.daily_salt.read().map_err(|_| eyre::eyre!("Failed to acquire read lock"))?;
-            salt.clone()
-        };
+        let (salt, updated_at) = &**self.daily_salt.load();
 
         // if the salt is older than 24 hours, replace it with a new one (utils::generate_salt)
         if (Utc::now() - updated_at) > chrono::Duration::hours(24) {
@@ -41,15 +38,11 @@ impl LiwanEvents {
             let now = Utc::now();
             let conn = self.sqlite.get()?;
             conn.execute("update salts set salt = ?, updated_at = ? where id = 1", rusqlite::params![&new_salt, now])?;
-
-            if let Ok(mut daily_salt) = self.daily_salt.try_write() {
-                daily_salt.0.clone_from(&new_salt);
-                daily_salt.1 = now;
-                return Ok(new_salt);
-            }
+            self.daily_salt.store((new_salt.clone(), now).into());
+            Ok(new_salt)
+        } else {
+            Ok(salt.clone())
         }
-
-        Ok(salt)
     }
 
     /// Append events in batch

data/licenses-cargo.json

@@ -1,14 +1,13 @@
 #![allow(dead_code)]
 
-use std::collections::HashMap;
 use std::io::{self};
 use std::net::IpAddr;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 
 use crate::app::SqlitePool;
-use crossbeam_utils::sync::ShardedLock;
+use arc_swap::ArcSwapOption;
 use eyre::{Context, OptionExt, Result};
 use futures_lite::StreamExt;
 use md5::{Digest, Md5};
@@ -19,6 +18,7 @@ const BASE_URL: &str = "https://updates.maxmind.com";
 const METADATA_ENDPOINT: &str = "/geoip/updates/metadata?edition_id=";
 const DOWNLOAD_ENDPOINT: &str = "/geoip/databases/";
 
+#[derive(Default)]
 pub struct LookupResult {
     pub city: Option<String>,
     pub country_code: Option<String>,
@@ -27,30 +27,26 @@ pub struct LookupResult {
 #[derive(Clone)]
 pub struct LiwanGeoIP {
     pool: SqlitePool,
-    reader: Arc<ShardedLock<Option<maxminddb::Reader<Vec<u8>>>>>,
+    reader: Arc<ArcSwapOption<maxminddb::Reader<Vec<u8>>>>,
 
     downloading: Arc<AtomicBool>,
-    config: crate::config::Config,
     geoip: crate::config::GeoIpConfig,
     path: PathBuf,
 }
 
 impl LiwanGeoIP {
-    pub fn try_new(config: crate::config::Config, pool: SqlitePool) -> Result<Option<Self>> {
-        let Some(geoip) = &config.geoip else {
-            tracing::trace!("GeoIP support disabled, skipping...");
-            return Ok(None);
-        };
-
+    pub fn try_new(config: crate::config::Config, pool: SqlitePool) -> Result<Self> {
+        let geoip = config.geoip;
         if geoip.maxmind_account_id.is_none() && geoip.maxmind_license_key.is_none() && geoip.maxmind_db_path.is_none()
         {
             tracing::trace!("GeoIP support disabled, skipping...");
-            return Ok(None);
+            return Ok(Self::noop(pool));
         }
 
-        let edition = geoip.maxmind_edition.as_deref().unwrap_or("GeoLite2-City");
+        let edition = &geoip.maxmind_edition;
         let default_path = PathBuf::from(config.data_dir.clone()).join(format!("./geoip/{edition}.mmdb"));
         let path = geoip.maxmind_db_path.as_ref().map_or(default_path, PathBuf::from);
+
         if let Some(parent) = path.parent() {
             std::fs::create_dir_all(parent)?;
         }
@@ -60,28 +56,38 @@ impl LiwanGeoIP {
         }
 
         tracing::info!(database = geoip.maxmind_db_path, "GeoIP support enabled, loading database");
-        let reader = if path.exists() {
-            Some(maxminddb::Reader::open_readfile(path.clone()).expect("Failed to open GeoIP database file"))
-        } else {
-            None
-        };
 
-        Ok(Some(Self {
-            geoip: geoip.clone(),
-            config,
+        let reader = path.exists().then(|| {
+            maxminddb::Reader::open_readfile(path.clone()).expect("Failed to open GeoIP database file").into()
+        });
+
+        Ok(Self { geoip, pool, reader: ArcSwapOption::new(reader).into(), path, downloading: Default::default() })
+    }
+
+    fn is_enabled(&self) -> bool {
+        self.reader.load().is_some() || self.downloading.load(Ordering::Acquire)
+    }
+
+    fn noop(pool: SqlitePool) -> Self {
+        Self {
+            geoip: Default::default(),
             pool,
-            reader: Arc::new(ShardedLock::new(reader)),
-            path,
-            downloading: Arc::new(AtomicBool::new(false)),
-        }))
+            reader: ArcSwapOption::new(None).into(),
+            downloading: Default::default(),
+            path: PathBuf::new(),
+        }
     }
 
     // Lookup the IP address in the GeoIP database
     pub fn lookup(&self, ip: &IpAddr) -> Result<LookupResult> {
-        let reader = self.reader.read().map_err(|_| eyre::eyre!("Failed to acquire GeoIP reader lock"))?;
-        let reader = reader.as_ref().ok_or_eyre("GeoIP database not found")?;
-        let lookup: maxminddb::geoip2::City =
-            reader.lookup(*ip)?.ok_or_else(|| eyre::eyre!("No data found for IP address"))?;
+        let Some(reader) = &*self.reader.load() else {
+            return Ok(Default::default());
+        };
+
+        let lookup = reader
+            .lookup::<maxminddb::geoip2::City>(*ip)?
+            .ok_or_else(|| eyre::eyre!("No data found for IP address"))?;
+
         let city = lookup.city.and_then(|city| city.names.and_then(|names| names.get("en").map(|s| (*s).to_string())));
         let country_code = lookup.country.and_then(|country| country.iso_code.map(ToString::to_string));
         Ok(LookupResult { city, country_code })
@@ -93,16 +99,17 @@ impl LiwanGeoIP {
             return Ok(());
         }
 
-        let maxmind_edition = self.geoip.maxmind_edition.clone().ok_or_eyre("MaxMind edition not found")?;
-        let maxmind_account_id = self.geoip.maxmind_account_id.clone().ok_or_eyre("MaxMind account ID not found")?;
-        let maxmind_license_key = self.geoip.maxmind_license_key.clone().ok_or_eyre("MaxMind license key not found")?;
+        let maxmind_edition = &self.geoip.maxmind_edition;
+        let maxmind_account_id = self.geoip.maxmind_account_id.as_deref().ok_or_eyre("MaxMind account ID not found")?;
+        let maxmind_license_key =
+            self.geoip.maxmind_license_key.as_deref().ok_or_eyre("MaxMind license key not found")?;
 
         let db_exists = self.path.exists();
         let db_md5 = if db_exists { file_md5(&self.path)? } else { String::new() };
 
-        let mut update = false;
+        let mut update = !db_exists;
         if db_exists {
-            match get_latest_md5(&maxmind_edition, &maxmind_account_id, &maxmind_license_key).await {
+            match get_latest_md5(maxmind_edition, maxmind_account_id, maxmind_license_key).await {
                 Ok(latest_md5) => {
                     if latest_md5 != db_md5 {
                         tracing::info!("GeoIP database outdated, downloading...");
@@ -115,11 +122,10 @@ impl LiwanGeoIP {
             };
         } else {
             tracing::info!("GeoIP database doesn't exist, attempting to download...");
-            update = true;
         }
 
         if update {
-            let file = match download_maxmind_db(&maxmind_edition, &maxmind_account_id, &maxmind_license_key).await {
+            let file = match download_maxmind_db(maxmind_edition, maxmind_account_id, maxmind_license_key).await {
                 Ok(file) => file,
                 Err(e) => {
                     tracing::warn!(error = ?e, "Failed to download GeoIP database, skipping update");
@@ -129,18 +135,15 @@ impl LiwanGeoIP {
             };
 
             // close the current reader to free up the file
-            {
-                let mut reader = self.reader.write().unwrap();
-                reader.take();
-            }
+            self.reader.swap(None);
 
             // move the downloaded file to the correct path
             std::fs::copy(&file, &self.path)?;
             std::fs::remove_file(file)?;
 
             // open the new reader
             let reader = maxminddb::Reader::open_readfile(self.path.clone())?;
-            *self.reader.write().unwrap() = Some(reader);
+            self.reader.store(Some(reader.into()));
 
             let path = std::fs::canonicalize(&self.path)?;
             tracing::info!(path = ?path, "GeoIP database updated successfully");
@@ -151,8 +154,10 @@ impl LiwanGeoIP {
     }
 }
 
-pub fn keep_updated(geoip: Option<LiwanGeoIP>) {
-    let Some(geoip) = geoip else { return };
+pub fn keep_updated(geoip: LiwanGeoIP) {
+    if !geoip.is_enabled() {
+        return;
+    }
 
     tokio::task::spawn(async move {
         if let Err(e) = geoip.check_for_updates().await {
@@ -182,7 +187,7 @@ async fn get_latest_md5(edition: &str, account_id: &str, license_key: &str) -> R
         .basic_auth(account_id, Some(license_key))
         .send()
         .await?
-        .json::<HashMap<String, Vec<HashMap<String, String>>>>()
+        .json::<ahash::HashMap<String, Vec<ahash::HashMap<String, String>>>>()
         .await?;
 
     Ok(response

data/licenses-npm.json

@@ -1,42 +1,34 @@
-use std::sync::Arc;
-
-use crossbeam_utils::sync::ShardedLock;
+use arc_swap::ArcSwapOption;
 use eyre::Result;
+use std::sync::Arc;
 
 use crate::{app::SqlitePool, utils::hash::onboarding_token};
 
 #[derive(Clone)]
 pub struct LiwanOnboarding {
-    token: Arc<ShardedLock<Option<String>>>,
+    token: Arc<ArcSwapOption<String>>,
 }
 
 impl LiwanOnboarding {
     pub fn try_new(pool: &SqlitePool) -> Result<Self> {
         let onboarding = {
             tracing::debug!("Checking if an onboarding token needs to be generated");
             let conn = pool.get()?;
-            let mut stmt = conn.prepare("select 1 from users limit 1")?;
-            ShardedLock::new(if stmt.exists([])? { None } else { Some(onboarding_token()) })
+            let onboarded = conn.prepare("select 1 from users limit 1")?.exists([])?;
+            ArcSwapOption::new(onboarded.then(|| onboarding_token().into()))
         };
 
         Ok(Self { token: onboarding.into() })
     }
 
     /// Get the onboarding token, if it exists
     pub fn token(&self) -> Result<Option<String>> {
-        Ok(self
-            .token
-            .read()
-            .map_err(|_| eyre::eyre!("Failed to acquire onboarding token read lock"))?
-            .as_ref()
-            .cloned())
+        Ok((self.token.load()).as_ref().map(|v| (**v).clone()))
     }
 
     /// Clear the onboarding token to prevent it from being used again
     pub fn clear(&self) -> Result<()> {
-        let mut onboarding =
-            self.token.write().map_err(|_| eyre::eyre!("Failed to acquire onboarding token write lock"))?;
-        *onboarding = None;
+        self.token.store(None);
         Ok(())
     }
 }

src/app/core/events.rs

@@ -10,7 +10,7 @@ use std::path::PathBuf;
 
 pub(super) fn init_duckdb(
     path: &PathBuf,
-    duckdb_config: Option<DuckdbConfig>,
+    duckdb_config: DuckdbConfig,
     mut migrations_runner: Runner,
 ) -> Result<r2d2::Pool<DuckdbConnectionManager>> {
     let mut flags = duckdb::Config::default()
@@ -19,14 +19,12 @@ pub(super) fn init_duckdb(
         .with("enable_fsst_vectors", "true")?
         .with("allocator_background_threads", "true")?;
 
-    if let Some(duckdb_config) = duckdb_config {
-        if let Some(memory_limit) = duckdb_config.memory_limit {
-            flags = flags.max_memory(&memory_limit)?;
-        }
+    if let Some(memory_limit) = duckdb_config.memory_limit {
+        flags = flags.max_memory(&memory_limit)?;
+    }
 
-        if let Some(threads) = duckdb_config.threads {
-            flags = flags.threads(threads.get().into())?;
-        }
+    if let Some(threads) = duckdb_config.threads {
+        flags = flags.threads(threads.get().into())?;
     }
 
     let conn = DuckdbConnectionManager::file_with_flags(path, flags).map_err(|e| {

src/app/core/geoip.rs

@@ -25,7 +25,7 @@ pub struct Liwan {
     pub onboarding: core::onboarding::LiwanOnboarding,
     pub entities: core::entities::LiwanEntities,
     pub projects: core::projects::LiwanProjects,
-    pub geoip: Option<core::geoip::LiwanGeoIP>,
+    pub geoip: core::geoip::LiwanGeoIP,
 
     pub config: Config,
 }

src/app/core/onboarding.rs

@@ -1,6 +1,6 @@
 use crate::{
     app::{Liwan, models::UserRole},
-    config::{Config, DEFAULT_CONFIG},
+    config::{Config, DEFAULT_CONFIG, GeoIpConfig},
 };
 use argh::FromArgs;
 use eyre::Result;
@@ -86,7 +86,7 @@ pub struct AddUser {
 }
 
 pub fn handle_command(mut config: Config, cmd: Command) -> Result<()> {
-    config.geoip = None; // disable GeoIP support in CLI
+    config.geoip = GeoIpConfig::default(); // disable GeoIP in CLI commands
 
     match cmd {
         Command::UpdatePassword(update) => {