Skip to content

Commit 15f1fbe

Browse files
committed
feat: add --datadir.min-free-disk flag like in geth
1 parent 56e60a3 commit 15f1fbe

File tree

5 files changed

+279
-2
lines changed

5 files changed

+279
-2
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

crates/cli/commands/src/node.rs

Lines changed: 78 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,14 @@ use reth_node_core::{
1212
DatabaseArgs, DatadirArgs, DebugArgs, DevArgs, EngineArgs, EraArgs, MetricArgs,
1313
NetworkArgs, PayloadBuilderArgs, PruningArgs, RpcServerArgs, StaticFilesArgs, TxPoolArgs,
1414
},
15+
dirs::DataDirPath,
1516
node_config::NodeConfig,
17+
utils::is_disk_space_low,
1618
version,
1719
};
20+
use reth_tasks::{shutdown::Shutdown, TaskExecutor};
1821
use std::{ffi::OsString, fmt, path::PathBuf, sync::Arc};
22+
use tracing::error;
1923

2024
/// Start the node
2125
#[derive(Debug, Parser)]
@@ -172,6 +176,9 @@ where
172176
ext,
173177
} = self;
174178

179+
// Save min_free_disk before moving datadir
180+
let min_free_disk_mb = datadir.min_free_disk;
181+
175182
// set up node config
176183
let mut node_config = NodeConfig {
177184
datadir,
@@ -195,16 +202,45 @@ where
195202
let data_dir = node_config.datadir();
196203
let db_path = data_dir.db();
197204

205+
// Check disk space at startup if min_free_disk is configured
206+
if min_free_disk_mb > 0 {
207+
use reth_node_core::utils::is_disk_space_low;
208+
if is_disk_space_low(data_dir.data_dir(), min_free_disk_mb) {
209+
eyre::bail!(
210+
"Insufficient disk space: available space is below minimum threshold of {} MB",
211+
min_free_disk_mb
212+
);
213+
}
214+
}
215+
198216
tracing::info!(target: "reth::cli", path = ?db_path, "Opening database");
199217
let database = Arc::new(init_db(db_path.clone(), self.db.database_args())?.with_metrics());
200218

201219
if with_unused_ports {
202220
node_config = node_config.with_unused_ports();
203221
}
204222

223+
// Spawn disk space monitoring task if min_free_disk is configured
224+
// Start monitoring immediately after database is opened
225+
if min_free_disk_mb > 0 {
226+
let data_dir_for_monitor = data_dir.clone();
227+
let shutdown = ctx.task_executor.on_shutdown_signal().clone();
228+
let task_executor = ctx.task_executor.clone();
229+
230+
ctx.task_executor.spawn_critical(
231+
"disk space monitor",
232+
Box::pin(disk_space_monitor_task(
233+
data_dir_for_monitor,
234+
min_free_disk_mb,
235+
shutdown,
236+
task_executor,
237+
)),
238+
);
239+
}
240+
205241
let builder = NodeBuilder::new(node_config)
206242
.with_database(database)
207-
.with_launch_context(ctx.task_executor);
243+
.with_launch_context(ctx.task_executor.clone());
208244

209245
launcher.entrypoint(builder, ext).await
210246
}
@@ -217,6 +253,47 @@ impl<C: ChainSpecParser, Ext: clap::Args + fmt::Debug> NodeCommand<C, Ext> {
217253
}
218254
}
219255

256+
/// Disk space monitoring task that periodically checks disk space and triggers shutdown if below threshold.
257+
async fn disk_space_monitor_task(
258+
data_dir: reth_node_core::dirs::ChainPath<DataDirPath>,
259+
min_free_disk_mb: u64,
260+
mut shutdown: Shutdown,
261+
task_executor: TaskExecutor,
262+
) {
263+
use tokio::time::{interval, Duration as TokioDuration};
264+
265+
let mut interval = interval(TokioDuration::from_secs(60)); // Check every minute
266+
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
267+
268+
loop {
269+
tokio::select! {
270+
_ = interval.tick() => {
271+
if is_disk_space_low(data_dir.data_dir(), min_free_disk_mb) {
272+
error!(
273+
target: "reth::cli",
274+
?data_dir,
275+
available_threshold = min_free_disk_mb,
276+
"Disk space below minimum threshold, initiating shutdown"
277+
);
278+
// Trigger graceful shutdown
279+
if let Err(e) = task_executor.initiate_graceful_shutdown() {
280+
error!(
281+
target: "reth::cli",
282+
%e,
283+
"Failed to initiate graceful shutdown"
284+
);
285+
}
286+
return;
287+
}
288+
}
289+
_ = &mut shutdown => {
290+
// Normal shutdown signal received
291+
return;
292+
}
293+
}
294+
}
295+
}
296+
220297
/// No Additional arguments
221298
#[derive(Debug, Clone, Copy, Default, Args)]
222299
#[non_exhaustive]

