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 #[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
55pub struct AuditLogManager {
57 writer: Arc<aster_forge_runtime::BufferedBatchWriter<AuditLogCreate>>,
58}
59
60impl AuditLogManager {
61 #[must_use]
63 pub fn new(db: DatabaseConnection) -> Self {
64 Self::with_config(db, AuditLogBufferConfig::default())
65 }
66
67 #[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 pub async fn record(&self, request: AuditLogCreate) {
96 self.writer.record(request).await;
97 }
98
99 pub async fn flush(&self) {
101 self.writer.flush().await;
102 }
103
104 pub fn cancel(&self) {
106 self.writer.cancel();
107 }
108}
109
110pub 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
124pub fn global_audit_log_manager() -> Option<&'static Arc<AuditLogManager>> {
126 GLOBAL_AUDIT_LOG_MANAGER.get()
127}
128
129pub 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
138pub async fn flush_global_audit_log_manager() {
140 if let Some(manager) = global_audit_log_manager() {
141 manager.flush().await;
142 }
143}
144
145pub 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
153pub 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}