linkrs_storage/engine/
wal_manager.rs

1//! WAL Manager
2//!
3//! Unified WAL (Write-Ahead Log) manager that properly integrates with LocalWalWriter.
4//! This module provides a single source of truth for LSN management and WAL operations.
5
6use crate::index::shard_runtime::IndexBarrierRegistry;
7use linkrs_core::types::{CommitLsn, Timestamp, TransactionId};
8use linkrs_core::wal::types::WalOpType;
9use linkrs_core::wal::OutboxIntent;
10use linkrs_core::{StorageError, StorageResult};
11use linkrs_transaction::wal::writer::WalWriter;
12use linkrs_transaction::wal::TransactionWalEntry;
13use linkrs_transaction::wal::{LocalWalWriter, Lsn, WalConfig};
14use parking_lot::Mutex;
15use postcard::to_allocvec;
16use serde::Serialize;
17use std::path::Path;
18use std::sync::atomic::{AtomicU64, Ordering};
19use std::sync::Arc;
20
21/// Unified WAL manager that wraps LocalWalWriter
22///
23/// This manager ensures LSN consistency by delegating all LSN operations
24/// to the underlying LocalWalWriter, avoiding the dual LSN tracking issue.
25pub struct WalManager {
26    local_writer: Option<Arc<Mutex<LocalWalWriter>>>,
27    barrier_registry: Option<IndexBarrierRegistry>,
28    config: WalConfig,
29    sync_count: AtomicU64,
30    sync_failures: AtomicU64,
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct WalMetrics {
35    pub accepted_lsn: Lsn,
36    pub durable_lsn: Lsn,
37    pub sync_count: u64,
38    pub sync_failures: u64,
39}
40
41impl WalManager {
42    pub fn new() -> Self {
43        Self {
44            local_writer: None,
45            barrier_registry: None,
46            config: WalConfig::default(),
47            sync_count: AtomicU64::new(0),
48            sync_failures: AtomicU64::new(0),
49        }
50    }
51
52    pub fn with_config(config: WalConfig) -> Self {
53        Self {
54            local_writer: None,
55            barrier_registry: None,
56            config,
57            sync_count: AtomicU64::new(0),
58            sync_failures: AtomicU64::new(0),
59        }
60    }
61
62    pub fn open(&mut self, wal_dir: &Path, thread_id: u32) -> StorageResult<()> {
63        let wal_uri = wal_dir.to_string_lossy().to_string();
64        let mut writer = LocalWalWriter::with_config(&wal_uri, thread_id, self.config.clone());
65        writer
66            .open()
67            .map_err(|e| StorageError::wal_error(format!("Failed to open WAL: {:?}", e)))?;
68        if self.config.group_commit_enabled {
69            writer.enable_group_commit().map_err(|e| {
70                StorageError::wal_error(format!("Failed to enable group commit: {:?}", e))
71            })?;
72        }
73        // Production async-flush path: stage appends in the per-thread buffer
74        // and drain them from a background thread. Disabled via
75        // `WalConfig::enable_async_flush = false` for exact sync behavior.
76        if self.config.enable_async_flush {
77            writer.enable_async_flush();
78            if writer.is_async_enabled() {
79                // A missing file handle (e.g. in-memory mode) makes background
80                // flush impossible; fall back to synchronous drains in sync().
81                let _ = writer.start_background_flush();
82            }
83        }
84        self.local_writer = Some(Arc::new(Mutex::new(writer)));
85        Ok(())
86    }
87
88    pub(crate) fn set_index_barrier_registry(&mut self, registry: IndexBarrierRegistry) {
89        self.barrier_registry = Some(registry);
90    }
91
92    fn truncation_barrier_lsn(&self) -> Option<Lsn> {
93        self.barrier_registry.as_ref().and_then(|registry| {
94            registry
95                .read()
96                .values()
97                .map(|barrier| Lsn::new(barrier.get()))
98                .min()
99        })
100    }
101
102    pub fn current_lsn(&self) -> Lsn {
103        if let Some(ref writer) = self.local_writer {
104            writer.lock().current_lsn()
105        } else {
106            Lsn::ZERO
107        }
108    }
109
110    /// Return the latest WAL position confirmed durable by the writer.
111    pub fn durable_lsn(&self) -> Lsn {
112        if let Some(ref writer) = self.local_writer {
113            writer.lock().durable_lsn()
114        } else {
115            Lsn::ZERO
116        }
117    }
118
119    pub fn metrics(&self) -> WalMetrics {
120        WalMetrics {
121            accepted_lsn: self.current_lsn(),
122            durable_lsn: self.durable_lsn(),
123            sync_count: self.sync_count.load(Ordering::Relaxed),
124            sync_failures: self.sync_failures.load(Ordering::Relaxed),
125        }
126    }
127
128    /// Whether the background async-flush path is active.
129    pub fn is_async_enabled(&self) -> bool {
130        if let Some(ref writer) = self.local_writer {
131            writer.lock().is_async_enabled()
132        } else {
133            false
134        }
135    }
136
137    /// Drain buffered appends and stop the background flush thread.
138    /// Idempotent; safe to call on an unopened or sync-mode manager.
139    pub fn shutdown(&self) {
140        if let Some(ref writer) = self.local_writer {
141            let mut guard = writer.lock();
142            let _ = guard.flush_and_sync();
143            let _ = guard.disable_async_flush();
144        }
145    }
146
147    pub fn sync(&self) -> StorageResult<()> {
148        if let Some(ref writer) = self.local_writer {
149            let result = writer
150                .lock()
151                .sync()
152                .map_err(|e| StorageError::wal_error(format!("Failed to sync WAL: {:?}", e)));
153            match result {
154                Ok(()) => {
155                    self.sync_count.fetch_add(1, Ordering::Relaxed);
156                }
157                Err(error) => {
158                    self.sync_failures.fetch_add(1, Ordering::Relaxed);
159                    return Err(error);
160                }
161            }
162        }
163        Ok(())
164    }
165
166    pub fn append_redo<T: Serialize>(
167        &self,
168        op_type: WalOpType,
169        timestamp: Timestamp,
170        redo: &T,
171    ) -> StorageResult<()> {
172        let Some(writer) = self.local_writer.as_ref() else {
173            return Err(StorageError::wal_error(
174                "WAL writer is not initialized".to_string(),
175            ));
176        };
177
178        let payload = to_allocvec(redo).map_err(|e| {
179            StorageError::serialize_error(format!("Failed to serialize WAL redo: {}", e))
180        })?;
181
182        writer
183            .lock()
184            .append_entry(op_type, timestamp, &payload)
185            .map_err(|e| StorageError::wal_error(format!("Failed to append WAL entry: {}", e)))?;
186
187        Ok(())
188    }
189
190    pub fn append_entry(
191        &self,
192        op_type: WalOpType,
193        timestamp: Timestamp,
194        payload: &[u8],
195    ) -> StorageResult<()> {
196        let Some(writer) = self.local_writer.as_ref() else {
197            return Err(StorageError::wal_error(
198                "WAL writer is not initialized".to_string(),
199            ));
200        };
201        writer
202            .lock()
203            .append_entry(op_type, timestamp, payload)
204            .map_err(|e| StorageError::wal_error(format!("Failed to append WAL entry: {}", e)))
205    }
206
207    pub fn append_transaction_with_durability(
208        &self,
209        transaction_id: TransactionId,
210        entries: Vec<TransactionWalEntry>,
211        intents: &[OutboxIntent],
212        durability: linkrs_core::types::DurabilityLevel,
213    ) -> StorageResult<CommitLsn> {
214        let Some(writer) = self.local_writer.as_ref() else {
215            return Err(StorageError::wal_error(
216                "WAL writer is not initialized".to_string(),
217            ));
218        };
219
220        // Append the commit batch under the writer lock, then release it
221        // before waiting for durability. For Sync durability this lets
222        // concurrent commits batch into a single fsync through the
223        // group-commit coordinator instead of serializing one fsync per
224        // transaction at the writer lock.
225        let lsn = {
226            let mut guard = writer.lock();
227            guard
228                .append_transaction_batch_no_wait(transaction_id, entries, intents)
229                .map_err(|error| {
230                    StorageError::wal_error(format!(
231                        "Failed to append committed WAL transaction: {}",
232                        error
233                    ))
234                })?
235        };
236
237        // Record the append in the group-commit coordinator for every
238        // durability level when group commit is enabled. For Sync this is the
239        // waiting path below; for Async this lets the transaction's WAL bytes
240        // be covered by the next leader's fsync (free durability piggyback,
241        // amortizing fsyncs across Sync + Async writers).
242        let coordinator = writer.lock().group_commit_coordinator().cloned();
243        if let Some(ref coordinator) = coordinator {
244            coordinator.record_appended(lsn.get());
245        }
246
247        if matches!(durability, linkrs_core::types::DurabilityLevel::Sync) {
248            if let Some(ref coordinator) = coordinator {
249                coordinator.append_and_wait(lsn.get()).map_err(|error| {
250                    StorageError::wal_error(format!(
251                        "Failed to await durable WAL transaction: {}",
252                        error
253                    ))
254                })?;
255            } else {
256                writer.lock().wait_for_durable(lsn.get()).map_err(|error| {
257                    StorageError::wal_error(format!(
258                        "Failed to sync committed WAL transaction: {}",
259                        error
260                    ))
261                })?;
262            }
263        }
264
265        Ok(lsn)
266    }
267
268    pub fn set_checkpoint_seq(&self, seq: u64) -> StorageResult<()> {
269        if let Some(ref writer) = self.local_writer {
270            writer.lock().set_checkpoint_seq(seq).map_err(|e| {
271                StorageError::wal_error(format!("Failed to update checkpoint seq: {:?}", e))
272            })?;
273        }
274        Ok(())
275    }
276
277    /// Restore the logical WAL baseline covered by a durable checkpoint.
278    pub fn set_recovery_baseline_lsn(&self, lsn: Lsn) -> StorageResult<()> {
279        if let Some(ref writer) = self.local_writer {
280            writer.lock().set_recovery_baseline_lsn(lsn).map_err(|e| {
281                StorageError::wal_error(format!("Failed to restore WAL recovery baseline: {:?}", e))
282            })?;
283        }
284        Ok(())
285    }
286
287    pub fn truncate(&self, lsn: Lsn) -> StorageResult<()> {
288        let truncation_lsn = self
289            .truncation_barrier_lsn()
290            .map_or(lsn, |barrier| lsn.min(barrier));
291        if let Some(ref writer) = self.local_writer {
292            writer.lock().truncate(truncation_lsn).map_err(|e| {
293                StorageError::wal_error(format!(
294                    "Failed to truncate WAL at {} (requested {}): {:?}",
295                    truncation_lsn, lsn, e
296                ))
297            })?;
298        }
299        Ok(())
300    }
301}
302
303impl Default for WalManager {
304    fn default() -> Self {
305        Self::new()
306    }
307}
308
309impl Drop for WalManager {
310    fn drop(&mut self) {
311        self.shutdown();
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318    use tempfile::TempDir;
319
320    #[test]
321    fn test_wal_manager_open_and_current_lsn() {
322        let temp_dir = TempDir::new().expect("Failed to create temp dir");
323        let mut manager = WalManager::new();
324
325        manager
326            .open(temp_dir.path(), 0)
327            .expect("Failed to open WAL");
328
329        assert_eq!(manager.current_lsn(), Lsn::ZERO);
330    }
331
332    #[test]
333    fn test_wal_manager_open_starts_async_flush_by_default() {
334        let temp_dir = TempDir::new().expect("Failed to create temp dir");
335        let mut manager = WalManager::new();
336        manager
337            .open(temp_dir.path(), 0)
338            .expect("Failed to open WAL");
339        assert!(
340            manager.is_async_enabled(),
341            "default WalConfig enables async flush, so open must start it"
342        );
343        manager.shutdown();
344    }
345
346    #[test]
347    fn test_wal_manager_open_sync_mode_stays_synchronous() {
348        let temp_dir = TempDir::new().expect("Failed to create temp dir");
349        let cfg = WalConfig {
350            enable_async_flush: false,
351            ..Default::default()
352        };
353        let mut manager = WalManager::with_config(cfg);
354        manager
355            .open(temp_dir.path(), 0)
356            .expect("Failed to open WAL");
357        assert!(
358            !manager.is_async_enabled(),
359            "async disabled in config must preserve exact sync behavior"
360        );
361    }
362
363    #[test]
364    fn test_wal_metrics_track_syncs() {
365        let temp_dir = TempDir::new().expect("Failed to create temp dir");
366        let mut manager = WalManager::new();
367        manager
368            .open(temp_dir.path(), 0)
369            .expect("Failed to open WAL");
370
371        manager.sync().expect("WAL sync should succeed");
372        let metrics = manager.metrics();
373        assert_eq!(metrics.sync_count, 1);
374        assert_eq!(metrics.sync_failures, 0);
375        assert_eq!(metrics.accepted_lsn, manager.current_lsn());
376        assert_eq!(metrics.durable_lsn, manager.durable_lsn());
377    }
378
379    #[test]
380    fn truncation_barrier_uses_the_oldest_published_index() {
381        use parking_lot::RwLock;
382        let mut manager = WalManager::new();
383        let registry = Arc::new(RwLock::new(std::collections::HashMap::from([
384            ((1, 1), CommitLsn::new(80)),
385            ((1, 2), CommitLsn::new(120)),
386        ])));
387        manager.set_index_barrier_registry(registry);
388
389        assert_eq!(
390            manager.truncation_barrier_lsn(),
391            Some(Lsn::new(80)),
392            "WAL truncation must stop at the oldest active index barrier"
393        );
394    }
395
396    #[test]
397    fn test_async_append_records_coordinator_for_piggyback_sync() {
398        let temp_dir = TempDir::new().expect("Failed to create temp dir");
399        let mut manager = WalManager::new();
400        manager
401            .open(temp_dir.path(), 0)
402            .expect("Failed to open WAL");
403
404        // Async append: recorded in the coordinator but must not block on fsync.
405        let lsn = manager
406            .append_transaction_with_durability(
407                linkrs_core::types::TransactionId::new(1),
408                Vec::new(),
409                &[],
410                linkrs_core::types::DurabilityLevel::Async,
411            )
412            .expect("async append must not wait");
413        assert!(lsn.get() > 0, "async append must still advance the WAL LSN");
414
415        // A later explicit sync must make the async append durable: the
416        // coordinator now knows about the appended LSN, so the sync's fsync
417        // covers it (this is the free-durability piggyback path).
418        manager.sync().expect("sync should succeed");
419        assert!(
420            manager.durable_lsn().as_u64() >= lsn.get(),
421            "sync after async append must cover the appended LSN"
422        );
423    }
424
425    #[test]
426    fn test_none_durability_does_not_block() {
427        let temp_dir = TempDir::new().expect("Failed to create temp dir");
428        let mut manager = WalManager::new();
429        manager
430            .open(temp_dir.path(), 0)
431            .expect("Failed to open WAL");
432
433        let lsn = manager
434            .append_transaction_with_durability(
435                linkrs_core::types::TransactionId::new(2),
436                Vec::new(),
437                &[],
438                linkrs_core::types::DurabilityLevel::None,
439            )
440            .expect("none durability append must not wait");
441        assert!(lsn.get() > 0);
442    }
443}
linkrs_storage/engine/wal_manager.rs

linkrs_storage/engine/
wal_manager.rs