crates/node/core/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ ipnet.workspace = true
6060
# io
6161
dirs-next.workspace = true
6262
shellexpand.workspace = true
63+
sysinfo = { workspace = true, features = ["disk"] }
6364

6465
# obs
6566
tracing.workspace = true

crates/node/core/src/args/datadir_args.rs

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,16 @@ pub struct DatadirArgs {
2727
verbatim_doc_comment
2828
)]
2929
pub static_files_path: Option<PathBuf>,
30+
31+
/// Minimum free disk space in MB, once reached triggers auto shut down (default = 0, disabled).
32+
#[arg(
33+
long = "datadir.min-free-disk",
34+
alias = "datadir.min_free_disk",
35+
value_name = "MB",
36+
default_value_t = 0,
37+
verbatim_doc_comment
38+
)]
39+
pub min_free_disk: u64,
3040
}
3141

3242
impl DatadirArgs {
@@ -55,4 +65,48 @@ mod tests {
5565
let args = CommandParser::<DatadirArgs>::parse_from(["reth"]).args;
5666
assert_eq!(args, default_args);
5767
}
68+
69+
#[test]
70+
fn test_parse_min_free_disk_flag() {
71+
// Test with hyphen format
72+
let args = CommandParser::<DatadirArgs>::parse_from([
73+
"reth",
74+
"--datadir.min-free-disk",
75+
"1000",
76+
])
77+
.args;
78+
assert_eq!(args.min_free_disk, 1000);
79+
80+
// Test with underscore format (alias)
81+
let args = CommandParser::<DatadirArgs>::parse_from([
82+
"reth",
83+
"--datadir.min_free_disk",
84+
"500",
85+
])
86+
.args;
87+
assert_eq!(args.min_free_disk, 500);
88+
89+
// Test default value (0 = disabled)
90+
let args = CommandParser::<DatadirArgs>::parse_from(["reth"]).args;
91+
assert_eq!(args.min_free_disk, 0);
92+
}
93+
94+
#[test]
95+
fn test_min_free_disk_default() {
96+
let args = DatadirArgs::default();
97+
assert_eq!(args.min_free_disk, 0, "Default should be 0 (disabled)");
98+
}
99+
100+
#[test]
101+
fn test_min_free_disk_with_datadir() {
102+
let args = CommandParser::<DatadirArgs>::parse_from([
103+
"reth",
104+
"--datadir",
105+
"/tmp/test-datadir",
106+
"--datadir.min-free-disk",
107+
"2000",
108+
])
109+
.args;
110+
assert_eq!(args.min_free_disk, 2000);
111+
}
58112
}

crates/node/core/src/utils.rs

Lines changed: 145 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use std::{
1414
env::VarError,
1515
path::{Path, PathBuf},
1616
};
17-
use tracing::{debug, info};
17+
use tracing::{debug, info, warn};
1818

1919
/// Parses a user-specified path with support for environment variables and common shorthands (e.g.
2020
/// ~ for the user's home directory).
@@ -89,3 +89,147 @@ where
8989

