Skip to main content

miden_note_transport_node/database/
maintenance.rs

1use 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
15/// Perform periodic maintenance of the database
16pub struct DatabaseMaintenance {
17    database: Arc<Database>,
18    config: DatabaseConfig,
19    state: State,
20    metrics: MetricsDatabase,
21}
22
23impl DatabaseMaintenance {
24    /// Main constructor
25    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    /// Database maintenance running-task
35    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        let deleted = self.database.cleanup_old_notes(self.config.retention_days).await?;
48        info!(
49            notes_deleted = deleted,
50            retention_days = self.config.retention_days,
51            "Maintenance cleanup completed"
52        );
53
54        timer.finish("ok");
55
56        sleep(Duration::from_secs(600)).await;
57
58        Ok(())
59    }
60
61    fn is_active(&self) -> bool {
62        matches!(self.state, State::Running)
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use chrono::Utc;
69    use serial_test::serial;
70
71    use super::*;
72    use crate::metrics::Metrics;
73    use crate::test_utils::test_note_header;
74    use crate::types::StoredNote;
75
76    fn note_at(age: Duration) -> StoredNote {
77        StoredNote {
78            header: test_note_header(),
79            details: vec![1, 2, 3, 4],
80            created_at: Utc::now() - age,
81            seq: 0, // ignored on INSERT
82            after_block_num: None,
83        }
84    }
85
86    #[tokio::test]
87    #[serial]
88    async fn test_cleanup_old_notes_no_retention() {
89        let config = DatabaseConfig { retention_days: 0, ..Default::default() };
90
91        let db = Arc::new(Database::connect(config.clone(), Metrics::default().db).await.unwrap());
92        db.store_note(&note_at(Duration::from_secs(30))).await.unwrap();
93
94        let maintenance = DatabaseMaintenance::new(db.clone(), config, Metrics::default().db);
95        tokio::spawn(maintenance.entrypoint());
96        sleep(Duration::from_secs(2)).await;
97
98        let (total_notes, _) = db.get_stats().await.unwrap();
99        assert_eq!(total_notes, 0);
100    }
101
102    #[tokio::test]
103    #[serial]
104    async fn test_cleanup_old_notes_retention() {
105        let config = DatabaseConfig { retention_days: 7, ..Default::default() };
106
107        let db = Arc::new(Database::connect(config.clone(), Metrics::default().db).await.unwrap());
108        db.store_note(&note_at(Duration::from_secs(30))).await.unwrap();
109
110        let maintenance = DatabaseMaintenance::new(db.clone(), config, Metrics::default().db);
111        tokio::spawn(maintenance.entrypoint());
112        sleep(Duration::from_secs(2)).await;
113
114        let (total_notes, _) = db.get_stats().await.unwrap();
115        assert_eq!(total_notes, 1);
116    }
117
118    #[tokio::test]
119    #[serial]
120    async fn test_cleanup_old_notes_mixed_ages() {
121        let config = DatabaseConfig { retention_days: 1, ..Default::default() };
122
123        let db = Arc::new(Database::connect(config.clone(), Metrics::default().db).await.unwrap());
124        db.store_note(&note_at(Duration::from_secs(30))).await.unwrap();
125        db.store_note(&note_at(Duration::from_secs(3600 * 26))).await.unwrap();
126
127        let maintenance = DatabaseMaintenance::new(db.clone(), config, Metrics::default().db);
128        tokio::spawn(maintenance.entrypoint());
129        sleep(Duration::from_secs(2)).await;
130
131        let (total_notes, _) = db.get_stats().await.unwrap();
132        assert_eq!(total_notes, 1);
133    }
134}