b33171b487 2018-06-01 arcade@b1t.na: extern crate config;
b33171b487 2018-06-01 arcade@b1t.na: extern crate postgres;
b33171b487 2018-06-01 arcade@b1t.na: extern crate redis;
b33171b487 2018-06-01 arcade@b1t.na: extern crate regex;
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: use postgres::{Connection, TlsMode};
b33171b487 2018-06-01 arcade@b1t.na: use redis::Commands;
b33171b487 2018-06-01 arcade@b1t.na: use regex::Regex;
b33171b487 2018-06-01 arcade@b1t.na: use std::collections::HashMap;
b33171b487 2018-06-01 arcade@b1t.na: use std::time::Duration;
b33171b487 2018-06-01 arcade@b1t.na: use std::thread;
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: fn main() {
b33171b487 2018-06-01 arcade@b1t.na: loop {
b33171b487 2018-06-01 arcade@b1t.na: let mut settings = config::Config::default();
b33171b487 2018-06-01 arcade@b1t.na: settings.merge(config::File::with_name("offload")).expect("Can't read configuration file");
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: let reddb = redis::Client::open("redis://127.0.0.1/").expect("Can't connect to the database");
b33171b487 2018-06-01 arcade@b1t.na: let red = reddb.get_connection().expect("Can't initialize new connection");
b33171b487 2018-06-01 arcade@b1t.na:
f721596292 2018-10-28 arcade: let re = Regex::new("^([0-9A-Z]+)[|_]([a-zA-Z./0-9-_#]+)[|_]([0-9.]+)$").expect("Can't parse regexp");
b33171b487 2018-06-01 arcade@b1t.na: let conn = Connection::connect(settings.get_str("pg").expect("Postgres connection absent in config"), TlsMode::None).expect("Can't connect to postgres");
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: let schemas = settings.get_array("schemas").expect("Schema list not found in config").into_iter().map(|value| config::Value::into_str(value).expect("We require string here"));
b33171b487 2018-06-01 arcade@b1t.na: for schema in schemas {
b33171b487 2018-06-01 arcade@b1t.na: let data_key = schema.to_owned() + "_counter_pending";
b33171b487 2018-06-01 arcade@b1t.na: let cache_key = schema.to_owned() + "_counter_pending_now";
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: if red.exists(&cache_key).expect("Can't query cache existance") {
b33171b487 2018-06-01 arcade@b1t.na: } else if red.exists(&data_key).expect("Can't query data existance") {
b33171b487 2018-06-01 arcade@b1t.na: let _ : bool = red.rename_nx(&data_key, &cache_key).expect("Can't readd unflushed hash");
b33171b487 2018-06-01 arcade@b1t.na: } else {
b33171b487 2018-06-01 arcade@b1t.na: continue;
b33171b487 2018-06-01 arcade@b1t.na: }
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: let stats : HashMap<String, i16> = red.hgetall(&cache_key).expect("Can't query for stored stat keys");
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: //println!("# {:?}", stats);
b33171b487 2018-06-01 arcade@b1t.na:
b33171b487 2018-06-01 arcade@b1t.na: let trans = conn.transaction().expect("Can't start transaction");
b33171b487 2018-06-01 arcade@b1t.na: let stmt = trans.prepare(&format!("select {}.merge_counter($1::text, $2::text, ($3::text)::inet, $4::smallint);", &schema)).expect("Can't prepare statement");
b33171b487 2018-06-01 arcade@b1t.na: for (client, count) in stats {
b33171b487 2018-06-01 arcade@b1t.na: //println!("# {:?}: {:?}", &client, &count);
f721596292 2018-10-28 arcade: let cap = re.captures(&client).expect(&format!("Client match failed: {}", &client));
b33171b487 2018-06-01 arcade@b1t.na: //let addr = IpAddr::V4(cap[3].parse().expect("Can't parse IP"));
b33171b487 2018-06-01 arcade@b1t.na: //println!("insert into x values({:?}, {:?}, {:?}, {:?});", &cap[1], &cap[2], &cap[3], &count);
b33171b487 2018-06-01 arcade@b1t.na: stmt.execute(&[&cap[1].to_string(), &cap[2].to_string(), &cap[3].to_string(), &count]).expect("Can't execute prepared");
b33171b487 2018-06-01 arcade@b1t.na: //trans.query(&format!("select {}.merge_counter($1::text, $2::text, $3::inet, $4::smallint);", &schema), &[&cap[1].to_string(), &cap[2].to_string(), &cap[3].to_string(), &count]).expect("Can't prepare statement");
b33171b487 2018-06-01 arcade@b1t.na: }
b33171b487 2018-06-01 arcade@b1t.na: trans.commit().expect("Can't commit transaction");
b33171b487 2018-06-01 arcade@b1t.na: red.del(cache_key).expect("Can't remove stale key")
b33171b487 2018-06-01 arcade@b1t.na: }
b33171b487 2018-06-01 arcade@b1t.na: thread::sleep(Duration::new(settings.get("delay").expect("Delay specification absent in config"), 0));
b33171b487 2018-06-01 arcade@b1t.na: }
b33171b487 2018-06-01 arcade@b1t.na: }