126 lines
3.7 KiB
Rust
126 lines
3.7 KiB
Rust
use std::{collections::HashSet, time::Duration};
|
||
|
||
use clap::{Parser};
|
||
use config::{Config, ConfigError, Environment};
|
||
use reqwest_middleware::ClientBuilder;
|
||
use reqwest_retry::{RetryTransientMiddleware, policies::ExponentialBackoff};
|
||
use serde::Deserialize;
|
||
use tokio::time::interval;
|
||
|
||
#[derive(Debug, Deserialize, Clone)]
|
||
pub struct Settings {
|
||
pub log_level: String,
|
||
pub interval: u64,
|
||
pub server_url: String,
|
||
}
|
||
|
||
#[derive(Parser, Debug)]
|
||
#[command(author, version, about = "PT 的硬盘状态汇报程序")]
|
||
struct Args {
|
||
#[arg(short, long)]
|
||
interval: Option<u64>,
|
||
|
||
#[arg(short = 'l', long)]
|
||
log_level: Option<String>,
|
||
}
|
||
|
||
impl Settings {
|
||
pub fn new() -> Result<Self, ConfigError> {
|
||
let s = Config::builder()
|
||
.set_default("log_level", "info")?
|
||
.set_default("interval", 300)?
|
||
.set_default("server_url", "")?
|
||
.add_source(Environment::with_prefix("PT").separator("__"))
|
||
.build()?;
|
||
|
||
s.try_deserialize()
|
||
}
|
||
}
|
||
|
||
#[tracing::instrument]
|
||
async fn task(settings: &Settings) -> anyhow::Result<()> {
|
||
let disks = sysinfo::Disks::new_with_refreshed_list();
|
||
|
||
let mut available = 0;
|
||
let mut seen_device = HashSet::new();
|
||
|
||
for disk in disks.list().iter().filter(|d| !d.is_removable()) {
|
||
let mount_path = disk.mount_point().to_string_lossy();
|
||
if mount_path.contains("/snap") || mount_path.contains("/docker") {
|
||
continue;
|
||
}
|
||
let dname = disk.name().to_string_lossy().to_string();
|
||
if !seen_device.insert(dname.clone()) {
|
||
continue;
|
||
}
|
||
if dname == "drivers" || dname == "none" {
|
||
continue;
|
||
}
|
||
let davailable = disk.available_space();
|
||
tracing::info!(disk=dname, available=davailable, "检查一块硬盘");
|
||
available += davailable;
|
||
}
|
||
|
||
let report = ((available >> 20) * 100) >> 10;
|
||
tracing::info!(report = report, "获取剩余硬盘空间（单位 .01GB）");
|
||
|
||
if settings.server_url.is_empty() {
|
||
tracing::warn!("没有配置服务器地址，不会上报给服务端");
|
||
} else {
|
||
let _rep = format!("{}", report);
|
||
let params = [
|
||
("status", "up"),
|
||
("msg", "OK"),
|
||
("ping", &_rep),
|
||
];
|
||
|
||
let retry_policy = ExponentialBackoff::builder().build_with_max_retries(3);
|
||
let client = ClientBuilder::new(reqwest::Client::new())
|
||
.with(RetryTransientMiddleware::new_with_policy(retry_policy))
|
||
.build();
|
||
|
||
let url = reqwest::Url::parse_with_params(&settings.server_url, ¶ms)?;
|
||
let urls: String = url.clone().into();
|
||
let resp = client.get(url).send().await?.text().await?;
|
||
tracing::info!(url=urls, resp=resp, "上报了硬盘可用空间");
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[tokio::main]
|
||
async fn main() -> anyhow::Result<()> {
|
||
let args = Args::parse();
|
||
let mut settings = Settings::new()?;
|
||
if let Some(i) = args.interval {
|
||
settings.interval = i;
|
||
}
|
||
if let Some(level) = args.log_level {
|
||
settings.log_level = level;
|
||
}
|
||
|
||
tracing_subscriber::fmt()
|
||
.with_env_filter(&settings.log_level)
|
||
.with_thread_ids(true)
|
||
.init();
|
||
|
||
tracing::info!("磁盘状态汇报程序启动");
|
||
let mut ticker = interval(Duration::from_secs(300));
|
||
|
||
loop {
|
||
tokio::select! {
|
||
_ = ticker.tick() => {
|
||
if let Err(e) = task(&settings).await {
|
||
tracing::error!(error = ?e, "执行任务时出错了！");
|
||
}
|
||
}
|
||
_ = tokio::signal::ctrl_c() => {
|
||
tracing::info!("收到 Ctrl+C 信号，程序退出");
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|