1use std::fs::{File, OpenOptions};
41use std::io::{Read, Seek, SeekFrom, Write};
42use std::path::{Path, PathBuf};
43use std::sync::atomic::{AtomicU64, Ordering};
44
45#[cfg(target_os = "linux")]
46use std::os::unix::fs::OpenOptionsExt as _;
47
48use crate::align::{AlignedBuf, DEFAULT_ALIGNMENT, is_aligned};
49use crate::error::{Result, WalError};
50use crate::record::{HEADER_SIZE, WalRecord};
51#[cfg(not(target_arch = "wasm32"))]
52use crate::record::{RecordHeader, WAL_MAGIC};
53
54const DWB_CAPACITY: usize = 64;
62
63const DWB_SLOT_PAYLOAD_MAX: usize = 64 * 1024;
65
66const DWB_SLOT_RAW: usize = 4 + HEADER_SIZE + DWB_SLOT_PAYLOAD_MAX;
68
69const DWB_SLOT_STRIDE: usize = round_up_const(DWB_SLOT_RAW, DEFAULT_ALIGNMENT);
73
74const DWB_HEADER_STRIDE: usize = DEFAULT_ALIGNMENT;
78const DWB_HEADER_FIELDS: usize = 12;
79const DWB_MAGIC: u32 = 0x4457_4246; static DWB_BYTES_WRITTEN_TOTAL: AtomicU64 = AtomicU64::new(0);
85
86pub fn wal_dwb_bytes_written_total() -> u64 {
88 DWB_BYTES_WRITTEN_TOTAL.load(Ordering::Relaxed)
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum DwbMode {
94 Off,
97 Buffered,
100 Direct,
103}
104
105impl DwbMode {
106 pub fn default_for_parent(parent_uses_direct_io: bool) -> Self {
110 if parent_uses_direct_io {
111 Self::Direct
112 } else {
113 Self::Buffered
114 }
115 }
116}
117
118const fn round_up_const(value: usize, align: usize) -> usize {
119 (value + align - 1) & !(align - 1)
120}
121
122pub const fn slot_stride() -> usize {
125 DWB_SLOT_STRIDE
126}
127
128fn slot_offset(idx: u32) -> u64 {
130 DWB_HEADER_STRIDE as u64 + (idx as u64 % DWB_CAPACITY as u64) * DWB_SLOT_STRIDE as u64
131}
132
133pub struct DoubleWriteBuffer {
135 file: File,
136 path: PathBuf,
137 mode: DwbMode,
138 write_pos: u32,
140 count: u32,
142 dirty: bool,
144 slot_buf: Option<AlignedBuf>,
147 header_buf: Option<AlignedBuf>,
149}
150
151impl std::fmt::Debug for DoubleWriteBuffer {
152 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
153 f.debug_struct("DoubleWriteBuffer")
154 .field("path", &self.path)
155 .field("mode", &self.mode)
156 .field("write_pos", &self.write_pos)
157 .field("count", &self.count)
158 .finish()
159 }
160}
161
162impl DoubleWriteBuffer {
163 pub fn open(path: &Path, mode: DwbMode) -> Result<Self> {
168 if mode == DwbMode::Off {
169 return Err(WalError::DwbOffNotOpenable);
170 }
171
172 let mut opts = OpenOptions::new();
173 opts.read(true).write(true).create(true).truncate(false);
174 #[cfg(target_os = "linux")]
175 if mode == DwbMode::Direct {
176 opts.custom_flags(libc::O_DIRECT);
177 }
178
179 let file = opts.open(path).map_err(|e| {
180 tracing::warn!(path = %path.display(), error = %e, mode = ?mode, "failed to open double-write buffer");
181 WalError::Io(e)
182 })?;
183
184 let (slot_buf, header_buf) = if mode == DwbMode::Direct {
185 (
186 Some(AlignedBuf::new(DWB_SLOT_STRIDE, DEFAULT_ALIGNMENT)?),
187 Some(AlignedBuf::new(DWB_HEADER_STRIDE, DEFAULT_ALIGNMENT)?),
188 )
189 } else {
190 (None, None)
191 };
192
193 let mut dwb = Self {
194 file,
195 path: path.to_path_buf(),
196 mode,
197 write_pos: 0,
198 count: 0,
199 dirty: false,
200 slot_buf,
201 header_buf,
202 };
203
204 let file_len = dwb.file.metadata().map(|m| m.len()).unwrap_or(0);
206 if file_len >= DWB_HEADER_STRIDE as u64 {
207 let mut block = vec![0u8; DWB_HEADER_STRIDE];
208 dwb.file.seek(SeekFrom::Start(0)).map_err(WalError::Io)?;
209 if dwb.file.read_exact(&mut block).is_ok() {
210 let mut arr4 = [0u8; 4];
211 arr4.copy_from_slice(&block[0..4]);
212 let magic = u32::from_le_bytes(arr4);
213 if magic == DWB_MAGIC {
214 arr4.copy_from_slice(&block[4..8]);
215 dwb.count = u32::from_le_bytes(arr4);
216 arr4.copy_from_slice(&block[8..12]);
217 dwb.write_pos = u32::from_le_bytes(arr4);
218 }
219 }
220 }
221
222 Ok(dwb)
223 }
224
225 pub fn mode(&self) -> DwbMode {
227 self.mode
228 }
229
230 pub fn write_record(&mut self, record: &WalRecord) -> Result<()> {
236 self.write_record_deferred(record)?;
237 self.flush()
238 }
239
240 pub fn write_record_deferred(&mut self, record: &WalRecord) -> Result<()> {
248 let total_size = HEADER_SIZE + record.payload.len();
249
250 if total_size > DWB_SLOT_PAYLOAD_MAX {
253 return Ok(()); }
255
256 let header_bytes = record.header.to_bytes();
257 let offset = slot_offset(self.write_pos);
258
259 match self.mode {
260 DwbMode::Off => unreachable!("Off never opens a DoubleWriteBuffer"),
261 DwbMode::Buffered => {
262 self.file
263 .seek(SeekFrom::Start(offset))
264 .map_err(WalError::Io)?;
265 self.file
266 .write_all(&(total_size as u32).to_le_bytes())
267 .map_err(WalError::Io)?;
268 self.file.write_all(&header_bytes).map_err(WalError::Io)?;
269 self.file.write_all(&record.payload).map_err(WalError::Io)?;
270 DWB_BYTES_WRITTEN_TOTAL.fetch_add(
271 (4 + header_bytes.len() + record.payload.len()) as u64,
272 Ordering::Relaxed,
273 );
274 }
275 DwbMode::Direct => {
276 let buf = self
277 .slot_buf
278 .as_mut()
279 .expect("slot_buf present in Direct mode");
280 buf.clear();
281 buf.write(&(total_size as u32).to_le_bytes());
282 buf.write(&header_bytes);
283 buf.write(&record.payload);
284 zero_tail(buf);
287 let slice = full_capacity_slice(buf);
288 debug_assert_eq!(slice.len(), DWB_SLOT_STRIDE);
289 debug_assert!(is_aligned(offset as usize, DEFAULT_ALIGNMENT));
290 pwrite_all(&self.file, slice, offset)?;
291 DWB_BYTES_WRITTEN_TOTAL.fetch_add(slice.len() as u64, Ordering::Relaxed);
292 }
293 }
294
295 self.write_pos = self.write_pos.wrapping_add(1);
296 self.count = self.count.saturating_add(1).min(DWB_CAPACITY as u32);
297 self.dirty = true;
298
299 Ok(())
300 }
301
302 pub fn flush(&mut self) -> Result<()> {
308 if !self.dirty {
309 return Ok(());
310 }
311
312 let mut header = [0u8; DWB_HEADER_FIELDS];
313 header[0..4].copy_from_slice(&DWB_MAGIC.to_le_bytes());
314 header[4..8].copy_from_slice(&self.count.to_le_bytes());
315 header[8..12].copy_from_slice(&self.write_pos.to_le_bytes());
316
317 match self.mode {
318 DwbMode::Off => unreachable!("invariant: flush() is gated on mode != Off by caller"),
319 DwbMode::Buffered => {
320 self.file.seek(SeekFrom::Start(0)).map_err(WalError::Io)?;
321 self.file.write_all(&header).map_err(WalError::Io)?;
322 DWB_BYTES_WRITTEN_TOTAL.fetch_add(header.len() as u64, Ordering::Relaxed);
323 }
324 DwbMode::Direct => {
325 let buf = self
326 .header_buf
327 .as_mut()
328 .expect("header_buf present in Direct mode");
329 buf.clear();
330 buf.write(&header);
331 zero_tail(buf);
332 let slice = full_capacity_slice(buf);
333 debug_assert_eq!(slice.len(), DWB_HEADER_STRIDE);
334 pwrite_all(&self.file, slice, 0)?;
335 DWB_BYTES_WRITTEN_TOTAL.fetch_add(slice.len() as u64, Ordering::Relaxed);
336 }
337 }
338
339 self.file.sync_all().map_err(WalError::Io)?;
340 self.dirty = false;
341
342 Ok(())
343 }
344
345 pub fn path(&self) -> &Path {
347 &self.path
348 }
349
350 pub fn recover_record(&mut self, target_lsn: u64) -> Result<Option<WalRecord>> {
358 #[cfg(target_arch = "wasm32")]
359 {
360 let _ = target_lsn;
361 return Ok(None);
362 }
363
364 #[cfg(not(target_arch = "wasm32"))]
365 {
366 use std::os::unix::io::AsRawFd as _;
367
368 let mut slot = AlignedBuf::new(DWB_SLOT_STRIDE, DEFAULT_ALIGNMENT)?;
371
372 for i in 0..DWB_CAPACITY as u32 {
373 let offset = slot_offset(i);
374 let read = unsafe {
376 libc::pread(
377 self.file.as_raw_fd(),
378 slot.as_mut_ptr() as *mut libc::c_void,
379 DWB_SLOT_STRIDE,
380 offset as libc::off_t,
381 )
382 };
383 if read <= 0 {
384 continue;
385 }
386 let bytes: &[u8] =
388 unsafe { std::slice::from_raw_parts(slot.as_ptr(), read as usize) };
389 if bytes.len() < 4 + HEADER_SIZE {
390 continue;
391 }
392
393 let mut arr4 = [0u8; 4];
394 arr4.copy_from_slice(&bytes[0..4]);
395 let total_size = u32::from_le_bytes(arr4) as usize;
396 if !(HEADER_SIZE..=DWB_SLOT_PAYLOAD_MAX).contains(&total_size)
397 || bytes.len() < 4 + total_size
398 {
399 continue;
400 }
401
402 let mut header_buf = [0u8; HEADER_SIZE];
403 header_buf.copy_from_slice(&bytes[4..4 + HEADER_SIZE]);
404 let header = RecordHeader::from_bytes(&header_buf);
405 if header.magic != WAL_MAGIC || header.lsn != target_lsn {
406 continue;
407 }
408
409 let payload_len = total_size - HEADER_SIZE;
410 let payload = bytes[4 + HEADER_SIZE..4 + HEADER_SIZE + payload_len].to_vec();
411 let record = WalRecord { header, payload };
412 if record.verify_checksum().is_ok() {
413 return Ok(Some(record));
414 }
415 }
416
417 Ok(None)
418 }
419 }
420}
421
422fn zero_tail(buf: &mut AlignedBuf) {
425 let written = buf.len();
426 let cap = buf.capacity();
427 if written < cap {
428 unsafe {
431 std::ptr::write_bytes(buf.as_mut_ptr().add(written), 0, cap - written);
432 }
433 }
434}
435
436fn full_capacity_slice(buf: &AlignedBuf) -> &[u8] {
439 unsafe { std::slice::from_raw_parts(buf.as_ptr(), buf.capacity()) }
442}
443
444fn pwrite_all(file: &File, data: &[u8], offset: u64) -> Result<()> {
446 #[cfg(not(target_arch = "wasm32"))]
447 {
448 use std::os::unix::io::AsRawFd as _;
449 let fd = file.as_raw_fd();
450 let mut remaining = data;
451 let mut write_offset = offset;
452 while !remaining.is_empty() {
453 let n = unsafe {
454 libc::pwrite(
455 fd,
456 remaining.as_ptr() as *const libc::c_void,
457 remaining.len(),
458 write_offset as libc::off_t,
459 )
460 };
461 if n < 0 {
462 return Err(WalError::Io(std::io::Error::last_os_error()));
463 }
464 let n = n as usize;
465 remaining = &remaining[n..];
466 write_offset += n as u64;
467 }
468 Ok(())
469 }
470 #[cfg(target_arch = "wasm32")]
471 {
472 let _ = (file, data, offset);
473 Err(WalError::Unsupported {
474 detail: "O_DIRECT pwrite not available on wasm32",
475 })
476 }
477}
478
479#[cfg(test)]
480mod tests {
481 use super::*;
482 use crate::record::{RecordType, WalRecordArgs};
483
484 fn open_buffered(path: &Path) -> DoubleWriteBuffer {
485 DoubleWriteBuffer::open(path, DwbMode::Buffered).unwrap()
486 }
487
488 #[test]
489 fn write_and_recover() {
490 let dir = tempfile::tempdir().unwrap();
491 let dwb_path = dir.path().join("test.dwb");
492
493 let mut dwb = open_buffered(&dwb_path);
494
495 let record = WalRecord::new(WalRecordArgs {
496 record_type: RecordType::Put as u32,
497 lsn: 42,
498 tenant_id: 1,
499 vshard_id: 0,
500 database_id: 0,
501 payload: b"hello double-write".to_vec(),
502 encryption_key: None,
503 preamble_bytes: None,
504 })
505 .unwrap();
506
507 dwb.write_record(&record).unwrap();
508
509 let recovered = dwb.recover_record(42).unwrap();
511 assert!(recovered.is_some());
512 let rec = recovered.unwrap();
513 assert_eq!(rec.header.lsn, 42);
514 assert_eq!(rec.payload, b"hello double-write");
515 }
516
517 #[test]
518 fn recover_nonexistent_returns_none() {
519 let dir = tempfile::tempdir().unwrap();
520 let dwb_path = dir.path().join("test2.dwb");
521
522 let mut dwb = open_buffered(&dwb_path);
523 let result = dwb.recover_record(999).unwrap();
524 assert!(result.is_none());
525 }
526
527 #[test]
528 fn survives_reopen() {
529 let dir = tempfile::tempdir().unwrap();
530 let dwb_path = dir.path().join("reopen.dwb");
531
532 {
533 let mut dwb = open_buffered(&dwb_path);
534 let record = WalRecord::new(WalRecordArgs {
535 record_type: RecordType::Put as u32,
536 lsn: 7,
537 tenant_id: 1,
538 vshard_id: 0,
539 database_id: 0,
540 payload: b"durable".to_vec(),
541 encryption_key: None,
542 preamble_bytes: None,
543 })
544 .unwrap();
545 dwb.write_record(&record).unwrap();
546 }
547
548 let mut dwb = open_buffered(&dwb_path);
549 let recovered = dwb.recover_record(7).unwrap();
550 assert!(recovered.is_some());
551 assert_eq!(recovered.unwrap().payload, b"durable");
552 }
553
554 #[test]
555 fn batch_deferred_writes_and_flush() {
556 let dir = tempfile::tempdir().unwrap();
557 let dwb_path = dir.path().join("batch.dwb");
558
559 let mut dwb = open_buffered(&dwb_path);
560
561 for lsn in 1..=5u64 {
562 let record = WalRecord::new(WalRecordArgs {
563 record_type: RecordType::Put as u32,
564 lsn,
565 tenant_id: 1,
566 vshard_id: 0,
567 database_id: 0,
568 payload: format!("batch-{lsn}").into_bytes(),
569 encryption_key: None,
570 preamble_bytes: None,
571 })
572 .unwrap();
573 dwb.write_record_deferred(&record).unwrap();
574 }
575
576 assert!(dwb.dirty);
577 dwb.flush().unwrap();
578 assert!(!dwb.dirty);
579
580 for lsn in 1..=5u64 {
581 let recovered = dwb.recover_record(lsn).unwrap();
582 assert!(recovered.is_some(), "LSN {lsn} should be recoverable");
583 assert_eq!(
584 recovered.unwrap().payload,
585 format!("batch-{lsn}").into_bytes()
586 );
587 }
588 }
589
590 #[test]
591 fn flush_is_idempotent() {
592 let dir = tempfile::tempdir().unwrap();
593 let dwb_path = dir.path().join("idem.dwb");
594
595 let mut dwb = open_buffered(&dwb_path);
596
597 dwb.flush().unwrap();
598 assert!(!dwb.dirty);
599
600 let record = WalRecord::new(WalRecordArgs {
601 record_type: RecordType::Put as u32,
602 lsn: 1,
603 tenant_id: 1,
604 vshard_id: 0,
605 database_id: 0,
606 payload: b"data".to_vec(),
607 encryption_key: None,
608 preamble_bytes: None,
609 })
610 .unwrap();
611 dwb.write_record_deferred(&record).unwrap();
612 dwb.flush().unwrap();
613 dwb.flush().unwrap();
614 assert!(!dwb.dirty);
615 }
616
617 #[test]
618 fn slot_stride_is_o_direct_aligned() {
619 assert!(
625 is_aligned(DWB_SLOT_STRIDE, DEFAULT_ALIGNMENT),
626 "DWB slot stride {DWB_SLOT_STRIDE} bytes is not a multiple of {DEFAULT_ALIGNMENT}"
627 );
628 assert!(is_aligned(DWB_HEADER_STRIDE, DEFAULT_ALIGNMENT));
629 for i in 0..DWB_CAPACITY as u32 {
630 assert!(is_aligned(slot_offset(i) as usize, DEFAULT_ALIGNMENT));
631 }
632 }
633
634 #[test]
635 fn recover_after_wraparound() {
636 let dir = tempfile::tempdir().unwrap();
637 let dwb_path = dir.path().join("wrap.dwb");
638
639 let mut dwb = open_buffered(&dwb_path);
640
641 let total = DWB_CAPACITY as u64 + 5;
642 for lsn in 1..=total {
643 let record = WalRecord::new(WalRecordArgs {
644 record_type: RecordType::Put as u32,
645 lsn,
646 tenant_id: 1,
647 vshard_id: 0,
648 database_id: 0,
649 payload: format!("wrap-{lsn}").into_bytes(),
650 encryption_key: None,
651 preamble_bytes: None,
652 })
653 .unwrap();
654 dwb.write_record_deferred(&record).unwrap();
655 }
656 dwb.flush().unwrap();
657
658 for lsn in (total - 4)..=total {
659 let recovered = dwb.recover_record(lsn).unwrap();
660 assert!(
661 recovered.is_some(),
662 "LSN {lsn} should be recoverable after wrap-around"
663 );
664 assert_eq!(
665 recovered.unwrap().payload,
666 format!("wrap-{lsn}").into_bytes()
667 );
668 }
669
670 for lsn in 1..=5u64 {
671 let recovered = dwb.recover_record(lsn).unwrap();
672 assert!(
673 recovered.is_none(),
674 "LSN {lsn} should have been overwritten by wrap-around"
675 );
676 }
677 }
678
679 #[test]
680 fn bytes_written_counter_increments() {
681 let dir = tempfile::tempdir().unwrap();
682 let dwb_path = dir.path().join("counter.dwb");
683 let before = wal_dwb_bytes_written_total();
684
685 let mut dwb = open_buffered(&dwb_path);
686 let rec = WalRecord::new(WalRecordArgs {
687 record_type: RecordType::Put as u32,
688 lsn: 1,
689 tenant_id: 1,
690 vshard_id: 0,
691 database_id: 0,
692 payload: b"counted".to_vec(),
693 encryption_key: None,
694 preamble_bytes: None,
695 })
696 .unwrap();
697 dwb.write_record(&rec).unwrap();
698
699 assert!(wal_dwb_bytes_written_total() > before);
700 }
701}