caveman/hunter/src/main.rs

94 lines
3.1 KiB
Rust

use std::time::Duration;
use tokio::time::timeout;
use cave::config::LoadConfig;
mod config;
mod scheduler;
mod feed;
mod worker;
mod redis_store;
use worker::Message;
#[tokio::main]
async fn main() {
cave::init::exit_on_panic();
cave::init::init_logger();
let config = config::Config::load();
cave::systemd::status("Starting redis client");
let redis_client = redis::Client::open(config.redis)
.expect("redis::Client");
let mut redis_man = redis::aio::ConnectionManager::new(redis_client).await
.expect("redis::aio::ConnectionManager");
cave::systemd::status("Starting scheduler");
let mut scheduler = scheduler::Scheduler::new();
for host in config.hosts.into_iter() {
scheduler.introduce(host).await;
}
for host in crate::redis_store::get_hosts(&mut redis_man).await.into_iter() {
scheduler.introduce(host).await;
}
cave::systemd::status("Starting HTTP client");
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.pool_max_idle_per_host(0)
.user_agent(concat!(
env!("CARGO_PKG_NAME"),
"/",
env!("CARGO_PKG_VERSION"),
))
.deflate(true)
.gzip(true)
.build()
.expect("reqwest::Client");
cave::systemd::status("Ready");
systemd::daemon::notify(false, [(systemd::daemon::STATE_READY, "1")].iter())
.unwrap();
let mut workers_active = 0usize;
let (message_tx, mut message_rx) = tokio::sync::mpsc::unbounded_channel();
loop {
log::trace!("{} workers active, queued {} of {}", workers_active, scheduler.queue_len(), scheduler.size());
cave::systemd::status(&format!("{} workers active, queued {} of {}", workers_active, scheduler.queue_len(), scheduler.size()));
let next_task = if workers_active < config.max_workers {
scheduler.dequeue()
} else {
Err(Duration::from_secs(5))
};
match next_task {
Err(duration) => {
let _ = timeout(duration, async {
let message = message_rx.recv().await.unwrap();
match message {
Message::Fetched { host, next_interval, new_posts } => {
workers_active -= 1;
scheduler.enqueue(host, new_posts > 0, next_interval);
}
Message::IntroduceHosts { hosts } => {
for host in hosts.into_iter() {
scheduler.introduce(host).await;
}
}
}
}).await;
}
Ok(host) => {
workers_active += 1;
worker::fetch_and_process(
message_tx.clone(),
redis_man.clone(),
client.clone(),
host
);
systemd::daemon::notify(false, [(systemd::daemon::STATE_WATCHDOG, "1")].iter())
.unwrap();
}
}
}
}