use std::future::Future;
use std::pin::Pin;
use serde::Serialize;
use crate::AutumnResult;
use crate::state::AppState;
use crate::task::TaskInfo;
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct RetentionSweepReport {
pub model: String,
pub table: String,
pub rows_swept: u64,
pub duration_ms: u64,
pub dry_run: bool,
}
type DryRunFn =
fn(AppState) -> Pin<Box<dyn Future<Output = AutumnResult<RetentionSweepReport>> + Send>>;
#[doc(hidden)]
pub struct RetentionSweepDescriptor {
pub model_name: &'static str,
pub table_name: &'static str,
pub task_info: fn() -> TaskInfo,
pub dry_run: DryRunFn,
}
inventory::collect!(RetentionSweepDescriptor);
#[must_use]
pub fn collect_retention_tasks() -> Vec<TaskInfo> {
let tasks: Vec<TaskInfo> = inventory::iter::<RetentionSweepDescriptor>()
.map(|descriptor| (descriptor.task_info)())
.collect();
if let Some(message) = duplicate_retention_task_name(&tasks) {
panic!("{message}");
}
tasks
}
fn duplicate_retention_task_name(tasks: &[TaskInfo]) -> Option<String> {
let mut seen = std::collections::HashSet::with_capacity(tasks.len());
tasks.iter().find_map(|task| {
(!seen.insert(task.name.as_str())).then(|| {
format!(
"two #[repository(..., retention(...))] policies both produced the retention \
task name {:?} — most likely two different repositories declaring \
retention(...) on the same table. Only one repository per table may declare \
retention(...) for now.",
task.name
)
})
})
}
#[must_use]
pub fn has_retention_descriptors() -> bool {
inventory::iter::<RetentionSweepDescriptor>
.into_iter()
.next()
.is_some()
}
pub async fn run_retention_dry_run(
state: &AppState,
model_filter: Option<&str>,
) -> AutumnResult<Vec<RetentionSweepReport>> {
let descriptors = resolve_retention_descriptors(model_filter)?;
let mut reports = Vec::with_capacity(descriptors.len());
for descriptor in descriptors {
reports.push((descriptor.dry_run)(state.clone()).await?);
}
reports.sort_by(|a, b| a.model.cmp(&b.model).then_with(|| a.table.cmp(&b.table)));
Ok(reports)
}
#[doc(hidden)]
pub fn resolve_retention_descriptors(
model_filter: Option<&str>,
) -> AutumnResult<Vec<&'static RetentionSweepDescriptor>> {
let all: Vec<&RetentionSweepDescriptor> =
inventory::iter::<RetentionSweepDescriptor>().collect();
validate_resolved_descriptors(&all);
let matches: Vec<&RetentionSweepDescriptor> = match model_filter {
None => all,
Some(filter) => {
let table_matches: Vec<&RetentionSweepDescriptor> = all
.iter()
.copied()
.filter(|descriptor| descriptor.table_name == filter)
.collect();
let found: Vec<&RetentionSweepDescriptor> = if table_matches.len() == 1 {
table_matches
} else {
all.into_iter()
.filter(|descriptor| {
descriptor.model_name == filter || descriptor.table_name == filter
})
.collect()
};
if found.is_empty() {
return Err(crate::AutumnError::not_found_msg(format!(
"no #[repository(..., retention(...))] policy is registered for model \
{filter:?}"
)));
}
if found.len() > 1 {
let table_names: Vec<&str> = found.iter().map(|d| d.table_name).collect();
return Err(crate::AutumnError::bad_request_msg(format!(
"{filter:?} matches more than one retention policy ({table_names:?}); pass \
the table name instead of the model name to disambiguate, e.g. --model {}",
table_names[0]
)));
}
found
}
};
Ok(matches)
}
#[doc(hidden)]
#[must_use]
pub fn all_retention_descriptors() -> Vec<&'static RetentionSweepDescriptor> {
inventory::iter::<RetentionSweepDescriptor>().collect()
}
fn validate_resolved_descriptors(matches: &[&RetentionSweepDescriptor]) {
let task_infos: Vec<TaskInfo> = matches.iter().map(|d| (d.task_info)()).collect();
if let Some(message) = duplicate_retention_task_name(&task_infos) {
panic!("{message}");
}
}
pub fn log_retention_sweep(report: &RetentionSweepReport) {
tracing::info!(
model = %report.model,
table = %report.table,
rows_swept = report.rows_swept,
duration_ms = report.duration_ms,
dry_run = report.dry_run,
"retention: sweep complete"
);
if !report.dry_run {
crate::metrics::counter("retention_sweep_rows_total")
.with_label("model", report.model.clone())
.with_label("table", report.table.clone())
.increment(report.rows_swept);
crate::metrics::timer("retention_sweep_duration_seconds")
.with_label("model", report.model.clone())
.with_label("table", report.table.clone())
.record(std::time::Duration::from_millis(report.duration_ms));
}
}
#[cfg(test)]
mod tests {
use super::*;
fn sample_dry_run(
_state: AppState,
) -> Pin<Box<dyn Future<Output = AutumnResult<RetentionSweepReport>> + Send>> {
Box::pin(async {
Ok(RetentionSweepReport {
model: "Widget".to_string(),
table: "widgets".to_string(),
rows_swept: 3,
duration_ms: 5,
dry_run: true,
})
})
}
fn sample_task_info() -> TaskInfo {
task_info_named("retention-sweep-widget")
}
fn task_info_named(name: &str) -> TaskInfo {
TaskInfo {
name: name.to_string(),
schedule: crate::task::Schedule::FixedDelay(std::time::Duration::from_secs(3600)),
coordination: crate::task::TaskCoordination::Fleet,
handler: |_state| Box::pin(async { Ok(()) }),
}
}
fn by_model_name_task_info() -> TaskInfo {
task_info_named("retention-sweep-__retention_runtime_test_widgets")
}
fn by_table_name_task_info() -> TaskInfo {
task_info_named("retention-sweep-__retention_runtime_test_widgets_by_table")
}
fn ambiguous_a_task_info() -> TaskInfo {
task_info_named("retention-sweep-__retention_runtime_ambiguous_widgets_a")
}
fn ambiguous_b_task_info() -> TaskInfo {
task_info_named("retention-sweep-__retention_runtime_ambiguous_widgets_b")
}
fn table_shadows_model_task_info_a() -> TaskInfo {
task_info_named("retention-sweep-__retention_runtime_table_shadows_model_a")
}
fn table_shadows_model_task_info_b() -> TaskInfo {
task_info_named("retention-sweep-__retention_runtime_table_shadows_model_b")
}
#[test]
fn report_serializes_with_expected_fields() {
let report = RetentionSweepReport {
model: "Widget".to_string(),
table: "widgets".to_string(),
rows_swept: 42,
duration_ms: 7,
dry_run: false,
};
let json = serde_json::to_value(&report).expect("report should serialize");
assert_eq!(json["model"], "Widget");
assert_eq!(json["table"], "widgets");
assert_eq!(json["rows_swept"], 42);
assert_eq!(json["duration_ms"], 7);
assert_eq!(json["dry_run"], false);
}
#[test]
fn collect_retention_tasks_calls_every_registered_descriptor() {
let tasks = collect_retention_tasks();
assert!(
tasks.iter().all(|t| !t.name.is_empty()),
"every collected retention task must have a name"
);
}
#[test]
fn duplicate_retention_task_name_detects_a_collision() {
let tasks = vec![sample_task_info(), sample_task_info()];
let message = duplicate_retention_task_name(&tasks)
.expect("two tasks with the same name must be flagged as a collision");
assert!(message.contains(&sample_task_info().name), "{message}");
}
#[test]
fn duplicate_retention_task_name_accepts_unique_names() {
let mut b = sample_task_info();
b.name = "retention-sweep-other-table".to_string();
let tasks = vec![sample_task_info(), b];
assert!(
duplicate_retention_task_name(&tasks).is_none(),
"two tasks with different names must not be flagged as a collision"
);
}
#[tokio::test]
async fn run_retention_dry_run_filters_by_model_name() {
struct Fixture;
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeTestWidget",
table_name: "__retention_runtime_test_widgets",
task_info: by_model_name_task_info,
dry_run: sample_dry_run,
}
}
let _ = Fixture;
let state = AppState::for_test();
let reports = run_retention_dry_run(&state, Some("__RetentionRuntimeTestWidget"))
.await
.expect("dry run should succeed for a registered model");
assert_eq!(reports.len(), 1);
assert_eq!(reports[0].model, "Widget");
assert!(reports[0].dry_run);
}
#[tokio::test]
async fn run_retention_dry_run_filters_by_table_name() {
struct Fixture;
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeTestWidgetByTable",
table_name: "__retention_runtime_test_widgets_by_table",
task_info: by_table_name_task_info,
dry_run: sample_dry_run,
}
}
let _ = Fixture;
let state = AppState::for_test();
let reports =
run_retention_dry_run(&state, Some("__retention_runtime_test_widgets_by_table"))
.await
.expect("dry run should succeed when filtering by table name");
assert_eq!(reports.len(), 1);
assert_eq!(reports[0].model, "Widget");
}
#[tokio::test]
async fn run_retention_dry_run_rejects_ambiguous_model_filter() {
struct FixtureA;
struct FixtureB;
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeAmbiguousWidget",
table_name: "__retention_runtime_ambiguous_widgets_a",
task_info: ambiguous_a_task_info,
dry_run: sample_dry_run,
}
}
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeAmbiguousWidget",
table_name: "__retention_runtime_ambiguous_widgets_b",
task_info: ambiguous_b_task_info,
dry_run: sample_dry_run,
}
}
let _ = (FixtureA, FixtureB);
let state = AppState::for_test();
let error = run_retention_dry_run(&state, Some("__RetentionRuntimeAmbiguousWidget"))
.await
.expect_err("two policies sharing a model name must be rejected as ambiguous");
assert!(
error
.to_string()
.contains("__retention_runtime_ambiguous_widgets_a")
|| error
.to_string()
.contains("__retention_runtime_ambiguous_widgets_b"),
"error should name the disambiguating table names: {error}"
);
}
#[test]
fn resolve_retention_descriptors_prefers_exact_table_match_over_ambiguous_model_match() {
struct FixtureA;
struct FixtureB;
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeTableShadowsModel",
table_name: "__retention_runtime_table_shadows_model_a",
task_info: table_shadows_model_task_info_a,
dry_run: sample_dry_run,
}
}
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeTableShadowsModelOther",
table_name: "__RetentionRuntimeTableShadowsModel",
task_info: table_shadows_model_task_info_b,
dry_run: sample_dry_run,
}
}
let _ = (FixtureA, FixtureB);
let descriptors =
resolve_retention_descriptors(Some("__RetentionRuntimeTableShadowsModel")).expect(
"an exact table-name match must resolve unambiguously even though another \
policy's model_name equals the same string",
);
assert_eq!(descriptors.len(), 1);
assert_eq!(
descriptors[0].table_name,
"__RetentionRuntimeTableShadowsModel"
);
assert_eq!(
descriptors[0].model_name,
"__RetentionRuntimeTableShadowsModelOther"
);
}
#[tokio::test]
async fn run_retention_dry_run_rejects_unknown_model_filter() {
let state = AppState::for_test();
let error = run_retention_dry_run(&state, Some("__NoSuchRetentionModel"))
.await
.expect_err("an unregistered model filter must error");
assert!(error.to_string().contains("__NoSuchRetentionModel"));
}
static COUNTING_TASK_INFO_CALLS: std::sync::atomic::AtomicUsize =
std::sync::atomic::AtomicUsize::new(0);
fn counting_task_info() -> TaskInfo {
COUNTING_TASK_INFO_CALLS.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
task_info_named("retention-sweep-__retention_runtime_test_counting")
}
#[test]
fn resolve_retention_descriptors_invokes_task_info_validation() {
struct Fixture;
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeTestCounting",
table_name: "__retention_runtime_test_counting",
task_info: counting_task_info,
dry_run: sample_dry_run,
}
}
let _ = Fixture;
let before = COUNTING_TASK_INFO_CALLS.load(std::sync::atomic::Ordering::SeqCst);
let descriptors = resolve_retention_descriptors(Some("__RetentionRuntimeTestCounting"))
.expect("resolving a registered model must succeed");
let after = COUNTING_TASK_INFO_CALLS.load(std::sync::atomic::Ordering::SeqCst);
assert_eq!(descriptors.len(), 1);
assert!(
after > before,
"resolve_retention_descriptors must call task_info() on every matched descriptor \
to run the same validation collect_retention_tasks() runs at boot, so a dry run \
cannot report success for a policy real boot would panic on"
);
}
static UNSELECTED_TASK_INFO_CALLS: std::sync::atomic::AtomicUsize =
std::sync::atomic::AtomicUsize::new(0);
fn unselected_counting_task_info() -> TaskInfo {
UNSELECTED_TASK_INFO_CALLS.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
task_info_named("retention-sweep-__retention_runtime_test_unselected")
}
fn filter_selected_task_info() -> TaskInfo {
task_info_named("retention-sweep-__retention_runtime_test_filter_selected")
}
#[test]
fn resolve_retention_descriptors_validates_the_full_registry_even_when_filtered() {
struct SelectedFixture;
struct UnselectedFixture;
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeTestFilterSelected",
table_name: "__retention_runtime_test_filter_selected",
task_info: filter_selected_task_info,
dry_run: sample_dry_run,
}
}
inventory::submit! {
RetentionSweepDescriptor {
model_name: "__RetentionRuntimeTestFilterUnselected",
table_name: "__retention_runtime_test_unselected",
task_info: unselected_counting_task_info,
dry_run: sample_dry_run,
}
}
let _ = (SelectedFixture, UnselectedFixture);
let before = UNSELECTED_TASK_INFO_CALLS.load(std::sync::atomic::Ordering::SeqCst);
let descriptors =
resolve_retention_descriptors(Some("__RetentionRuntimeTestFilterSelected"))
.expect("resolving a registered model must succeed");
let after = UNSELECTED_TASK_INFO_CALLS.load(std::sync::atomic::Ordering::SeqCst);
assert_eq!(
descriptors.len(),
1,
"the returned/counted set must still be narrowed to the --model filter"
);
assert!(
after > before,
"resolve_retention_descriptors must validate every registered descriptor, \
including ones --model does not select, since real boot has no filter concept"
);
}
#[test]
fn validate_resolved_descriptors_panics_on_duplicate_task_names() {
let a = RetentionSweepDescriptor {
model_name: "__ValidateResolvedDescriptorsA",
table_name: "__validate_resolved_descriptors_dup",
task_info: || task_info_named("retention-sweep-__validate_resolved_descriptors_dup"),
dry_run: sample_dry_run,
};
let b = RetentionSweepDescriptor {
model_name: "__ValidateResolvedDescriptorsB",
table_name: "__validate_resolved_descriptors_dup",
task_info: || task_info_named("retention-sweep-__validate_resolved_descriptors_dup"),
dry_run: sample_dry_run,
};
let matches: Vec<&RetentionSweepDescriptor> = vec![&a, &b];
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
validate_resolved_descriptors(&matches);
}));
assert!(
result.is_err(),
"validate_resolved_descriptors must panic when two descriptors produce the same \
task name"
);
}
#[test]
fn validate_resolved_descriptors_accepts_unique_task_names() {
let a = RetentionSweepDescriptor {
model_name: "__ValidateResolvedDescriptorsUniqueA",
table_name: "__validate_resolved_descriptors_unique_a",
task_info: || {
task_info_named("retention-sweep-__validate_resolved_descriptors_unique_a")
},
dry_run: sample_dry_run,
};
let b = RetentionSweepDescriptor {
model_name: "__ValidateResolvedDescriptorsUniqueB",
table_name: "__validate_resolved_descriptors_unique_b",
task_info: || {
task_info_named("retention-sweep-__validate_resolved_descriptors_unique_b")
},
dry_run: sample_dry_run,
};
let matches: Vec<&RetentionSweepDescriptor> = vec![&a, &b];
validate_resolved_descriptors(&matches);
}
}