move host tool added

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(())
+}