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