miden_note_transport_node/database/
maintenance.rs1use std::sync::Arc;
2
3use tokio::time::{Duration, sleep};
4use tracing::{error, info};
5
6use super::{Database, DatabaseConfig};
7use crate::Result;
8use crate::metrics::MetricsDatabase;
9
10enum State {
11 Stopped,
12 Running,
13}
14
15pub struct DatabaseMaintenance {
17 database: Arc<Database>,
18 config: DatabaseConfig,
19 state: State,
20 metrics: MetricsDatabase,
21}
22
23impl DatabaseMaintenance {
24 pub fn new(database: Arc<Database>, config: DatabaseConfig, metrics: MetricsDatabase) -> Self {
26 Self {
27 database,
28 config,
29 state: State::Stopped,
30 metrics,
31 }
32 }
33
34 pub async fn entrypoint(mut self) {
36 self.state = State::Running;
37 while self.is_active() {
38 if let Err(e) = self.step().await {
39 error!("Database maintenance error: {e}");
40 }
41 }
42 }
43
44 async fn step(&mut self) -> Result<()> {
45 let timer = self.metrics.db_maintenance_cleanup_notes();
46
47 self.database.cleanup_old_notes(self.config.retention_days).await?;
48 info!("Cleaned up old notes");
49
50 timer.finish("ok");
51
52 sleep(Duration::from_secs(600)).await;
53
54 Ok(())
55 }
56
57 fn is_active(&self) -> bool {
58 matches!(self.state, State::Running)
59 }
60}
61
62#[cfg(test)]
63mod tests {
64 use chrono::Utc;
65 use serial_test::serial;
66
67 use super::*;
68 use crate::metrics::Metrics;
69 use crate::test_utils::test_note_header;
70 use crate::types::StoredNote;
71
72 fn note_at(age: Duration) -> StoredNote {
73 StoredNote {
74 header: test_note_header(),
75 details: vec![1, 2, 3, 4],
76 created_at: Utc::now() - age,
77 seq: 0, }
79 }
80
81 #[tokio::test]
82 #[serial]
83 async fn test_cleanup_old_notes_no_retention() {
84 let config = DatabaseConfig { retention_days: 0, ..Default::default() };
85
86 let db = Arc::new(Database::connect(config.clone(), Metrics::default().db).await.unwrap());
87 db.store_note(¬e_at(Duration::from_secs(30))).await.unwrap();
88
89 let maintenance = DatabaseMaintenance::new(db.clone(), config, Metrics::default().db);
90 tokio::spawn(maintenance.entrypoint());
91 sleep(Duration::from_secs(2)).await;
92
93 let (total_notes, _) = db.get_stats().await.unwrap();
94 assert_eq!(total_notes, 0);
95 }
96
97 #[tokio::test]
98 #[serial]
99 async fn test_cleanup_old_notes_retention() {
100 let config = DatabaseConfig { retention_days: 7, ..Default::default() };
101
102 let db = Arc::new(Database::connect(config.clone(), Metrics::default().db).await.unwrap());
103 db.store_note(¬e_at(Duration::from_secs(30))).await.unwrap();
104
105 let maintenance = DatabaseMaintenance::new(db.clone(), config, Metrics::default().db);
106 tokio::spawn(maintenance.entrypoint());
107 sleep(Duration::from_secs(2)).await;
108
109 let (total_notes, _) = db.get_stats().await.unwrap();
110 assert_eq!(total_notes, 1);
111 }
112
113 #[tokio::test]
114 #[serial]
115 async fn test_cleanup_old_notes_mixed_ages() {
116 let config = DatabaseConfig { retention_days: 1, ..Default::default() };
117
118 let db = Arc::new(Database::connect(config.clone(), Metrics::default().db).await.unwrap());
119 db.store_note(¬e_at(Duration::from_secs(30))).await.unwrap();
120 db.store_note(¬e_at(Duration::from_secs(3600 * 26))).await.unwrap();
121
122 let maintenance = DatabaseMaintenance::new(db.clone(), config, Metrics::default().db);
123 tokio::spawn(maintenance.entrypoint());
124 sleep(Duration::from_secs(2)).await;
125
126 let (total_notes, _) = db.get_stats().await.unwrap();
127 assert_eq!(total_notes, 1);
128 }
129}