use anyhow::Context; use redb::{Database, ReadableTable, TableDefinition}; use rocket::tokio::sync::Mutex; use serde::Deserialize; use std::{collections::HashMap, net::IpAddr, path::PathBuf}; #[derive(Deserialize)] pub struct AdInfo { pub image: PathBuf, pub target: String, } #[derive(Deserialize)] pub struct Config { bloom_filter_size: usize, impression_weight_falloff: f64, pub image_base: PathBuf, database_path: PathBuf, pub port: u16, pub ads: HashMap, } pub struct Logic { pub config: Config, database: Database, impressions_by_addr: Mutex>, pub ad_keys: Vec, } static T_TOTAL: TableDefinition<'static, (), u128> = TableDefinition::new("t"); static T_IMPRESSIONS_RAW: TableDefinition<'static, &str, u128> = TableDefinition::new("ir"); static T_IMPRESSIONS_WEIGHTED: TableDefinition<'static, &str, f64> = TableDefinition::new("iw"); static T_ADS: TableDefinition<'static, &str, u128> = TableDefinition::new("a"); impl Logic { pub fn new(config: Config) -> Self { Self { impressions_by_addr: vec![0; config.bloom_filter_size].into(), database: { let db = Database::create(&config.database_path).expect("database open failed"); { let txn = db.begin_write().unwrap(); txn.open_table(T_IMPRESSIONS_RAW).unwrap(); txn.open_table(T_IMPRESSIONS_WEIGHTED).unwrap(); txn.open_table(T_ADS).unwrap(); } db }, ad_keys: config.ads.keys().map(String::from).collect(), config, } } pub async fn register_impression( &self, site: &str, adid: &str, address: IpAddr, ) -> anyhow::Result<()> { let num_impressions = { let mut bloom = self.impressions_by_addr.lock().await; let ind = (xorshift(xorshift(xorshift( match address { IpAddr::V4(a) => a.to_ipv6_mapped(), IpAddr::V6(a) => a, } .to_bits() as u64, ))) % bloom.len() as u64) as usize; bloom[ind] = bloom[ind].saturating_add(1); bloom[ind] } as f64; let weight = self.config.impression_weight_falloff.powf(num_impressions); let txn = self.database.begin_write().context("database failure")?; { let mut raw = txn.open_table(T_TOTAL)?; let v = raw.get(())?.map(|g| g.value()).unwrap_or_default(); raw.insert((), v + 1)?; } { let mut raw = txn.open_table(T_IMPRESSIONS_RAW)?; let v = raw.get(site)?.map(|g| g.value()).unwrap_or_default(); raw.insert(site, v + 1)?; } { let mut raw = txn.open_table(T_IMPRESSIONS_WEIGHTED)?; let v = raw.get(site)?.map(|g| g.value()).unwrap_or_default(); raw.insert(site, v + weight)?; } { let mut raw = txn.open_table(T_ADS)?; let v = raw.get(site)?.map(|g| g.value()).unwrap_or_default(); raw.insert(adid, v + 1)?; } txn.commit().context("database failure")?; Ok(()) } } fn xorshift(mut x: u64) -> u64 { x ^= x << 13; x ^= x >> 7; x ^= x << 17; x }