use std::fs::File;
use std::io::{Seek, SeekFrom, Write};
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex};
use linkrs_core::wal::types::{WalError, WalResult};
use super::super::buffer;
use super::LocalWalWriter;
pub(crate) struct BufferedFlushState {
pub(crate) file: File,
pub(crate) offset: u64,
pub(crate) file_size: usize,
}
impl LocalWalWriter {
pub fn enable_async_buffer(&mut self, buffer: Arc<buffer::WalBuffer>) {
if self.flush_coordinator.is_none() {
let coordinator = buffer::WalFlushCoordinator::new(self.config.flush_interval());
coordinator.register_buffer(Arc::clone(&buffer));
self.flush_coordinator = Some(coordinator);
} else if let Some(ref coordinator) = self.flush_coordinator {
coordinator.register_buffer(Arc::clone(&buffer));
}
self.buffer = Some(buffer);
self.buffered_lsn
.store(self.current_lsn.load(Ordering::SeqCst), Ordering::SeqCst);
self.ensure_flush_state();
}
pub fn enable_async_flush(&mut self) {
if !self.config.enable_async_flush {
return;
}
let buf = Arc::new(buffer::WalBuffer::new(self.config.buffer_size));
self.enable_async_buffer(buf);
}
pub fn disable_async_flush(&mut self) -> WalResult<()> {
self.flush_buffered_to_file()?;
self.buffer = None;
if let Some(ref coordinator) = self.flush_coordinator {
coordinator.shutdown();
}
self.flush_coordinator = None;
self.flush_state = None;
self.background_started.store(false, Ordering::SeqCst);
Ok(())
}
pub fn is_async_enabled(&self) -> bool {
self.buffer.is_some() && self.config.enable_async_flush
}
pub fn request_flush(&self) {
if let Some(ref coordinator) = self.flush_coordinator {
coordinator.signal_flush();
}
}
pub(crate) fn ensure_flush_state(&mut self) {
if self.flush_state.is_some() || self.buffer.is_none() {
return;
}
if let Some(ref file) = self.file {
if let Ok(cloned) = file.try_clone() {
self.flush_state = Some(Arc::new(Mutex::new(BufferedFlushState {
file: cloned,
offset: self.file_used as u64,
file_size: self.file_size,
})));
}
}
}
fn write_to_flush_state(state: &mut BufferedFlushState, data: &[u8]) -> WalResult<()> {
let end = state.offset + data.len() as u64;
if end as usize > state.file_size {
state.file.set_len(end)?;
state.file_size = end as usize;
}
state.file.seek(SeekFrom::Start(state.offset))?;
state.file.write_all(data)?;
state.offset = end;
Ok(())
}
pub(crate) fn drain_buffer_to_state(&self) -> WalResult<bool> {
let buf = match self.buffer.as_ref() {
Some(b) => Arc::clone(b),
None => return Ok(false),
};
if buf.pending_bytes() == 0 {
return Ok(false);
}
let state = match self.flush_state.as_ref() {
Some(s) => Arc::clone(s),
None => return Ok(false),
};
let data = buf.drain();
if data.is_empty() {
return Ok(false);
}
let mut guard = state
.lock()
.map_err(|e| WalError::InvalidOperation(format!("flush state lock poisoned: {}", e)))?;
Self::write_to_flush_state(&mut guard, &data)?;
let staged = self.buffered_lsn.load(Ordering::SeqCst);
drop(guard);
if let Some(ref coordinator) = self.group_commit {
coordinator.record_appended(staged);
}
Ok(true)
}
fn syncback_file_used(&mut self) {
if let Some(ref state) = self.flush_state {
if let Ok(guard) = state.lock() {
self.file_used = guard.offset as usize;
self.file_size = guard.file_size;
}
}
}
pub fn flush_buffered_to_file(&mut self) -> WalResult<()> {
if self.buffer.is_none() {
return Ok(());
}
self.ensure_flush_state();
if self.flush_state.is_some() {
let _ = self.drain_buffer_to_state()?;
self.syncback_file_used();
self.request_flush();
return Ok(());
}
Ok(())
}
pub fn flush_and_sync(&mut self) -> WalResult<()> {
self.flush_buffered_to_file()?;
self.sync_via_file()
}
pub fn start_background_flush(&mut self) -> WalResult<()> {
if !self.is_async_enabled() {
return Err(WalError::InvalidOperation(
"async buffer not enabled".to_string(),
));
}
if self
.background_started
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
return Ok(());
}
self.ensure_flush_state();
let coordinator = self.flush_coordinator.clone().ok_or_else(|| {
WalError::InvalidOperation("flush coordinator not enabled".to_string())
})?;
let state = self
.flush_state
.clone()
.ok_or_else(|| WalError::InvalidOperation("flush state not initialized".to_string()))?;
let staged = Arc::clone(&self.buffered_lsn);
let group = self.group_commit.clone();
coordinator.start_flush_thread(move |data: &[u8]| {
let mut guard = state.lock().map_err(|e| {
WalError::InvalidOperation(format!("flush state lock poisoned: {}", e))
})?;
Self::write_to_flush_state(&mut guard, data)?;
let lsn = staged.load(Ordering::SeqCst);
drop(guard);
if let Some(ref coordinator) = group {
coordinator.record_appended(lsn);
}
Ok(())
});
Ok(())
}
pub fn buffer(&self) -> Option<&Arc<buffer::WalBuffer>> {
self.buffer.as_ref()
}
}