#![allow(dead_code)]
use crate::provider::ProviderRegistry;
use anyhow::Result;
use std::path::PathBuf;
use std::sync::{
atomic::{AtomicU64, Ordering},
Arc, RwLock,
};
use std::time::{Duration, SystemTime};
pub struct HotReloadRegistry {
inner: RwLock<Arc<ProviderRegistry>>,
rules_dir: PathBuf,
poll_interval: Duration,
last_mtime_nanos: AtomicU64,
reload_count: AtomicU64,
}
impl HotReloadRegistry {
pub fn start(rules_dir: impl Into<PathBuf>, poll_interval: Duration) -> Result<Arc<Self>> {
let rules_dir = rules_dir.into();
let registry = build_registry(&rules_dir)?;
let initial_mtime = scan_max_mtime_nanos(&rules_dir);
let this = Arc::new(Self {
inner: RwLock::new(Arc::new(registry)),
rules_dir: rules_dir.clone(),
poll_interval,
last_mtime_nanos: AtomicU64::new(initial_mtime),
reload_count: AtomicU64::new(0),
});
let weak = Arc::downgrade(&this);
tokio::spawn(async move {
loop {
tokio::time::sleep(poll_interval).await;
let Some(strong) = weak.upgrade() else {
return; };
strong.poll_once();
}
});
Ok(this)
}
pub fn current(&self) -> Arc<ProviderRegistry> {
self.inner
.read()
.expect("rule_watcher inner lock poisoned")
.clone()
}
pub fn reload_count(&self) -> u64 {
self.reload_count.load(Ordering::Relaxed)
}
pub fn poll_once(&self) -> bool {
let current_mtime = scan_max_mtime_nanos(&self.rules_dir);
let last = self.last_mtime_nanos.load(Ordering::Relaxed);
if current_mtime <= last {
return false;
}
match build_registry(&self.rules_dir) {
Ok(new_reg) => {
let mut guard = self
.inner
.write()
.expect("rule_watcher inner lock poisoned");
*guard = Arc::new(new_reg);
self.last_mtime_nanos
.store(current_mtime, Ordering::Relaxed);
self.reload_count.fetch_add(1, Ordering::Relaxed);
tracing::info!(
rules_dir = %self.rules_dir.display(),
reload_count = self.reload_count.load(Ordering::Relaxed),
"captchaforge rule watcher: registry hot-reloaded"
);
true
}
Err(e) => {
tracing::warn!(
rules_dir = %self.rules_dir.display(),
error = %e,
"captchaforge rule watcher: rebuild FAILED — keeping previous registry active"
);
self.last_mtime_nanos
.store(current_mtime, Ordering::Relaxed);
false
}
}
}
}
fn build_registry(rules_dir: &std::path::Path) -> Result<ProviderRegistry> {
use crate::detect::rules::RuleSet;
let mut registry = ProviderRegistry::with_built_in_rules()?;
if !rules_dir.exists() {
return Ok(registry);
}
let mut entries: Vec<PathBuf> = Vec::new();
for entry in std::fs::read_dir(rules_dir)? {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|s| s.to_str()) == Some("toml") {
entries.push(path);
}
}
entries.sort();
for path in entries {
let extra = RuleSet::load_from_path(&path)?;
for rule in extra.providers {
let detector = crate::detect::rules::RuleDetector::from(rule);
registry.add_provider(Box::new(detector));
}
}
registry.sort_by_priority();
Ok(registry)
}
fn scan_max_mtime_nanos(rules_dir: &std::path::Path) -> u64 {
let Ok(read) = std::fs::read_dir(rules_dir) else {
return 0;
};
let mut max_nanos = 0u64;
for entry in read.flatten() {
let path = entry.path();
if path.extension().and_then(|s| s.to_str()) != Some("toml") {
continue;
}
if let Ok(meta) = entry.metadata() {
if let Ok(mtime) = meta.modified() {
if let Ok(d) = mtime.duration_since(SystemTime::UNIX_EPOCH) {
let nanos = d.as_nanos() as u64;
if nanos > max_nanos {
max_nanos = nanos;
}
}
}
}
}
max_nanos
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
fn write_rule(dir: &std::path::Path, filename: &str, body: &str) {
std::fs::write(dir.join(filename), body).unwrap();
}
fn dummy_rule_toml(name: &str, priority: i32) -> String {
format!(
r#"
[[provider]]
name = "{name}"
priority = {priority}
[provider.triggers]
selectors = [".dummy"]
"#,
)
}
#[tokio::test]
async fn start_loads_initial_rules_from_dir() {
let tmp = tempdir().unwrap();
write_rule(
tmp.path(),
"vendor_a.toml",
&dummy_rule_toml("vendor_a", 100),
);
let h = HotReloadRegistry::start(tmp.path(), Duration::from_secs(60)).unwrap();
let reg = h.current();
let names: Vec<&'static str> = reg.providers().iter().map(|p| p.name()).collect();
assert!(
names.contains(&"vendor_a"),
"initial load missed vendor_a; got {names:?}"
);
}
#[tokio::test]
async fn poll_once_reloads_after_new_file_lands() {
let tmp = tempdir().unwrap();
write_rule(tmp.path(), "v1.toml", &dummy_rule_toml("v1", 100));
let h = HotReloadRegistry::start(tmp.path(), Duration::from_secs(60)).unwrap();
assert_eq!(h.reload_count(), 0);
tokio::time::sleep(Duration::from_millis(1100)).await;
write_rule(tmp.path(), "v2.toml", &dummy_rule_toml("v2", 101));
let reloaded = h.poll_once();
assert!(reloaded, "poll_once should detect new file");
assert_eq!(h.reload_count(), 1);
let names: Vec<&'static str> = h.current().providers().iter().map(|p| p.name()).collect();
assert!(names.contains(&"v2"), "reload missed v2; got {names:?}");
}
#[tokio::test]
async fn poll_once_is_no_op_when_dir_unchanged() {
let tmp = tempdir().unwrap();
write_rule(tmp.path(), "v.toml", &dummy_rule_toml("v", 100));
let h = HotReloadRegistry::start(tmp.path(), Duration::from_secs(60)).unwrap();
let count_before = h.reload_count();
assert!(!h.poll_once());
assert_eq!(h.reload_count(), count_before);
}
#[tokio::test]
async fn poll_once_keeps_previous_registry_on_bad_toml() {
let tmp = tempdir().unwrap();
write_rule(tmp.path(), "good.toml", &dummy_rule_toml("good", 100));
let h = HotReloadRegistry::start(tmp.path(), Duration::from_secs(60)).unwrap();
let names_before: Vec<&'static str> =
h.current().providers().iter().map(|p| p.name()).collect();
tokio::time::sleep(Duration::from_millis(1100)).await;
write_rule(tmp.path(), "bad.toml", "this is not valid TOML }{");
let reloaded = h.poll_once();
assert!(
!reloaded,
"broken TOML should NOT count as a successful reload"
);
let names_after: Vec<&'static str> =
h.current().providers().iter().map(|p| p.name()).collect();
assert_eq!(names_before, names_after);
}
#[tokio::test]
async fn missing_rules_dir_is_treated_as_no_extra_rules() {
let tmp = tempdir().unwrap();
let nonexistent = tmp.path().join("rules-not-yet-created");
let h = HotReloadRegistry::start(&nonexistent, Duration::from_secs(60)).unwrap();
assert!(!h.current().providers().is_empty());
}
#[tokio::test]
async fn watcher_terminates_when_outer_arc_dropped() {
let tmp = tempdir().unwrap();
write_rule(tmp.path(), "v.toml", &dummy_rule_toml("v", 100));
let h = HotReloadRegistry::start(tmp.path(), Duration::from_millis(50)).unwrap();
drop(h);
tokio::time::sleep(Duration::from_millis(100)).await;
}
}