use clap::Args;
use sc_client_db::DatabaseSource;
use sp_core::traits::SpawnEssentialNamed;
use std::{
io,
path::{Path, PathBuf},
time::Duration,
};
const LOG_TARGET: &str = "storage-monitor";
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("IO Error")]
IOError(#[from] io::Error),
#[error("Out of storage space: available {0}MiB, required {1}MiB")]
StorageOutOfSpace(u64, u64),
}
#[derive(Default, Debug, Clone, Args)]
pub struct StorageMonitorParams {
#[arg(long = "db-storage-threshold", value_name = "MiB", default_value_t = 1024)]
pub threshold: u64,
#[arg(long = "db-storage-polling-period", value_name = "SECONDS", default_value_t = 5, value_parser = clap::value_parser!(u32).range(1..))]
pub polling_period: u32,
}
pub struct StorageMonitorService {
path: PathBuf,
threshold: u64,
polling_period: Duration,
}
impl StorageMonitorService {
pub fn try_spawn(
parameters: StorageMonitorParams,
database: DatabaseSource,
spawner: &impl SpawnEssentialNamed,
) -> Result<()> {
Ok(match (parameters.threshold, database.path()) {
(0, _) => {
log::info!(
target: LOG_TARGET,
"StorageMonitorService: threshold `0` given, storage monitoring disabled",
);
},
(_, None) => {
log::warn!(
target: LOG_TARGET,
"StorageMonitorService: no database path to observe",
);
},
(threshold, Some(path)) => {
log::debug!(
target: LOG_TARGET,
"Initializing StorageMonitorService for db path: {path:?}",
);
Self::check_free_space(&path, threshold)?;
let storage_monitor_service = StorageMonitorService {
path: path.to_path_buf(),
threshold,
polling_period: Duration::from_secs(parameters.polling_period.into()),
};
spawner.spawn_essential(
"storage-monitor",
None,
Box::pin(storage_monitor_service.run()),
);
},
})
}
async fn run(self) {
loop {
tokio::time::sleep(self.polling_period).await;
if Self::check_free_space(&self.path, self.threshold).is_err() {
break
};
}
}
fn free_space(path: &Path) -> Result<u64> {
Ok(fs4::available_space(path).map(|s| s / 1024 / 1024)?)
}
fn check_free_space(path: &Path, threshold: u64) -> Result<()> {
match StorageMonitorService::free_space(path) {
Ok(available_space) => {
log::trace!(
target: LOG_TARGET,
"free: {available_space} , threshold: {threshold}.",
);
if available_space < threshold {
log::error!(target: LOG_TARGET, "Available space {available_space}MiB for path `{}` dropped below threshold: {threshold}MiB , terminating...", path.display());
Err(Error::StorageOutOfSpace(available_space, threshold))
} else {
Ok(())
}
},
Err(e) => {
log::error!(target: LOG_TARGET, "Could not read available space: {e:?}.");
Err(e)
},
}
}
}