9090
Ok(block)
9191
}
92+
93+
/// Check available disk space for the given path using sysinfo.
94+
///
95+
/// Returns the available space in MB, or None if the check fails.
96+
pub fn get_available_disk_space_mb(path: &Path) -> Option<u64> {
97+
use sysinfo::Disks;
98+
99+
// Find the disk that contains the given path
100+
let path_canonical = match std::fs::canonicalize(path) {
101+
Ok(p) => p,
102+
Err(_) => return None,
103+
};
104+
105+
let disks = Disks::new_with_refreshed_list();
106+
107+
// Find the disk that contains the given path
108+
for disk in disks.iter() {
109+
let mount_point = disk.mount_point();
110+
if path_canonical.starts_with(mount_point) {
111+
// Get available space in bytes, convert to MB
112+
let available_bytes = disk.available_space();
113+
return Some(available_bytes / (1024 * 1024));
114+
}
115+
}
116+
117+
None
118+
}
119+
120+
/// Check if disk space is below the minimum threshold.
121+
///
122+
/// Returns true if the available disk space is below the minimum threshold (in MB).
123+
pub fn is_disk_space_low(path: &Path, min_free_disk_mb: u64) -> bool {
124+
if min_free_disk_mb == 0 {
125+
return false; // Feature disabled
126+
}
127+
128+
match get_available_disk_space_mb(path) {
129+
Some(available_mb) => {
130+
if available_mb <= min_free_disk_mb {
131+
warn!(
132+
target: "reth::cli",
133+
?path,
134+
available_mb,
135+
min_free_disk_mb,
136+
"Disk space below minimum threshold"
137+
);
138+
return true;
139+
}
140+
false
141+
}
142+
None => {
143+
warn!(
144+
target: "reth::cli",
145+
?path,
146+
"Failed to check disk space, continuing anyway"
147+
);
148+
false
149+
}
150+
}
151+
}
152+
153+
#[cfg(test)]
154+
mod tests {
155+
use super::*;
156+
use std::path::Path;
157+
158+
#[test]
159+
fn test_is_disk_space_low_disabled() {
160+
// When min_free_disk is 0, feature is disabled
161+
let temp_dir = std::env::temp_dir();
162+
assert!(!is_disk_space_low(&temp_dir, 0), "Should return false when disabled");
163+
}
164+
165+
#[test]
166+
fn test_is_disk_space_low_with_valid_path() {
167+
// Test with a valid path (should not panic)
168+
// We can't easily mock sysinfo, so we just test that it doesn't panic
169+
// and returns a boolean value
170+
let temp_dir = std::env::temp_dir();
171+
let result = is_disk_space_low(&temp_dir, 1_000_000_000); // Very large threshold
172+
// Should return either true or false, but not panic
173+
assert!(result == true || result == false);
174+
}
175+
176+
#[test]
177+
fn test_is_disk_space_low_with_invalid_path() {
178+
// Test with a non-existent path
179+
let invalid_path = Path::new("/nonexistent/path/that/does/not/exist");
180+
// Should not panic, but may return false (feature disabled) or handle gracefully
181+
let result = is_disk_space_low(invalid_path, 1000);
182+
// Should not panic, result depends on sysinfo behavior
183+
assert!(result == true || result == false);
184+
}
185+
186+
#[test]
187+
fn test_get_available_disk_space_mb_with_valid_path() {
188+
// Test with a valid path
189+
let temp_dir = std::env::temp_dir();
190+
let result = get_available_disk_space_mb(&temp_dir);
191+
// Should return Some(u64) if successful, or None if it fails
192+
match result {
193+
Some(mb) => {
194+
// If successful, should be a reasonable value (not 0 for temp dir)
195+
assert!(mb > 0 || mb == 0, "Available space should be a valid number");
196+
}
197+
None => {
198+
// It's okay if it fails, sysinfo might not work in all test environments
199+
}
200+
}
201+
}
202+
203+
#[test]
204+
fn test_get_available_disk_space_mb_with_invalid_path() {
205+
// Test with a non-existent path
206+
let invalid_path = Path::new("/nonexistent/path/that/does/not/exist");
207+
let result = get_available_disk_space_mb(invalid_path);
208+
// Should return None for invalid paths
209+
assert_eq!(result, None);
210+
}
211+
212+
#[test]
213+
fn test_is_disk_space_low_threshold_comparison() {
214+
// Test that the function correctly compares available space with threshold
215+
let temp_dir = std::env::temp_dir();
216+
217+
if let Some(available_mb) = get_available_disk_space_mb(&temp_dir) {
218+
// Test with a threshold smaller than available space (should pass - space is sufficient)
219+
if available_mb > 0 {
220+
let result_small = is_disk_space_low(&temp_dir, available_mb.saturating_sub(1));
221+
assert!(!result_small, "Should pass (return false) when threshold is less than available space");
222+
}
223+
224+
// Test with a threshold larger than available space (should fail - space is insufficient)
225+
let result_large = is_disk_space_low(&temp_dir, available_mb.saturating_add(1));
226+
assert!(result_large, "Should fail (return true) when threshold is greater than available space");
227+
228+
// Test with threshold equal to available space (edge case)
229+
// When available == threshold, should trigger shutdown (return true) because "once reached" includes equality
230+
let result_equal = is_disk_space_low(&temp_dir, available_mb);
231+
assert!(result_equal, "Should fail (return true) when threshold equals available space, as 'once reached' includes equality");
232+
}
233+
// If we can't get disk space, that's okay - test environment might not support it
234+
}
235+
}

0 commit comments

Comments
 (0)