use std::time::Duration;
use ursula_raft::LeadershipShedReason;
use ursula_raft::RaftGroupHandleRegistry;
use ursula_raft::RaftGroupMetricsSnapshot;
use ursula_runtime::ShardRuntime;
use ursula_runtime::SharedSnapshotStore;
use ursula_runtime::default_snapshot_store;
use ursula_shard::RaftGroupId;
pub(crate) fn resolve_snapshot_drive_interval_ms(
configured: Option<usize>,
snapshot_store_configured: bool,
) -> usize {
configured.unwrap_or(if snapshot_store_configured { 5_000 } else { 0 })
}
pub(crate) fn next_snapshot_to_drive(
snapshots: &[RaftGroupMetricsSnapshot],
next_pos: usize,
logs_since_last: u64,
) -> Option<(usize, &RaftGroupMetricsSnapshot)> {
if snapshots.is_empty() {
return None;
}
let start = next_pos % snapshots.len();
snapshots
.iter()
.enumerate()
.cycle()
.skip(start)
.take(snapshots.len())
.find(|(_, snapshot)| should_drive_snapshot_for_group(snapshot, logs_since_last))
}
pub(crate) fn should_drive_snapshot_for_group(
snapshot: &RaftGroupMetricsSnapshot,
logs_since_last: u64,
) -> bool {
let Some(last_applied) = snapshot.last_applied else {
return false;
};
let Some(current) = snapshot.snapshot else {
return true;
};
last_applied.index.saturating_sub(current.index) >= logs_since_last.max(1)
}
pub(crate) fn unpurged_log_entries(snapshot: &RaftGroupMetricsSnapshot) -> u64 {
let Some(last_log_index) = snapshot.last_log_index else {
return 0;
};
snapshot.purged.map_or_else(
|| last_log_index.saturating_add(1),
|purged| last_log_index.saturating_sub(purged.index),
)
}
fn snapshot_advance_entries(snapshot: &RaftGroupMetricsSnapshot) -> u64 {
let Some(last_applied) = snapshot.last_applied else {
return 0;
};
snapshot.snapshot.map_or_else(
|| last_applied.index.saturating_add(1),
|current| last_applied.index.saturating_sub(current.index),
)
}
pub(crate) fn pressure_snapshot_groups(
snapshots: &[RaftGroupMetricsSnapshot],
max_groups: usize,
) -> Vec<&RaftGroupMetricsSnapshot> {
let mut candidates = snapshots
.iter()
.filter_map(|snapshot| {
let advance = snapshot_advance_entries(snapshot);
(advance > 0).then_some((advance, snapshot))
})
.collect::<Vec<_>>();
candidates.sort_unstable_by(|(left_advance, left), (right_advance, right)| {
right_advance
.cmp(left_advance)
.then_with(|| left.raft_group_id.cmp(&right.raft_group_id))
});
candidates
.into_iter()
.take(max_groups)
.map(|(_, snapshot)| snapshot)
.collect()
}
pub fn spawn_snapshot_driver(
runtime: &ShardRuntime,
registry: &RaftGroupHandleRegistry,
snapshot_store: Option<SharedSnapshotStore>,
s3_cfg: Option<&ursula_config::S3Config>,
interval_ms: usize,
logs_since_last: u64,
pressure_unpurged_logs: u64,
pressure_max_groups_per_tick: usize,
) {
if interval_ms == 0 {
return;
}
let snapshot_store = snapshot_store.unwrap_or_else(default_snapshot_store);
let probe_timeout = Duration::from_millis(
s3_cfg
.map(|c| c.probe_timeout.as_duration().as_millis() as u64)
.unwrap_or(2_000),
);
let unhealthy_ticks = s3_cfg.map(|c| c.unhealthy_ticks).unwrap_or(1).max(1);
let heal_ticks = s3_cfg.map(|c| c.heal_ticks).unwrap_or(2).max(1);
let runtime = runtime.clone();
let registry = registry.clone();
tokio::spawn(async move {
let interval = Duration::from_millis(u64::try_from(interval_ms).unwrap_or(u64::MAX));
let mut consecutive_bad = 0usize;
let mut consecutive_good = 0usize;
let mut yielded = false;
let mut last_flush_errors = runtime.metrics().snapshot().cold_flush_write_errors;
let mut next_snapshot_drive_pos = 0usize;
loop {
let snaps = registry.metrics_snapshot();
let probe_healthy = matches!(
tokio::time::timeout(probe_timeout, snapshot_store.health_check()).await,
Ok(Ok(()))
);
let flush_errors_now = runtime.metrics().snapshot().cold_flush_write_errors;
let flush_failing = flush_errors_now > last_flush_errors;
last_flush_errors = flush_errors_now;
let bad_tick = !probe_healthy || flush_failing;
if bad_tick {
consecutive_bad += 1;
consecutive_good = 0;
} else {
consecutive_bad = 0;
consecutive_good += 1;
}
if !yielded && consecutive_bad >= unhealthy_ticks {
yielded = true;
registry.mark_leadership_shed(LeadershipShedReason::SnapshotDriverS3);
for snapshot in &snaps {
let Some(raft) = registry.get(RaftGroupId(snapshot.raft_group_id)) else {
continue;
};
if snapshot.current_leader == Some(snapshot.node_id)
&& let Some(target) = snapshot
.voter_ids
.iter()
.copied()
.find(|voter| *voter != snapshot.node_id)
{
match raft.trigger().transfer_leader(target).await {
Ok(()) => tracing::warn!(
"s3-unhealthy: node {} yielded leadership of group {} to node {}",
snapshot.node_id,
snapshot.raft_group_id,
target,
),
Err(err) => tracing::error!(
"s3-unhealthy: transfer_leader group {} -> {} failed: {err}",
snapshot.raft_group_id,
target,
),
}
}
}
} else if yielded && consecutive_good >= heal_ticks {
yielded = false;
registry.clear_leadership_shed(LeadershipShedReason::SnapshotDriverS3);
}
if !bad_tick {
let unpurged_logs = snaps
.iter()
.map(unpurged_log_entries)
.fold(0u64, u64::saturating_add);
if unpurged_logs >= pressure_unpurged_logs.max(1) {
let candidates =
pressure_snapshot_groups(&snaps, pressure_max_groups_per_tick.max(1));
let mut triggered = 0u64;
for snapshot in candidates {
let gid = snapshot.raft_group_id;
let Some(raft) = registry.get(RaftGroupId(gid)) else {
continue;
};
match raft.trigger().snapshot().await {
Ok(()) => triggered = triggered.saturating_add(1),
Err(err) => tracing::error!(
"snapshot pressure driver trigger group {gid} error: {err}"
),
}
}
runtime.metrics().record_raft_snapshot_pressure(triggered);
tracing::info!(
unpurged_logs,
pressure_unpurged_logs,
triggered,
"raft snapshot pressure pass completed"
);
} else if let Some((pos, snapshot)) =
next_snapshot_to_drive(&snaps, next_snapshot_drive_pos, logs_since_last)
{
next_snapshot_drive_pos = pos.wrapping_add(1);
if let Some(raft) = registry.get(RaftGroupId(snapshot.raft_group_id)) {
let gid = snapshot.raft_group_id;
if let Err(err) = raft.trigger().snapshot().await {
tracing::error!("snapshot driver trigger group {gid} error: {err}");
}
}
}
}
tokio::time::sleep(interval).await;
}
});
}