146 lines
4.1 KiB
Rust
146 lines
4.1 KiB
Rust
use std::{collections::HashMap, time::Duration};
|
|
|
|
use chrono::{Datelike, Timelike};
|
|
use tokio::task::JoinHandle;
|
|
|
|
use tracing::{Instrument, Level};
|
|
use tracing_subscriber::FmtSubscriber;
|
|
|
|
use tzfile::{ArcTz, Tz};
|
|
|
|
use serde::Deserialize;
|
|
use toml;
|
|
|
|
mod guac;
|
|
mod shotter;
|
|
|
|
fn duration_until_next_minute() -> Duration {
|
|
let now = chrono::Utc::now();
|
|
let interval = (60 - now.second()) * 1000;
|
|
|
|
// not sure if this ever triggers
|
|
if interval == 0 {
|
|
return Duration::from_secs(59);
|
|
}
|
|
|
|
return Duration::from_millis(interval as u64);
|
|
}
|
|
|
|
#[derive(Deserialize, Debug, Clone)]
|
|
struct Node {
|
|
url: String,
|
|
origin: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct Config {
|
|
root_path: std::path::PathBuf,
|
|
|
|
webp_quality: f32,
|
|
|
|
vms: HashMap<String, Node>,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> anyhow::Result<()> {
|
|
// Init tracing
|
|
#[cfg(debug_assertions)]
|
|
let subscriber = FmtSubscriber::builder()
|
|
.with_max_level(Level::TRACE)
|
|
.compact()
|
|
.finish();
|
|
|
|
#[cfg(not(debug_assertions))]
|
|
let subscriber = FmtSubscriber::builder()
|
|
.with_max_level(Level::INFO)
|
|
.compact()
|
|
.finish();
|
|
|
|
tracing::subscriber::set_global_default(subscriber)?;
|
|
|
|
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
|
|
|
let config: Config = toml::from_str(std::fs::read_to_string("./config.toml")?.as_str())?;
|
|
|
|
// Essentially this is meant to be a sentinel for "SCREENSHOT NOW!".
|
|
// None will stop the task immediately, so basically this is a really bad way of implementing cancellation
|
|
// (that isn't used currently)
|
|
let (tx, _) = tokio::sync::broadcast::channel::<Option<()>>(10);
|
|
|
|
let gmt_timezone = ArcTz::new(Tz::named("GMT")?);
|
|
|
|
for (id, node) in config.vms.iter() {
|
|
tracing::info!("Adding node {id} : {:?}", node);
|
|
|
|
let mut _rx = tx.subscribe();
|
|
|
|
// These clones are scary, but they're only done once.
|
|
// Additionally, in the case of the timezone, no actual clone is performed.
|
|
let id_clone = id.clone();
|
|
let node_clone = node.clone();
|
|
let root = config.root_path.clone();
|
|
let tz = gmt_timezone.clone();
|
|
|
|
// Spawn the per-node task.
|
|
let _: JoinHandle<anyhow::Result<()>> = tokio::spawn(async move {
|
|
while let Some(_) = _rx.recv().await? {
|
|
let span = tracing::span!(Level::INFO, "node screenshot", node = id_clone.as_str());
|
|
let now = chrono::Utc::now().with_timezone(&tz);
|
|
|
|
// start with /yyyy-mm-dd
|
|
let date_path = root.join(format!(
|
|
"{:04}-{:02}-{:02}",
|
|
now.year(),
|
|
now.month(),
|
|
now.day()
|
|
));
|
|
|
|
// then /yyyy-mm-dd/[node]
|
|
let node_path = date_path.join(&id_clone);
|
|
|
|
if !node_path.exists() {
|
|
std::fs::create_dir_all(&node_path)?;
|
|
}
|
|
|
|
// add target webp path finally
|
|
let file_path = node_path.join(format!(
|
|
"{:02}-{:02}-{:02}.webp",
|
|
now.hour(),
|
|
now.minute(),
|
|
now.second()
|
|
));
|
|
|
|
// do it!
|
|
match shotter::take_one_screenshot(
|
|
&node_clone.url,
|
|
&node_clone.origin,
|
|
&id_clone,
|
|
file_path.clone(),
|
|
config.webp_quality,
|
|
)
|
|
.instrument(span)
|
|
.await
|
|
{
|
|
Ok(_) => {}
|
|
|
|
Err(error) => {
|
|
// FIXME: On WebSocket errors, should we just try again?
|
|
tracing::error!("Error taking screenshot: {error}");
|
|
}
|
|
};
|
|
}
|
|
|
|
Ok(())
|
|
});
|
|
}
|
|
|
|
loop {
|
|
// Request a screenshot
|
|
tx.send(Some(()))?;
|
|
|
|
// Wait for the next :xx:00 minute
|
|
let dur = duration_until_next_minute();
|
|
tracing::info!("Waiting {:?} to take next screenshot", dur);
|
|
tokio::time::sleep(dur).await;
|
|
}
|
|
}
|