1use std::collections::VecDeque;
30use std::io::{self};
31use std::sync::atomic::{AtomicU64, Ordering};
32use std::sync::{Arc, Condvar, Mutex, MutexGuard};
33use std::time::{Duration, Instant};
34
35pub const DEFAULT_DRAIN_TIMEOUT: Duration = Duration::from_micros(200);
37
38pub const DEFAULT_LEADER_TIMEOUT: Duration = Duration::from_millis(50);
40
41pub trait WalLike: Send + Sync {
45 fn flush_to_disk(&self) -> io::Result<()>;
47}
48
49impl WalLike for std::sync::Mutex<crate::wal::WAL> {
50 fn flush_to_disk(&self) -> io::Result<()> {
51 let mut wal = self
52 .lock()
53 .map_err(|e| io::Error::other(format!("WAL lock poisoned: {e}")))?;
54 wal.flush_to_disk()
55 }
56}
57
58#[derive(Debug, Clone, Copy)]
60pub struct GroupCommitConfig {
61 pub enabled: bool,
64 pub drain_timeout: Duration,
66 pub leader_timeout: Duration,
68}
69
70impl Default for GroupCommitConfig {
71 fn default() -> Self {
72 Self {
73 enabled: true,
74 drain_timeout: DEFAULT_DRAIN_TIMEOUT,
75 leader_timeout: DEFAULT_LEADER_TIMEOUT,
76 }
77 }
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub struct GroupCommitResult {
83 pub lsn: u64,
85 pub durable_batch_lsn: u64,
88 pub group_size: usize,
90}
91
92#[derive(Debug, Default, Clone, Copy)]
94pub struct GroupCommitStats {
95 pub batches: u64,
97 pub commits: u64,
99 pub total_group_entries: u64,
101}
102
103impl GroupCommitStats {
104 pub fn avg_group_size(&self) -> f64 {
106 if self.batches == 0 {
107 0.0
108 } else {
109 self.total_group_entries as f64 / self.batches as f64
110 }
111 }
112}
113
114struct PendingCommit {
119 lsn: u64,
120 slot: Arc<(Mutex<Option<io::Result<GroupCommitResult>>>, Condvar)>,
121}
122
123fn new_slot() -> Arc<(Mutex<Option<io::Result<GroupCommitResult>>>, Condvar)> {
124 Arc::new((Mutex::new(None), Condvar::new()))
125}
126
127pub struct GroupCommit<W: WalLike> {
129 wal: Arc<W>,
130 config: GroupCommitConfig,
131 next_lsn: AtomicU64,
133 queue: Mutex<VecDeque<PendingCommit>>,
134 leader_lock: Mutex<()>,
136 stats: Mutex<GroupCommitStats>,
137}
138
139impl<W: WalLike> GroupCommit<W> {
140 pub fn new(wal: Arc<W>, config: GroupCommitConfig) -> Self {
142 Self {
143 wal,
144 config,
145 next_lsn: AtomicU64::new(0),
146 queue: Mutex::new(VecDeque::new()),
147 leader_lock: Mutex::new(()),
148 stats: Mutex::new(GroupCommitStats::default()),
149 }
150 }
151
152 pub fn config(&self) -> &GroupCommitConfig {
153 &self.config
154 }
155
156 pub fn stats(&self) -> GroupCommitStats {
158 *self.stats.lock().unwrap()
159 }
160
161 pub fn flush(&self) -> io::Result<GroupCommitResult> {
167 let lsn = self.next_lsn.fetch_add(1, Ordering::Relaxed) + 1;
168 let slot = new_slot();
169 {
170 let mut queue = self.queue.lock().unwrap();
171 queue.push_back(PendingCommit {
172 lsn,
173 slot: slot.clone(),
174 });
175 }
176
177 if let Ok(guard) = self.leader_lock.try_lock() {
180 self.run_leader(guard);
181 }
182 let deadline = Instant::now() + self.config.leader_timeout;
183 let (lock, cvar) = &*slot;
184 let mut result = lock.lock().unwrap();
185 loop {
186 if let Some(outcome) = result.take() {
187 return match outcome {
188 Ok(ok) => Ok(ok),
189 Err(e) => Err(io::Error::new(e.kind(), e.to_string())),
190 };
191 }
192 if Instant::now() >= deadline {
193 drop(result);
196 if let Ok(guard) = self.leader_lock.try_lock() {
197 self.run_leader(guard);
198 }
199 result = lock.lock().unwrap();
200 continue;
201 }
202 let (guard, _to) = cvar.wait_timeout(result, Duration::from_millis(1)).unwrap();
203 result = guard;
204 }
205 }
206
207 fn run_leader(&self, _guard: MutexGuard<'_, ()>) {
213 let deadline = Instant::now() + self.config.drain_timeout;
214 let mut batch: Vec<PendingCommit> = Vec::new();
215
216 loop {
217 let mut drained = {
218 let mut queue = self.queue.lock().unwrap();
219 let mut v = Vec::with_capacity(queue.len());
220 while let Some(entry) = queue.pop_front() {
221 v.push(entry);
222 }
223 drop(queue);
224 v
225 };
226 if !drained.is_empty() {
227 batch.append(&mut drained);
228 }
229 if batch.is_empty() {
230 return;
233 }
234 if Instant::now() >= deadline {
235 break;
236 }
237 std::thread::sleep(Duration::from_micros(100));
238 }
239
240 let group_size = batch.len();
241 let durable_batch_lsn = batch.iter().map(|e| e.lsn).max().unwrap_or(0);
242 let io_result = self.wal.flush_to_disk();
243
244 {
245 let mut stats = self.stats.lock().unwrap();
246 stats.batches += 1;
247 stats.commits += group_size as u64;
248 stats.total_group_entries += group_size as u64;
249 }
250
251 for entry in batch {
252 let outcome = match &io_result {
253 Ok(()) => Ok(GroupCommitResult {
254 lsn: entry.lsn,
255 durable_batch_lsn,
256 group_size,
257 }),
258 Err(e) => Err(io::Error::new(e.kind(), format!("group fsync failed: {e}"))),
259 };
260 let (lock, cvar) = &*entry.slot;
261 let mut slot = lock.lock().unwrap();
262 *slot = Some(outcome);
263 drop(slot);
264 cvar.notify_all();
265 }
266 }
267}
268
269#[cfg(test)]
270mod tests {
271 use super::*;
272 use std::sync::atomic::AtomicU64;
273 use std::sync::atomic::Ordering as AtomicOrdering;
274 use std::thread;
275
276 struct MockWal {
279 fsyncs: AtomicU64,
280 delay: Duration,
281 }
282
283 impl WalLike for MockWal {
284 fn flush_to_disk(&self) -> io::Result<()> {
285 if !self.delay.is_zero() {
286 thread::sleep(self.delay);
287 }
288 self.fsyncs.fetch_add(1, AtomicOrdering::SeqCst);
289 Ok(())
290 }
291 }
292
293 fn make_gc(drain: Duration) -> (Arc<GroupCommit<MockWal>>, Arc<MockWal>) {
294 let wal = Arc::new(MockWal {
295 fsyncs: AtomicU64::new(0),
296 delay: Duration::from_millis(2),
297 });
298 let gc = Arc::new(GroupCommit::new(
299 wal.clone(),
300 GroupCommitConfig {
301 enabled: true,
302 drain_timeout: drain,
303 leader_timeout: Duration::from_millis(100),
304 },
305 ));
306 (gc, wal)
307 }
308
309 #[test]
310 fn test_single_commit_flushes_once() {
311 let (gc, wal) = make_gc(Duration::from_millis(2));
312 let result = gc.flush().unwrap();
313 assert_eq!(result.lsn, 1);
314 assert_eq!(result.durable_batch_lsn, 1);
315 assert_eq!(result.group_size, 1);
316 assert_eq!(wal.fsyncs.load(AtomicOrdering::SeqCst), 1);
317 let stats = gc.stats();
318 assert_eq!(stats.batches, 1);
319 assert_eq!(stats.commits, 1);
320 }
321
322 #[test]
323 fn test_concurrent_commits_coalesce_into_one_fsync() {
324 let (gc, wal) = make_gc(Duration::from_millis(10));
325 let threads: Vec<_> = (0..12)
326 .map(|_| {
327 let gc = gc.clone();
328 thread::spawn(move || gc.flush().unwrap())
329 })
330 .collect();
331 let results: Vec<GroupCommitResult> = threads.into_iter().map(|t| t.join().unwrap()).collect();
332
333 assert_eq!(results.len(), 12);
334 let mut sorted: Vec<u64> = results.iter().map(|r| r.lsn).collect();
336 sorted.sort_unstable();
337 assert_eq!(sorted, (1..=12).collect::<Vec<u64>>());
338 for r in &results {
340 assert!(r.durable_batch_lsn >= r.lsn);
341 }
342 assert!(results.iter().any(|r| r.group_size > 1), "expected coalescing");
344 let fsyncs = wal.fsyncs.load(AtomicOrdering::SeqCst);
346 assert!(fsyncs < 12, "group commit should batch fsyncs, got {fsyncs}");
347 assert_eq!(fsyncs, gc.stats().batches);
348 assert_eq!(gc.stats().commits, 12);
349 assert!(gc.stats().avg_group_size() > 1.0);
350 }
351
352 #[test]
353 fn test_concurrent_commits_no_coalescing_guarantee_needed() {
354 let (gc, wal) = make_gc(Duration::from_micros(0));
357 let threads: Vec<_> = (0..8)
358 .map(|_| {
359 let gc = gc.clone();
360 thread::spawn(move || gc.flush().unwrap())
361 })
362 .collect();
363 let results: Vec<GroupCommitResult> = threads.into_iter().map(|t| t.join().unwrap()).collect();
364 assert_eq!(results.len(), 8);
365 let mut sorted: Vec<u64> = results.iter().map(|r| r.lsn).collect();
366 sorted.sort_unstable();
367 assert_eq!(sorted, (1..=8).collect::<Vec<u64>>());
368 assert!(wal.fsyncs.load(AtomicOrdering::SeqCst) <= 8);
369 }
370
371 #[test]
372 fn test_stale_follower_self_heals() {
373 let (gc, wal) = make_gc(Duration::from_millis(1));
377 let garbage: Arc<(Mutex<Option<io::Result<GroupCommitResult>>>, Condvar)> = new_slot();
378 {
380 let mut queue = gc.queue.lock().unwrap();
381 queue.push_back(PendingCommit {
382 lsn: 999,
383 slot: garbage,
384 });
385 }
386 let result = gc.flush().unwrap();
389 assert!(result.lsn >= 1);
390 assert!(wal.fsyncs.load(AtomicOrdering::SeqCst) >= 1);
391 }
392
393 #[test]
394 fn test_fsync_error_propagates() {
395 struct FailingWal;
396 impl WalLike for FailingWal {
397 fn flush_to_disk(&self) -> io::Result<()> {
398 Err(io::Error::other("disk on fire"))
399 }
400 }
401 let gc = Arc::new(GroupCommit::new(Arc::new(FailingWal), GroupCommitConfig::default()));
402 let err = gc.flush().unwrap_err();
403 assert!(err.to_string().contains("disk on fire"));
404 }
405}