feoxdb/storage/
write_buffer.rs

1use ahash::RandomState;
2use crossbeam_channel::{bounded, Receiver, Sender};
3use crossbeam_utils::CachePadded;
4use parking_lot::{Mutex, RwLock};
5use std::collections::VecDeque;
6use std::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering};
7use std::sync::Arc;
8use std::thread::{self, JoinHandle};
9use std::time::{Duration, Instant};
10
11use crate::constants::*;
12use crate::core::record::Record;
13use crate::error::{FeoxError, Result};
14use crate::stats::Statistics;
15use crate::storage::format::{get_format, RecordFormat};
16use crate::storage::free_space::FreeSpaceManager;
17use crate::storage::io::DiskIO;
18
19/// Sharded write buffer for reducing contention
20/// Each thread consistently uses the same shard to improve cache locality
21#[repr(align(64))] // Cache line alignment
22pub struct ShardedWriteBuffer {
23    /// Buffered writes pending flush
24    buffer: Mutex<VecDeque<WriteEntry>>,
25
26    /// Number of entries in buffer
27    count: AtomicUsize,
28
29    /// Total size of buffered data
30    size: AtomicUsize,
31}
32
33/// Write entry for buffered operations
34pub struct WriteEntry {
35    pub op: Operation,
36    pub record: Arc<Record>,
37    pub old_value_len: usize,
38    pub work_status: AtomicU32,
39    pub retry_count: AtomicU32,
40    pub timestamp: Instant,
41}
42
43/// Main write buffer coordinator
44pub struct WriteBuffer {
45    /// Sharded buffers to reduce contention between threads
46    sharded_buffers: Arc<Vec<CachePadded<ShardedWriteBuffer>>>,
47
48    /// Shared disk I/O handle
49    disk_io: Arc<RwLock<DiskIO>>,
50
51    /// Free space manager for sector allocation
52    free_space: Arc<RwLock<FreeSpaceManager>>,
53
54    /// Per-worker channels for targeted flush requests
55    worker_channels: Vec<Sender<FlushRequest>>,
56
57    /// Background worker handles
58    worker_handles: Vec<JoinHandle<()>>,
59
60    /// Periodic flush thread handle
61    periodic_flush_handle: Option<JoinHandle<()>>,
62
63    /// Shutdown flag
64    shutdown: Arc<AtomicBool>,
65
66    /// Shared statistics
67    stats: Arc<Statistics>,
68
69    /// Stable per-store key hasher for preserving per-key write order
70    shard_hasher: RandomState,
71
72    /// Format version for record serialization
73    format_version: u32,
74}
75
76#[derive(Debug)]
77struct FlushRequest {
78    response: Option<Sender<Result<()>>>,
79}
80
81struct WorkerContext {
82    worker_id: usize,
83    disk_io: Arc<RwLock<DiskIO>>,
84    free_space: Arc<RwLock<FreeSpaceManager>>,
85    sharded_buffers: Arc<Vec<CachePadded<ShardedWriteBuffer>>>,
86    shutdown: Arc<AtomicBool>,
87    stats: Arc<Statistics>,
88    format_version: u32,
89}
90
91impl ShardedWriteBuffer {
92    fn new(_shard_id: usize) -> Self {
93        Self {
94            buffer: Mutex::new(VecDeque::new()),
95            count: AtomicUsize::new(0),
96            size: AtomicUsize::new(0),
97        }
98    }
99
100    fn add_entry(&self, entry: WriteEntry) -> Result<()> {
101        let entry_size = entry.record.calculate_size();
102
103        let mut buffer = self.buffer.lock();
104        buffer.push_back(entry);
105
106        self.count.fetch_add(1, Ordering::AcqRel);
107        self.size.fetch_add(entry_size, Ordering::AcqRel);
108
109        Ok(())
110    }
111
112    fn drain_entries(&self) -> Vec<WriteEntry> {
113        let mut buffer = self.buffer.lock();
114        let entries: Vec<_> = buffer.drain(..).collect();
115
116        self.count.store(0, Ordering::Release);
117        self.size.store(0, Ordering::Release);
118
119        entries
120    }
121
122    fn is_full(&self) -> bool {
123        self.count.load(Ordering::Acquire) >= WRITE_BUFFER_SIZE
124            || self.size.load(Ordering::Acquire) >= FEOX_WRITE_BUFFER_SIZE
125    }
126}
127
128impl WriteBuffer {
129    pub fn new(
130        disk_io: Arc<RwLock<DiskIO>>,
131        free_space: Arc<RwLock<FreeSpaceManager>>,
132        stats: Arc<Statistics>,
133        format_version: u32,
134    ) -> Self {
135        // Use half CPU count for both shards and workers
136        let num_shards = (num_cpus::get() / 2).max(1);
137
138        let sharded_buffers = Arc::new(
139            (0..num_shards)
140                .map(|shard_id| CachePadded::new(ShardedWriteBuffer::new(shard_id)))
141                .collect(),
142        );
143
144        Self {
145            sharded_buffers,
146            disk_io,
147            free_space,
148            worker_channels: Vec::new(),
149            worker_handles: Vec::new(),
150            periodic_flush_handle: None,
151            shutdown: Arc::new(AtomicBool::new(false)),
152            stats,
153            shard_hasher: RandomState::new(),
154            format_version,
155        }
156    }
157
158    /// Add write operation to buffer (lock-free fast path)
159    pub fn add_write(
160        &self,
161        op: Operation,
162        record: Arc<Record>,
163        old_value_len: usize,
164    ) -> Result<()> {
165        if self.shutdown.load(Ordering::Acquire) {
166            return Err(FeoxError::ShuttingDown);
167        }
168
169        let shard_id = self.get_shard_id(&record.key);
170        let entry = WriteEntry {
171            op,
172            record,
173            old_value_len,
174            work_status: AtomicU32::new(0),
175            retry_count: AtomicU32::new(0),
176            timestamp: Instant::now(),
177        };
178
179        let buffer = &self.sharded_buffers[shard_id];
180
181        buffer.add_entry(entry)?;
182        self.stats.record_write_buffered();
183
184        // Check if we need to trigger flush for this specific shard
185        if buffer.is_full() && shard_id < self.worker_channels.len() {
186            let req = FlushRequest { response: None };
187            let _ = self.worker_channels[shard_id].try_send(req);
188        }
189
190        Ok(())
191    }
192
193    /// Start background worker threads
194    pub fn start_workers(&mut self, num_workers: usize) {
195        // Ensure we have the right number of workers for shards
196        let num_shards = self.sharded_buffers.len();
197        let actual_workers = num_workers.min(num_shards);
198
199        // Create per-worker channels
200        let mut receivers = Vec::new();
201        for _ in 0..actual_workers {
202            let (tx, rx) = bounded(2);
203            self.worker_channels.push(tx);
204            receivers.push(rx);
205        }
206
207        // Start workers, each owning one shard
208        for (worker_id, flush_rx) in receivers.into_iter().enumerate() {
209            let ctx = WorkerContext {
210                worker_id,
211                disk_io: self.disk_io.clone(),
212                free_space: self.free_space.clone(),
213                sharded_buffers: self.sharded_buffers.clone(),
214                shutdown: self.shutdown.clone(),
215                stats: self.stats.clone(),
216                format_version: self.format_version,
217            };
218
219            let handle = thread::spawn(move || {
220                write_buffer_worker(ctx, flush_rx);
221            });
222
223            self.worker_handles.push(handle);
224        }
225
226        // Start periodic flush coordinator
227        let worker_channels = self.worker_channels.clone();
228        let shutdown = self.shutdown.clone();
229        let sharded_buffers = self.sharded_buffers.clone();
230
231        let periodic_handle = thread::spawn(move || {
232            let interval = WRITE_BUFFER_FLUSH_INTERVAL;
233
234            while !shutdown.load(Ordering::Acquire) {
235                thread::sleep(interval);
236
237                // Check each shard and trigger its worker if needed
238                for (shard_id, buffer) in sharded_buffers.iter().enumerate() {
239                    let count = buffer.count.load(Ordering::Relaxed);
240                    if count > 0 && shard_id < worker_channels.len() {
241                        let req = FlushRequest { response: None };
242                        let _ = worker_channels[shard_id].try_send(req);
243                    }
244                }
245            }
246        });
247
248        self.periodic_flush_handle = Some(periodic_handle);
249    }
250
251    /// Force flush and wait for completion
252    pub fn force_flush(&self) -> Result<()> {
253        let mut responses = Vec::new();
254
255        // Send flush request to each worker and collect response channels
256        for worker_tx in &self.worker_channels {
257            let (tx, rx) = bounded(1);
258            let req = FlushRequest { response: Some(tx) };
259
260            worker_tx.send(req).map_err(|_| FeoxError::ChannelError)?;
261            responses.push(rx);
262        }
263
264        // Wait for all workers to complete
265        for rx in responses {
266            rx.recv().map_err(|_| FeoxError::ChannelError)??;
267        }
268
269        Ok(())
270    }
271
272    /// Shutdown write buffer
273    pub fn initiate_shutdown(&self) {
274        self.shutdown.store(true, Ordering::Release);
275
276        // Don't call force_flush here as it can block
277        // Workers will see the shutdown flag and exit gracefully
278    }
279
280    /// Complete shutdown - must be called after initiate_shutdown
281    pub fn complete_shutdown(&mut self) {
282        use std::time::Duration;
283
284        // Ensure shutdown flag is set
285        self.shutdown.store(true, Ordering::Release);
286
287        // Wait for periodic flush thread to finish with timeout
288        if let Some(handle) = self.periodic_flush_handle.take() {
289            // Spawn a thread to wait with timeout since JoinHandle doesn't have join_timeout
290            let (tx, rx) = crossbeam_channel::bounded(1);
291            thread::spawn(move || {
292                let _ = handle.join();
293                let _ = tx.send(());
294            });
295
296            if rx.recv_timeout(Duration::from_secs(5)).is_err() {
297                // Timeout waiting for periodic flush thread
298            }
299        }
300
301        // Signal workers to stop and wait
302        for handle in self.worker_handles.drain(..) {
303            let _ = handle.join();
304        }
305
306        // Note: disk_io shutdown is handled by the Store's Drop implementation
307        // to ensure proper ordering
308    }
309
310    /// Legacy shutdown for compatibility
311    pub fn shutdown(&mut self) {
312        self.complete_shutdown();
313    }
314
315    #[inline]
316    fn get_shard_id(&self, key: &[u8]) -> usize {
317        self.shard_hasher.hash_one(key) as usize % self.sharded_buffers.len()
318    }
319}
320
321/// Background worker for processing write buffer flushes
322fn write_buffer_worker(ctx: WorkerContext, flush_rx: Receiver<FlushRequest>) {
323    let worker_id = ctx.worker_id;
324    let format = get_format(ctx.format_version);
325
326    loop {
327        if ctx.shutdown.load(Ordering::Acquire) {
328            break;
329        }
330
331        // Wait for flush request with timeout to check shutdown periodically
332        let req = match flush_rx.recv_timeout(Duration::from_millis(500)) {
333            Ok(req) => req,
334            Err(crossbeam_channel::RecvTimeoutError::Timeout) => {
335                continue;
336            }
337            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
338                break;
339            }
340        };
341
342        // Drain only this worker's shard
343        if worker_id < ctx.sharded_buffers.len() {
344            let buffer = &ctx.sharded_buffers[worker_id];
345            let entries = buffer.drain_entries();
346
347            if !entries.is_empty() {
348                // Process writes for this shard
349                let result = process_write_batch(
350                    &ctx.disk_io,
351                    &ctx.free_space,
352                    entries,
353                    &ctx.stats,
354                    format.as_ref(),
355                );
356
357                ctx.stats.flush_count.fetch_add(1, Ordering::Relaxed);
358
359                // Send response if requested
360                if let Some(tx) = req.response {
361                    let _ = tx.send(result);
362                }
363            } else if let Some(tx) = req.response {
364                // No data to flush, but still respond if needed
365                let _ = tx.send(Ok(()));
366            }
367        }
368    }
369
370    // Before exiting, flush any remaining data from this worker's shard
371    if ctx.shutdown.load(Ordering::Acquire) && worker_id < ctx.sharded_buffers.len() {
372        let buffer = &ctx.sharded_buffers[worker_id];
373        let final_entries = buffer.drain_entries();
374
375        if !final_entries.is_empty() {
376            let _ = process_write_batch(
377                &ctx.disk_io,
378                &ctx.free_space,
379                final_entries,
380                &ctx.stats,
381                format.as_ref(),
382            );
383        }
384    }
385}
386
387/// Process a batch of write entries
388fn process_write_batch(
389    disk_io: &Arc<RwLock<DiskIO>>,
390    free_space: &Arc<RwLock<FreeSpaceManager>>,
391    entries: Vec<WriteEntry>,
392    stats: &Arc<Statistics>,
393    format: &dyn RecordFormat,
394) -> Result<()> {
395    let mut pending_writes = Vec::new();
396    let mut batch_writes = Vec::new();
397    let mut delete_operations = Vec::new();
398    let mut records_to_clear = Vec::new();
399
400    // Prepare all operations
401    for entry in entries {
402        match entry.op {
403            Operation::Insert | Operation::Update => {
404                // Check if record is still valid (not deleted)
405                if entry.record.refcount.load(Ordering::Acquire) > 0
406                    && entry.record.sector.load(Ordering::Acquire) == 0
407                {
408                    let data = prepare_record_data(&entry.record, format)?;
409                    let sectors_needed = data.len().div_ceil(FEOX_BLOCK_SIZE);
410                    pending_writes.push((data, sectors_needed, entry.record));
411                }
412            }
413            Operation::Delete => {
414                let sector = entry.record.sector.load(Ordering::Acquire);
415                if sector != 0 {
416                    let sectors_needed = format
417                        .total_size(entry.record.key.len(), entry.old_value_len)
418                        .div_ceil(FEOX_BLOCK_SIZE);
419                    delete_operations.push((sector, sectors_needed));
420                }
421            }
422            _ => {}
423        }
424    }
425
426    for (data, sectors_needed, record) in pending_writes {
427        let sector = match free_space.write().allocate_sectors(sectors_needed as u64) {
428            Ok(sector) => sector,
429            Err(error) => {
430                release_allocations(free_space, &records_to_clear, stats)?;
431                return Err(error);
432            }
433        };
434        stats
435            .disk_usage
436            .fetch_add((sectors_needed * FEOX_BLOCK_SIZE) as u64, Ordering::Relaxed);
437        batch_writes.push((sector, data));
438        records_to_clear.push((sector, sectors_needed, record));
439    }
440
441    // Process batch writes with io_uring
442    if !batch_writes.is_empty() {
443        if !delete_operations.is_empty() {
444            crash_at("before_replacement_write");
445        }
446
447        // Use io_uring with retry
448        let mut retries = 3;
449        let mut delay_us = 100;
450
451        while retries > 0 {
452            // Execute batch write
453            let result = disk_io.write().batch_write(batch_writes.clone());
454
455            match result {
456                Ok(()) => {
457                    if !delete_operations.is_empty() {
458                        crash_at("after_replacement_write");
459                    }
460                    for (sector, _, record) in &records_to_clear {
461                        record.sector.store(*sector, Ordering::Release);
462                        std::sync::atomic::fence(Ordering::Release);
463                        record.clear_value();
464                    }
465                    stats.record_write_flushed(records_to_clear.len() as u64);
466                    break;
467                }
468                Err(e) => {
469                    retries -= 1;
470                    if retries > 0 {
471                        // Exponential backoff with jitter ±10%
472                        let jitter = {
473                            use rand::Rng;
474                            let mut rng = rand::rng();
475                            (delay_us * rng.random_range(-10..=10)) / 100
476                        };
477                        let actual_delay = (delay_us + jitter).max(1);
478                        thread::sleep(Duration::from_micros(actual_delay as u64));
479                        delay_us *= 2;
480                    } else {
481                        stats.record_write_failed();
482                        release_allocations(free_space, &records_to_clear, stats)?;
483                        return Err(e);
484                    }
485                }
486            }
487        }
488    }
489
490    for (sector, sectors_needed) in delete_operations {
491        let mut deletion_marker = vec![0u8; FEOX_BLOCK_SIZE];
492        deletion_marker[..8].copy_from_slice(b"\0DELETED");
493
494        disk_io
495            .write()
496            .write_sectors_sync(sector, &deletion_marker)?;
497
498        free_space
499            .write()
500            .release_sectors(sector, sectors_needed as u64)?;
501
502        stats
503            .disk_usage
504            .fetch_sub((sectors_needed * FEOX_BLOCK_SIZE) as u64, Ordering::Relaxed);
505    }
506
507    Ok(())
508}
509
510#[cfg(test)]
511fn crash_at(point: &str) {
512    if std::env::var("FEOX_TEST_CRASH_POINT").as_deref() == Ok(point) {
513        std::process::exit(86);
514    }
515}
516
517#[cfg(not(test))]
518#[inline]
519fn crash_at(_: &str) {}
520
521fn release_allocations(
522    free_space: &Arc<RwLock<FreeSpaceManager>>,
523    allocations: &[(u64, usize, Arc<Record>)],
524    stats: &Statistics,
525) -> Result<()> {
526    for (sector, sectors_needed, _) in allocations {
527        free_space
528            .write()
529            .release_sectors(*sector, *sectors_needed as u64)?;
530        stats.disk_usage.fetch_sub(
531            (*sectors_needed * FEOX_BLOCK_SIZE) as u64,
532            Ordering::Relaxed,
533        );
534    }
535    Ok(())
536}
537
538fn prepare_record_data(record: &Record, format: &dyn RecordFormat) -> Result<Vec<u8>> {
539    // Get the total size using the format trait
540    let total_size = format.total_size(record.key.len(), record.value_len);
541
542    // Calculate padded size to sector boundary
543    let sectors_needed = total_size.div_ceil(FEOX_BLOCK_SIZE);
544    let padded_size = sectors_needed * FEOX_BLOCK_SIZE;
545
546    let mut data = Vec::with_capacity(padded_size);
547
548    // Sector header
549    data.extend_from_slice(&SECTOR_MARKER.to_le_bytes());
550    data.extend_from_slice(&0u16.to_le_bytes()); // seq_number
551
552    // Use format trait to serialize the record
553    let record_data = format.serialize_record(record, true);
554    data.extend_from_slice(&record_data);
555
556    // Pad to sector boundary
557    data.resize(padded_size, 0);
558
559    Ok(data)
560}
561
562impl Drop for WriteBuffer {
563    fn drop(&mut self) {
564        if !self.shutdown.load(Ordering::Acquire) {
565            self.complete_shutdown();
566        }
567    }
568}
feoxdb/storage/write_buffer.rs

feoxdb/storage/
write_buffer.rs
