1use 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
9pub const DEFAULT_AUDIT_LOG_QUEUE_CAPACITY: usize = 4096;
11pub const DEFAULT_AUDIT_LOG_BATCH_SIZE: usize = 100;
13pub const DEFAULT_AUDIT_LOG_DELAYED_FLUSH_AFTER: Duration = Duration::from_secs(1);
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct AuditLogBufferConfig {
19 pub queue_capacity: usize,
21 pub batch_size: usize,
23 pub delayed_flush_after: Duration,
25}
26
27impl AuditLogBufferConfig {
28 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
54pub struct AuditLogManager {
56 writer: Arc<aster_forge_runtime::BufferedBatchWriter<AuditLogCreate>>,
57}
58
59impl AuditLogManager {
60 pub fn new(db: DatabaseConnection) -> Self {
62 Self::with_config(db, AuditLogBufferConfig::default())
63 }
64
65 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 pub async fn record(&self, request: AuditLogCreate) {
93 self.writer.record(request).await;
94 }
95
96 pub async fn flush(&self) {
98 self.writer.flush().await;
99 }
100
101 pub fn cancel(&self) {
103 self.writer.cancel();
104 }
105}
106
107pub 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
121pub fn global_audit_log_manager() -> Option<&'static Arc<AuditLogManager>> {
123 GLOBAL_AUDIT_LOG_MANAGER.get()
124}
125
126pub 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
135pub async fn flush_global_audit_log_manager() {
137 if let Some(manager) = global_audit_log_manager() {
138 manager.flush().await;
139 }
140}
141
142pub 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
150pub 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}