Skip to main content

aster_forge_audit/
db_writer.rs

1//! Buffered audit-log persistence backed by [`aster_forge_db`].
2
3use std::sync::{Arc, OnceLock};
4use std::time::Duration;
5
6use aster_forge_db::{AuditLogCreate, create_audit_log_requests, create_audit_log_row};
7use sea_orm::DatabaseConnection;
8
9/// Default number of audit records retained in memory before direct-write fallback.
10pub const DEFAULT_AUDIT_LOG_QUEUE_CAPACITY: usize = 4096;
11/// Default number of audit records that triggers an immediate batch flush.
12pub const DEFAULT_AUDIT_LOG_BATCH_SIZE: usize = 100;
13/// Default delay before a partial audit batch is flushed.
14pub const DEFAULT_AUDIT_LOG_DELAYED_FLUSH_AFTER: Duration = Duration::from_secs(1);
15
16/// Buffering policy for [`AuditLogManager`].
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct AuditLogBufferConfig {
19    /// Maximum records retained in memory before overflow records are written directly.
20    pub queue_capacity: usize,
21    /// Number of queued records that triggers an immediate batch flush.
22    pub batch_size: usize,
23    /// Delay before a partial batch is flushed.
24    pub delayed_flush_after: Duration,
25}
26
27impl AuditLogBufferConfig {
28    /// Creates an audit-log buffering policy.
29    pub const fn new(
30        queue_capacity: usize,
31        batch_size: usize,
32        delayed_flush_after: Duration,
33    ) -> Self {
34        Self {
35            queue_capacity,
36            batch_size,
37            delayed_flush_after,
38        }
39    }
40}
41
42impl Default for AuditLogBufferConfig {
43    fn default() -> Self {
44        Self::new(
45            DEFAULT_AUDIT_LOG_QUEUE_CAPACITY,
46            DEFAULT_AUDIT_LOG_BATCH_SIZE,
47            DEFAULT_AUDIT_LOG_DELAYED_FLUSH_AFTER,
48        )
49    }
50}
51
52static GLOBAL_AUDIT_LOG_MANAGER: OnceLock<Arc<AuditLogManager>> = OnceLock::new();
53
54/// Best-effort buffered writer for Forge audit-log rows.
55pub struct AuditLogManager {
56    writer: Arc<aster_forge_runtime::BufferedBatchWriter<AuditLogCreate>>,
57}
58
59impl AuditLogManager {
60    /// Creates a manager with the default buffering policy.
61    pub fn new(db: DatabaseConnection) -> Self {
62        Self::with_config(db, AuditLogBufferConfig::default())
63    }
64
65    /// Creates a manager with a caller-provided buffering policy.
66    pub fn with_config(db: DatabaseConnection, config: AuditLogBufferConfig) -> Self {
67        let batch_size = config.batch_size.max(1);
68        let batch_db = db.clone();
69        let single_db = db;
70        let writer = aster_forge_runtime::BufferedBatchWriter::new(
71            aster_forge_runtime::BufferedBatchConfig::new(
72                config.queue_capacity.max(1),
73                batch_size,
74                config.delayed_flush_after,
75                "audit_log",
76            ),
77            move |batch| {
78                let db = batch_db.clone();
79                async move { write_audit_batch(&db, batch, batch_size).await }
80            },
81            move |request| {
82                let db = single_db.clone();
83                async move { write_audit_log_direct(&db, request).await }
84            },
85        );
86        Self {
87            writer: Arc::new(writer),
88        }
89    }
90
91    /// Queues an audit record, falling back to a direct write if the queue is full.
92    pub async fn record(&self, request: AuditLogCreate) {
93        self.writer.record(request).await;
94    }
95
96    /// Flushes all currently buffered audit records.
97    pub async fn flush(&self) {
98        self.writer.flush().await;
99    }
100
101    /// Cancels delayed flush tasks. Call [`AuditLogManager::flush`] afterwards on shutdown.
102    pub fn cancel(&self) {
103        self.writer.cancel();
104    }
105}
106
107/// Initializes the process-global audit manager.
108///
109/// Returns `true` when this call installed the manager and `false` when a manager was already
110/// initialized.
111pub fn init_global_audit_log_manager(db: DatabaseConnection) -> bool {
112    let installed = GLOBAL_AUDIT_LOG_MANAGER
113        .set(Arc::new(AuditLogManager::new(db)))
114        .is_ok();
115    if !installed {
116        tracing::warn!("global audit log manager is already initialized; ignoring");
117    }
118    installed
119}
120
121/// Returns the process-global audit manager when initialized.
122pub fn global_audit_log_manager() -> Option<&'static Arc<AuditLogManager>> {
123    GLOBAL_AUDIT_LOG_MANAGER.get()
124}
125
126/// Records through the global manager, or writes directly through `fallback_db` before startup.
127pub async fn record_audit_log(fallback_db: &DatabaseConnection, request: AuditLogCreate) {
128    if let Some(manager) = global_audit_log_manager() {
129        manager.record(request).await;
130    } else {
131        write_audit_log_direct(fallback_db, request).await;
132    }
133}
134
135/// Flushes the process-global audit manager when initialized.
136pub async fn flush_global_audit_log_manager() {
137    if let Some(manager) = global_audit_log_manager() {
138        manager.flush().await;
139    }
140}
141
142/// Cancels delayed writes and flushes the process-global audit manager when initialized.
143pub async fn shutdown_global_audit_log_manager() {
144    if let Some(manager) = global_audit_log_manager() {
145        manager.cancel();
146        manager.flush().await;
147    }
148}
149
150/// Writes one audit record directly, logging persistence failures as best-effort warnings.
151pub async fn write_audit_log_direct(db: &DatabaseConnection, request: AuditLogCreate) {
152    if let Err(error) = create_audit_log_row(db, request).await {
153        tracing::warn!(%error, "failed to write audit log");
154    }
155}
156
157async fn write_audit_batch(db: &DatabaseConnection, batch: Vec<AuditLogCreate>, batch_size: usize) {
158    let total = batch.len();
159    let mut requests = batch.into_iter();
160    loop {
161        let chunk = requests.by_ref().take(batch_size).collect::<Vec<_>>();
162        if chunk.is_empty() {
163            break;
164        }
165
166        let count = chunk.len();
167        if let Err(error) = create_audit_log_requests(db, chunk).await {
168            tracing::warn!(count, total, %error, "failed to write audit log batch");
169        }
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use std::sync::Arc;
176    use std::time::{Duration, Instant};
177
178    use aster_forge_db::{AuditLogCreate, audit_log};
179    use chrono::Utc;
180    use sea_orm::{ConnectionTrait, DatabaseConnection, EntityTrait, PaginatorTrait};
181
182    use super::{
183        AuditLogBufferConfig, AuditLogManager, DEFAULT_AUDIT_LOG_BATCH_SIZE,
184        DEFAULT_AUDIT_LOG_QUEUE_CAPACITY,
185    };
186
187    async fn test_db() -> DatabaseConnection {
188        let db = sea_orm::Database::connect("sqlite::memory:")
189            .await
190            .expect("audit writer test database should connect");
191        let backend = db.get_database_backend();
192        let create_table = aster_forge_db::create_audit_logs_table(backend);
193        db.execute(&create_table)
194            .await
195            .expect("audit writer test table should be created");
196        db
197    }
198
199    fn audit_request(index: i64) -> AuditLogCreate {
200        AuditLogCreate {
201            user_id: 42,
202            action: "file_upload".to_string(),
203            entity_type: "file".to_string(),
204            entity_id: Some(index),
205            entity_name: Some(format!("file-{index}.txt")),
206            details: None,
207            ip_address: Some("127.0.0.1".to_string()),
208            user_agent: Some("audit-writer-test".to_string()),
209            created_at: Utc::now(),
210        }
211    }
212
213    async fn audit_log_count(db: &DatabaseConnection) -> u64 {
214        audit_log::Entity::find()
215            .count(db)
216            .await
217            .expect("audit writer count query should succeed")
218    }
219
220    async fn wait_for_audit_log_count(db: &DatabaseConnection, expected: u64) {
221        let deadline = Instant::now() + Duration::from_secs(2);
222        loop {
223            let current = audit_log_count(db).await;
224            if current == expected {
225                return;
226            }
227            assert!(
228                current < expected,
229                "audit count exceeded {expected}: {current}"
230            );
231            assert!(
232                Instant::now() < deadline,
233                "timed out waiting for audit count {expected}; last count was {current}"
234            );
235            tokio::time::sleep(Duration::from_millis(10)).await;
236        }
237    }
238
239    fn test_manager(db: DatabaseConnection, delayed_flush_after: Duration) -> AuditLogManager {
240        AuditLogManager::with_config(
241            db,
242            AuditLogBufferConfig::new(
243                DEFAULT_AUDIT_LOG_QUEUE_CAPACITY,
244                DEFAULT_AUDIT_LOG_BATCH_SIZE,
245                delayed_flush_after,
246            ),
247        )
248    }
249
250    #[tokio::test]
251    async fn flushes_threshold_batch() {
252        let db = test_db().await;
253        let manager = test_manager(db.clone(), Duration::from_secs(5));
254
255        for index in 0..DEFAULT_AUDIT_LOG_BATCH_SIZE {
256            manager.record(audit_request(index as i64)).await;
257        }
258
259        wait_for_audit_log_count(&db, DEFAULT_AUDIT_LOG_BATCH_SIZE as u64).await;
260        manager.cancel();
261    }
262
263    #[tokio::test]
264    async fn flushes_partial_batch_after_delay() {
265        let db = test_db().await;
266        let manager = test_manager(db.clone(), Duration::from_millis(20));
267
268        for index in 0..3 {
269            manager.record(audit_request(index)).await;
270        }
271
272        wait_for_audit_log_count(&db, 3).await;
273        manager.cancel();
274    }
275
276    #[tokio::test]
277    async fn partial_batch_waits_for_configured_delay() {
278        let db = test_db().await;
279        let manager = test_manager(db.clone(), Duration::from_millis(120));
280
281        manager.record(audit_request(1)).await;
282        tokio::time::sleep(Duration::from_millis(30)).await;
283        assert_eq!(audit_log_count(&db).await, 0);
284
285        wait_for_audit_log_count(&db, 1).await;
286        manager.cancel();
287    }
288
289    #[tokio::test]
290    async fn cancelled_shutdown_flushes_buffer() {
291        let db = test_db().await;
292        let manager = test_manager(db.clone(), Duration::from_secs(5));
293
294        manager.record(audit_request(1)).await;
295        manager.cancel();
296        manager.flush().await;
297
298        assert_eq!(audit_log_count(&db).await, 1);
299    }
300
301    #[tokio::test]
302    async fn manual_flush_allows_later_delayed_flush() {
303        let db = test_db().await;
304        let manager = test_manager(db.clone(), Duration::from_millis(20));
305
306        manager.record(audit_request(1)).await;
307        manager.flush().await;
308        assert_eq!(audit_log_count(&db).await, 1);
309
310        manager.record(audit_request(2)).await;
311        wait_for_audit_log_count(&db, 2).await;
312        manager.cancel();
313    }
314
315    #[tokio::test]
316    async fn cancel_stops_delayed_flush_until_explicit_flush() {
317        let db = test_db().await;
318        let manager = test_manager(db.clone(), Duration::from_millis(20));
319
320        manager.record(audit_request(1)).await;
321        manager.cancel();
322        tokio::time::sleep(Duration::from_millis(60)).await;
323        assert_eq!(audit_log_count(&db).await, 0);
324
325        manager.flush().await;
326        assert_eq!(audit_log_count(&db).await, 1);
327    }
328
329    #[tokio::test]
330    async fn overflow_writes_extra_record_directly_then_flushes_buffer() {
331        let db = test_db().await;
332        let manager = Arc::new(test_manager(db.clone(), Duration::from_secs(5)));
333        let flush_guard = manager.writer.lock_flush_for_test().await;
334
335        for index in 0..DEFAULT_AUDIT_LOG_QUEUE_CAPACITY {
336            manager.record(audit_request(index as i64)).await;
337        }
338        manager.record(audit_request(10_000)).await;
339
340        assert_eq!(audit_log_count(&db).await, 1);
341        drop(flush_guard);
342
343        wait_for_audit_log_count(&db, (DEFAULT_AUDIT_LOG_QUEUE_CAPACITY + 1) as u64).await;
344        manager.cancel();
345    }
346
347    #[tokio::test]
348    async fn delayed_batch_follows_a_pending_immediate_flush() {
349        let db = test_db().await;
350        let manager = Arc::new(test_manager(db.clone(), Duration::from_millis(20)));
351        let flush_guard = manager.writer.lock_flush_for_test().await;
352
353        for index in 0..DEFAULT_AUDIT_LOG_BATCH_SIZE {
354            manager.record(audit_request(index as i64)).await;
355        }
356        manager
357            .record(audit_request(DEFAULT_AUDIT_LOG_BATCH_SIZE as i64))
358            .await;
359        assert_eq!(audit_log_count(&db).await, 0);
360
361        drop(flush_guard);
362        wait_for_audit_log_count(&db, (DEFAULT_AUDIT_LOG_BATCH_SIZE + 1) as u64).await;
363        manager.cancel();
364    }
365}