use crate::{
SequenceNumberCounter, TableId,
fs::{Fs, FsFile, FsOpenOptions, SyncMode},
path::{Path, PathBuf},
};
#[cfg(not(feature = "std"))]
use alloc::string::ToString;
use alloc::{boxed::Box, collections::VecDeque, sync::Arc, vec::Vec};
#[cfg(feature = "std")]
use parking_lot::Mutex;
#[cfg(not(feature = "std"))]
use spin::Mutex;
pub const DEFAULT_TABLE_FILE_POOL: usize = 16;
pub struct ClaimedFile {
pub id: TableId,
pub path: PathBuf,
pub file: Box<dyn FsFile>,
pub durable_entry: bool,
pub writing: Writing,
}
pub const WRITING_BELOW_CAP: usize = 16;
type Unwritten = Arc<Mutex<Vec<TableId>>>;
pub struct Writing {
folder: Unwritten,
id: TableId,
}
impl Writing {
fn new(folder: &Unwritten, id: TableId) -> Self {
folder.lock().push(id);
Self::listed(folder, id)
}
fn listed(folder: &Unwritten, id: TableId) -> Self {
Self {
folder: Arc::clone(folder),
id,
}
}
#[cfg(test)]
pub fn alone(id: TableId) -> Self {
Self::new(&Arc::new(Mutex::new(Vec::new())), id)
}
pub fn below(&self) -> Vec<TableId> {
let mut ids: Vec<TableId> = self
.folder
.lock()
.iter()
.copied()
.filter(|id| *id < self.id)
.collect();
ids.sort_unstable();
ids.truncate(WRITING_BELOW_CAP);
ids
}
}
impl Drop for Writing {
fn drop(&mut self) {
unlist(&self.folder, self.id);
}
}
fn unlist(folder: &Unwritten, id: TableId) {
let mut ids = folder.lock();
if let Some(at) = ids.iter().position(|listed| *listed == id) {
ids.swap_remove(at);
}
}
struct FolderPool {
routes: Vec<(Arc<dyn Fs>, PathBuf)>,
ids: VecDeque<TableId>,
synced: bool,
unpooled: bool,
writing: Unwritten,
refilling: bool,
}
struct Pools {
ids: SequenceNumberCounter,
batch: usize,
sync_mode: SyncMode,
background: bool,
folders: Mutex<Vec<FolderPool>>,
#[cfg(feature = "std")]
refilled: parking_lot::Condvar,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum RefillOn {
#[default]
Background,
ClaimingThread,
}
struct Refill {
at: usize,
folder: PathBuf,
fs: Arc<dyn Fs>,
ids: Vec<TableId>,
}
pub struct TableFilePools {
pools: Arc<Pools>,
#[cfg(feature = "std")]
refills: Mutex<Vec<std::thread::JoinHandle<()>>>,
}
impl TableFilePools {
#[cfg(test)]
pub fn new(ids: SequenceNumberCounter, batch: usize, sync_mode: SyncMode) -> Self {
Self::with(ids, batch, sync_mode, false)
}
#[cfg(all(test, feature = "std"))]
pub fn refilling_in_background(
ids: SequenceNumberCounter,
batch: usize,
sync_mode: SyncMode,
) -> Self {
Self::with(ids, batch, sync_mode, true)
}
pub fn for_tree(
ids: SequenceNumberCounter,
batch: usize,
sync_mode: SyncMode,
refill: RefillOn,
) -> Self {
let background = cfg!(feature = "std") && refill == RefillOn::Background;
Self::with(ids, batch, sync_mode, background)
}
fn with(
ids: SequenceNumberCounter,
batch: usize,
sync_mode: SyncMode,
background: bool,
) -> Self {
Self {
pools: Arc::new(Pools {
ids,
batch,
sync_mode,
background,
folders: Mutex::new(Vec::new()),
#[cfg(feature = "std")]
refilled: parking_lot::Condvar::new(),
}),
#[cfg(feature = "std")]
refills: Mutex::new(Vec::new()),
}
}
pub fn adopt(&self, folder: &Path, fs: &Arc<dyn Fs>, mut ids: Vec<TableId>) {
if self.pools.batch == 0 {
return;
}
let mut folders = self.pools.folders.lock();
ids.retain(|id| !folders.iter().any(|pool| pool.ids.contains(id)));
if ids.is_empty() {
return;
}
ids.sort_unstable();
let at = pool_of(&mut folders, folder, fs);
let Some(pool) = folders.get_mut(at) else {
return;
};
if pool.unpooled {
for id in ids {
discard(&**fs, &folder.join(id.to_string()));
}
return;
}
pool.writing.lock().extend(ids.iter().copied());
pool.ids.extend(ids);
pool.synced = false;
drop(folders);
}
pub fn claim(
&self,
folder: &Path,
fs: &Arc<dyn Fs>,
above: Option<TableId>,
may_refill: bool,
) -> crate::Result<ClaimedFile> {
let mut folders = self.pools.folders.lock();
let at = pool_of(&mut folders, folder, fs);
#[cfg(feature = "std")]
while folders
.get(at)
.is_some_and(|pool| pool.refilling && pool.ids.is_empty())
{
self.pools.refilled.wait(&mut folders);
}
let Some(pool) = folders.get_mut(at) else {
return Err(crate::Error::Unrecoverable);
};
let (claimed, due) = if pool.unpooled || self.pools.batch == 0 {
(Ok(None), false)
} else {
self.pools.claim_pooled(pool, folder, fs, above, may_refill)
};
let pooled = claimed?.ok_or_else(|| Writing::new(&pool.writing, self.pools.ids.next()));
let refill = due.then(|| self.pools.reserve(pool, at, folder, fs));
drop(folders);
let claimed = match pooled {
Ok(claimed) => Ok(claimed),
Err(made) => Pools::create(folder, fs, made),
};
if let Some(refill) = refill {
self.start(refill);
}
claimed
}
#[cfg(feature = "std")]
fn start(&self, refill: Refill) {
let (at, ids) = (refill.at, refill.ids.clone());
let pools = Arc::clone(&self.pools);
match std::thread::Builder::new()
.name("lsm-table-files".into())
.spawn(move || {
let ran =
std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| pools.run(&refill)));
if let Err(panic) = ran {
pools.withdraw(refill.at, &refill.ids);
std::panic::resume_unwind(panic);
}
}) {
Ok(handle) => {
let mut refills = self.refills.lock();
refills.retain(|refill| !refill.is_finished());
refills.push(handle);
}
Err(e) => {
log::warn!("could not start refilling a table file pool: {e}");
self.pools.withdraw(at, &ids);
}
}
}
#[cfg(not(feature = "std"))]
fn start(&self, refill: Refill) {
self.pools.withdraw(refill.at, &refill.ids);
}
pub fn wait_for_refills(&self) {
#[cfg(feature = "std")]
{
let refills = core::mem::take(&mut *self.refills.lock());
for refill in refills {
if refill.join().is_err() {
log::warn!("a table file pool refill panicked");
}
}
}
}
}
impl Drop for TableFilePools {
fn drop(&mut self) {
self.wait_for_refills();
}
}
impl Pools {
fn claim_pooled(
&self,
pool: &mut FolderPool,
folder: &Path,
fs: &Arc<dyn Fs>,
above: Option<TableId>,
may_refill: bool,
) -> (crate::Result<Option<ClaimedFile>>, bool) {
let mut refilled = false;
let claimed = loop {
if !pool.synced && !pool.ids.is_empty() {
if let Err(e) = crate::file::fsync_directory(folder, &**fs, self.sync_mode) {
break Err(e.into());
}
pool.synced = true;
}
let Some(id) = pool.ids.pop_front() else {
if self.background || refilled || !may_refill {
break Ok(None);
}
if let Err(e) = self.refill(pool, folder, &**fs) {
break Err(e);
}
refilled = true;
continue;
};
let path = folder.join(id.to_string());
if above.is_some_and(|above| id <= above) {
unlist(&pool.writing, id);
discard(&**fs, &path);
continue;
}
match open_unshared_empty(&**fs, &path) {
Ok(Opened::Unshared(file)) => {
break Ok(Some(ClaimedFile {
id,
path,
file,
durable_entry: true,
writing: Writing::listed(&pool.writing, id),
}));
}
Ok(Opened::Unusable) => unlist(&pool.writing, id),
Ok(Opened::LinksUncountable) => {
unlist(&pool.writing, id);
for id in pool.ids.drain(..) {
unlist(&pool.writing, id);
discard(&**fs, &folder.join(id.to_string()));
}
pool.unpooled = true;
break Ok(None);
}
Err(e) => {
unlist(&pool.writing, id);
break Err(e);
}
}
};
let due = self.background
&& may_refill
&& !pool.unpooled
&& !pool.refilling
&& pool.ids.len() <= self.batch / 2;
(claimed, due)
}
fn reserve(&self, pool: &mut FolderPool, at: usize, folder: &Path, fs: &Arc<dyn Fs>) -> Refill {
pool.refilling = true;
let ids: Vec<TableId> = (0..self.batch).map(|_| self.ids.next()).collect();
pool.writing.lock().extend(ids.iter().copied());
Refill {
at,
folder: folder.to_path_buf(),
fs: Arc::clone(fs),
ids,
}
}
#[cfg(feature = "std")]
fn run(&self, refill: &Refill) {
let mut made = Vec::with_capacity(refill.ids.len());
let mut synced = true;
for &id in &refill.ids {
let path = refill.folder.join(id.to_string());
match refill
.fs
.open(&path, &FsOpenOptions::new().write(true).create_new(true))
{
Ok(file) => {
drop(file);
made.push(id);
}
Err(e) => {
log::warn!("could not make pooled table file {}: {e}", path.display());
synced = false;
break;
}
}
}
if !made.is_empty()
&& let Err(e) =
crate::file::fsync_directory(&refill.folder, &*refill.fs, self.sync_mode)
{
log::warn!(
"could not sync {} for its pooled table files: {e}",
refill.folder.display()
);
synced = false;
}
if let Some(pool) = self.folders.lock().get_mut(refill.at) {
Self::pool_made(pool, refill, made, synced);
}
self.refilled.notify_all();
}
#[cfg(feature = "std")]
fn pool_made(pool: &mut FolderPool, refill: &Refill, made: Vec<TableId>, synced: bool) {
pool.refilling = false;
for id in refill.ids.iter().filter(|id| !made.contains(id)) {
unlist(&pool.writing, *id);
}
if pool.unpooled {
for id in made {
unlist(&pool.writing, id);
discard(&*refill.fs, &refill.folder.join(id.to_string()));
}
return;
}
if !made.is_empty() {
pool.synced &= synced;
pool.ids.extend(made);
}
}
fn withdraw(&self, at: usize, ids: &[TableId]) {
if let Some(pool) = self.folders.lock().get_mut(at) {
pool.refilling = false;
for id in ids {
unlist(&pool.writing, *id);
}
}
#[cfg(feature = "std")]
self.refilled.notify_all();
}
fn create(folder: &Path, fs: &Arc<dyn Fs>, writing: Writing) -> crate::Result<ClaimedFile> {
let id = writing.id;
let path = folder.join(id.to_string());
let file = fs.open(&path, &FsOpenOptions::new().write(true).create_new(true))?;
Ok(ClaimedFile {
id,
path,
file,
durable_entry: false,
writing,
})
}
fn refill(&self, pool: &mut FolderPool, folder: &Path, fs: &dyn Fs) -> crate::Result<()> {
pool.synced = false;
for _ in 0..self.batch {
let id = self.ids.next();
let path = folder.join(id.to_string());
drop(fs.open(&path, &FsOpenOptions::new().write(true).create_new(true))?);
pool.writing.lock().push(id);
pool.ids.push_back(id);
}
crate::file::fsync_directory(folder, fs, self.sync_mode)?;
pool.synced = true;
Ok(())
}
}
fn pool_of(folders: &mut Vec<FolderPool>, folder: &Path, fs: &Arc<dyn Fs>) -> usize {
let known = folders.iter().position(|pool| {
pool.routes
.iter()
.any(|(route_fs, route)| Arc::ptr_eq(route_fs, fs) && route == folder)
});
known.unwrap_or_else(|| {
let alias = folders.iter().position(|pool| {
pool.routes
.first()
.is_some_and(|(route_fs, route)| same_folder(route_fs, route, fs, folder))
});
let route = (Arc::clone(fs), folder.to_path_buf());
if let Some((at, pool)) = alias.and_then(|at| Some((at, folders.get_mut(at)?))) {
pool.routes.push(route);
at
} else {
folders.push(FolderPool {
routes: alloc::vec![route],
ids: VecDeque::new(),
synced: true,
unpooled: fs.capabilities(folder).directory_sync_noop,
writing: Arc::new(Mutex::new(Vec::new())),
refilling: false,
});
folders.len() - 1
}
})
}
fn same_folder(route_fs: &Arc<dyn Fs>, route: &Path, fs: &Arc<dyn Fs>, folder: &Path) -> bool {
let one_namespace = match (route_fs.backend_id(), fs.backend_id()) {
(Some(a), Some(b)) => a == b,
_ => Arc::ptr_eq(route_fs, fs),
};
one_namespace && fs.same_file(route, folder).unwrap_or(false)
}
enum Opened {
Unshared(Box<dyn FsFile>),
Unusable,
LinksUncountable,
}
fn open_unshared_empty(fs: &dyn Fs, path: &Path) -> crate::Result<Opened> {
let opts = FsOpenOptions::new()
.write(true)
.no_follow(true)
.no_wait(true);
let file = match fs.open(path, &opts) {
Ok(file) => file,
Err(e) if e.kind() == crate::io::ErrorKind::NotFound => return Ok(Opened::Unusable),
Err(e) => {
if let Ok(Some(_)) = fs.read_link(path) {
return Ok(Opened::Unusable);
}
return match fs.metadata(path) {
Ok(meta) if !meta.is_file => Ok(Opened::Unusable),
_ => Err(e.into()),
};
}
};
let meta = file.metadata()?;
if !meta.is_file || meta.len != 0 {
return Ok(Opened::Unusable);
}
match file.hard_link_count() {
Ok(1) => Ok(Opened::Unshared(file)),
Err(e) if e.kind() == crate::io::ErrorKind::Unsupported => {
drop(file);
discard(fs, path);
Ok(Opened::LinksUncountable)
}
Ok(_) | Err(_) => {
drop(file);
discard(fs, path);
Ok(Opened::Unusable)
}
}
}
fn discard(fs: &dyn Fs, path: &Path) {
if let Err(e) = fs.remove_file(path)
&& e.kind() != crate::io::ErrorKind::NotFound
{
log::warn!("could not remove pooled table file {}: {e}", path.display());
}
}
#[cfg(test)]
mod tests;