diff options
| -rw-r--r-- | src/main.rs | 49 | ||||
| -rw-r--r-- | src/state.rs | 10 | ||||
| -rw-r--r-- | src/tool.rs | 51 |
3 files changed, 85 insertions, 25 deletions
diff --git a/src/main.rs b/src/main.rs index 986b7f1..c9b80ab 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,9 +3,11 @@ pub mod embed; pub mod error; pub mod info; pub mod state; +pub mod tool; + +use anyhow::{anyhow, Context}; use embed::*; use info::*; - use markup::Render; use rocket::{ catch, catchers, @@ -15,24 +17,29 @@ use rocket::{ response::{self, Responder}, routes, shield::{self, Shield}, + tokio::{self}, Request, Response, }; use state::{AdInfo, Config, Logic}; -use std::{io::Cursor, net::IpAddr}; +use std::{env::args, fs::read_to_string, io::Cursor, net::IpAddr, sync::Arc}; +use tool::tool_main; -#[rocket::main] -async fn main() { +fn main() -> anyhow::Result<()> { env_logger::init_from_env("LOG"); + let mut args = args().skip(1); + let config = args.next().expect("first arg needs to be the config"); + let config = read_to_string(config).context(anyhow!("could not read config"))?; + let mut config: Config = toml::from_str(config.as_str()).context(anyhow!("config invalid"))?; - let config = std::env::args() - .nth(1) - .expect("first arg needs to be the config"); - let config = rocket::tokio::fs::read_to_string(config) - .await - .expect("could not read config"); - let mut config: Config = toml::from_str(config.as_str()).expect("config invalid"); + if let Some(action) = args.next() { + return tool_main(config, &action, args.collect::<Vec<_>>()); + } - for entry in config.ad_dir.read_dir().expect("cannot read ad directory") { + for entry in config + .ad_dir + .read_dir() + .context(anyhow!("cannot read ad directory"))? + { if let Ok(entry) = entry { if entry .path() @@ -45,12 +52,7 @@ async fn main() { let path = entry.path(); let imname = path.file_stem().unwrap().to_str().unwrap(); let basename = imname.split_once(".").unwrap().0; - let info: AdInfo = toml::from_str( - &rocket::tokio::fs::read_to_string(entry.path()) - .await - .unwrap(), - ) - .unwrap(); + let info: AdInfo = toml::from_str(&read_to_string(entry.path()).unwrap()).unwrap(); config.ads.insert( basename.to_string(), AdInfo { @@ -64,6 +66,14 @@ async fn main() { let state = Logic::new(config); + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build()? + .block_on(inner_main(state)); + + Ok(()) +} +async fn inner_main(state: Arc<Logic>) { let _ = rocket::build() .configure(rocket::Config { port: state.config.port, @@ -78,8 +88,7 @@ async fn main() { .mount("/", routes![r_index, r_embed, r_style, r_image, r_iptest]) .register("/", catchers![r_catch]) .launch() - .await - .unwrap(); + .await; } #[get("/myip")] diff --git a/src/state.rs b/src/state.rs index 4c91c6a..72a8aa4 100644 --- a/src/state.rs +++ b/src/state.rs @@ -24,7 +24,7 @@ pub struct Config { impression_weight_falloff: f64, leaderboard_weight_threshold: f64, pub ad_dir: PathBuf, - database_path: PathBuf, + pub database_path: PathBuf, pub port: u16, #[serde(default)] pub ads: HashMap<String, AdInfo>, @@ -43,10 +43,10 @@ struct ImpressionEvent { address_hash: u64, } -static T_TOTAL: TableDefinition<'static, (), u64> = TableDefinition::new("t"); -static T_IMPRESSIONS_RAW: TableDefinition<'static, &str, u64> = TableDefinition::new("ir"); -static T_IMPRESSIONS_WEIGHTED: TableDefinition<'static, &str, f64> = TableDefinition::new("iw"); -static T_IMPRESSIONS_ADS: TableDefinition<'static, &str, u64> = TableDefinition::new("ia"); +pub static T_TOTAL: TableDefinition<'static, (), u64> = TableDefinition::new("t"); +pub static T_IMPRESSIONS_RAW: TableDefinition<'static, &str, u64> = TableDefinition::new("ir"); +pub static T_IMPRESSIONS_WEIGHTED: TableDefinition<'static, &str, f64> = TableDefinition::new("iw"); +pub static T_IMPRESSIONS_ADS: TableDefinition<'static, &str, u64> = TableDefinition::new("ia"); impl Logic { pub fn new(config: Config) -> Arc<Self> { diff --git a/src/tool.rs b/src/tool.rs new file mode 100644 index 0000000..c431c09 --- /dev/null +++ b/src/tool.rs @@ -0,0 +1,51 @@ +use crate::state::{Config, T_IMPRESSIONS_RAW, T_IMPRESSIONS_WEIGHTED}; +use anyhow::{anyhow, bail, Result}; +use redb::{Database, ReadableTable}; + +pub fn tool_main(config: Config, action: &str, args: Vec<String>) -> Result<()> { + match action { + "move_host" => { + let oldhost = args.get(0).ok_or(anyhow!("2nd arg is old host"))?; + let newhost = args.get(1).ok_or(anyhow!("3rd arg is new host"))?; + move_host(config, oldhost, newhost)?; + } + _ => bail!("unknown action"), + } + Ok(()) +} + +pub fn move_host(config: Config, oldhost: &str, newhost: &str) -> Result<()> { + let db = Database::open(config.database_path)?; + + let txn = db.begin_write()?; + { + let mut t_impressions = txn.open_table(T_IMPRESSIONS_RAW)?; + let mut t_weighted = txn.open_table(T_IMPRESSIONS_WEIGHTED)?; + + let old_imp = t_impressions + .remove(oldhost)? + .map(|g| g.value()) + .unwrap_or_default(); + let old_wei = t_weighted + .remove(oldhost)? + .map(|g| g.value()) + .unwrap_or_default(); + + let new_imp = t_impressions + .get(newhost)? + .map(|g| g.value()) + .unwrap_or_default(); + let new_wei = t_weighted + .get(newhost)? + .map(|g| g.value()) + .unwrap_or_default(); + + println!("impressions old={old_imp} new={new_imp}"); + println!("weighted old={old_wei} new={new_wei}"); + t_impressions.insert(newhost, old_imp + new_imp)?; + t_weighted.insert(newhost, old_wei + new_wei)?; + } + txn.commit()?; + + Ok(()) +} |