353 lines
8.2 KiB
Rust
353 lines
8.2 KiB
Rust
use chrono::{DateTime, Utc};
|
|
use pile_config::{ConfigToml, Label, Source};
|
|
use pile_toolbox::cancelabletask::{CancelFlag, CancelableTaskError};
|
|
use rayon::{
|
|
ThreadPoolBuilder,
|
|
iter::{IntoParallelIterator, ParallelIterator},
|
|
};
|
|
use std::{
|
|
io::ErrorKind,
|
|
path::PathBuf,
|
|
sync::{
|
|
Arc,
|
|
atomic::{AtomicU64, Ordering},
|
|
mpsc::Receiver,
|
|
},
|
|
thread::JoinHandle,
|
|
time::Instant,
|
|
};
|
|
use tantivy::{Executor, Index, IndexWriter, TantivyError, collector::TopDocs};
|
|
use thiserror::Error;
|
|
use tracing::{debug, info, trace, warn};
|
|
|
|
use crate::{
|
|
DataSource, Item,
|
|
index::{DbFtsIndex, FtsLookupResult},
|
|
path_ts_earliest,
|
|
source::DirDataSource,
|
|
};
|
|
|
|
#[derive(Debug, Error)]
|
|
pub enum DatasetError {
|
|
#[error("{0}")]
|
|
IoError(#[from] std::io::Error),
|
|
|
|
#[error("{0}")]
|
|
TantivyError(#[from] TantivyError),
|
|
|
|
#[error("this dataset does not have an fts index")]
|
|
NoFtsIndex,
|
|
}
|
|
|
|
pub struct Dataset {
|
|
pub path_config: PathBuf,
|
|
pub path_parent: PathBuf,
|
|
pub path_workdir: PathBuf,
|
|
|
|
pub config: ConfigToml,
|
|
}
|
|
|
|
impl Dataset {
|
|
pub fn open(config: impl Into<PathBuf>) -> Result<Self, std::io::Error> {
|
|
let path_config = config.into();
|
|
let path_parent = path_config
|
|
.parent()
|
|
.ok_or(std::io::Error::new(
|
|
ErrorKind::NotADirectory,
|
|
format!("Config file {} has no parent", path_config.display()),
|
|
))?
|
|
.to_owned();
|
|
|
|
let config = {
|
|
let config = std::fs::read_to_string(&path_config)?;
|
|
let config: Result<ConfigToml, _> = toml::from_str(&config);
|
|
|
|
match config {
|
|
Ok(config) => {
|
|
trace!(message = "Loaded config", ?config);
|
|
config
|
|
}
|
|
|
|
Err(error) => {
|
|
return Err(std::io::Error::new(
|
|
ErrorKind::InvalidData,
|
|
format!("{} is invalid:\n{error}", path_config.display()),
|
|
));
|
|
}
|
|
}
|
|
};
|
|
|
|
let path_workdir = config
|
|
.dataset
|
|
.working_dir
|
|
.clone()
|
|
.unwrap_or(path_parent.join(".pile"))
|
|
.join(config.dataset.name.as_str());
|
|
|
|
return Ok(Self {
|
|
path_config,
|
|
path_parent,
|
|
path_workdir,
|
|
config,
|
|
});
|
|
}
|
|
|
|
//
|
|
// MARK: get
|
|
//
|
|
|
|
pub fn get(
|
|
&self,
|
|
source: &Label,
|
|
key: &PathBuf,
|
|
) -> Option<Box<dyn Item<Key = PathBuf> + 'static>> {
|
|
let s = self.config.dataset.source.get(source)?;
|
|
let s = match s {
|
|
Source::Flac { path } => DirDataSource::new(source, path.clone().to_vec()),
|
|
};
|
|
|
|
s.get(key).ok().flatten()
|
|
}
|
|
|
|
//
|
|
// MARK: fts
|
|
//
|
|
|
|
/// Refresh this dataset's fts index
|
|
pub fn fts_refresh(
|
|
&self,
|
|
threads: usize,
|
|
flag: Option<CancelFlag>,
|
|
) -> Result<(), CancelableTaskError<DatasetError>> {
|
|
let fts_tmp_dir = self.path_workdir.join(".tmp-fts");
|
|
let fts_dir = self.path_workdir.join("fts");
|
|
|
|
if fts_tmp_dir.is_dir() {
|
|
warn!("Removing temporary index in {}", fts_dir.display());
|
|
std::fs::remove_dir_all(&fts_tmp_dir).map_err(DatasetError::from)?;
|
|
}
|
|
|
|
std::fs::create_dir_all(&fts_tmp_dir).map_err(DatasetError::from)?;
|
|
|
|
let db_index = DbFtsIndex::new(&fts_tmp_dir, &self.config);
|
|
let mut index = Index::create_in_dir(&fts_tmp_dir, db_index.schema.clone())
|
|
.map_err(DatasetError::from)?;
|
|
index.set_executor(Executor::multi_thread(10, "build-fts").map_err(DatasetError::from)?);
|
|
let mut index_writer: IndexWriter =
|
|
index.writer(50 * 1024 * 1024).map_err(DatasetError::from)?;
|
|
|
|
let batch_size = 1000;
|
|
let (_read_task, read_rx) = start_read_task(&self.config, batch_size);
|
|
|
|
#[expect(clippy::unwrap_used)]
|
|
let write_pool = ThreadPoolBuilder::new()
|
|
.num_threads(threads.max(1))
|
|
.thread_name(|x| format!("fts_refresh_thread_{x}"))
|
|
.build()
|
|
.unwrap();
|
|
|
|
let mut total = 0u64;
|
|
while let Ok(batch) = read_rx.recv() {
|
|
let batch = batch?;
|
|
if let Some(flag) = &flag
|
|
&& flag.is_cancelled()
|
|
{
|
|
return Err(CancelableTaskError::Cancelled);
|
|
}
|
|
|
|
let this = AtomicU64::new(0);
|
|
let start = Instant::now();
|
|
write_pool
|
|
.install(|| {
|
|
batch
|
|
.into_par_iter()
|
|
.filter_map(|(key, item)| match db_index.entry_to_document(&*item) {
|
|
Ok(Some(doc)) => Some((key, doc)),
|
|
Ok(None) => {
|
|
warn!("Skipping {key:?}, document is empty");
|
|
None
|
|
}
|
|
Err(err) => {
|
|
warn!("Could not read {key:?}, skipping. {err}");
|
|
None
|
|
}
|
|
})
|
|
.map(|(key, doc)| {
|
|
this.fetch_add(1, Ordering::Relaxed);
|
|
index_writer
|
|
.add_document(doc)
|
|
.map_err(|err| (key, err))
|
|
.map(|_| ())
|
|
})
|
|
.find_first(|x| x.is_err())
|
|
.unwrap_or(Ok(()))
|
|
})
|
|
.map_err(|(_key, err)| DatasetError::from(err))?;
|
|
|
|
let this = this.load(Ordering::Relaxed);
|
|
total += this;
|
|
let time_ms = start.elapsed().as_millis();
|
|
debug!("Added a batch of {this} in {time_ms} ms ({total} total)");
|
|
}
|
|
|
|
if let Some(flag) = flag.as_ref()
|
|
&& flag.is_cancelled()
|
|
{
|
|
return Err(CancelableTaskError::Cancelled);
|
|
}
|
|
|
|
info!("Committing index");
|
|
index_writer.commit().map_err(DatasetError::from)?;
|
|
|
|
if fts_dir.is_dir() {
|
|
warn!("Removing existing index in {}", fts_dir.display());
|
|
std::fs::remove_dir_all(&fts_dir).map_err(DatasetError::from)?;
|
|
}
|
|
std::fs::rename(&fts_tmp_dir, &fts_dir).map_err(DatasetError::from)?;
|
|
|
|
return Ok(());
|
|
}
|
|
|
|
pub fn fts_lookup(
|
|
&self,
|
|
query: &str,
|
|
top_n: usize,
|
|
) -> Result<Vec<FtsLookupResult>, DatasetError> {
|
|
let fts_dir = self.path_workdir.join("fts");
|
|
|
|
if !fts_dir.exists() {
|
|
return Err(DatasetError::NoFtsIndex);
|
|
}
|
|
if !fts_dir.is_dir() {
|
|
return Err(std::io::Error::new(
|
|
ErrorKind::NotADirectory,
|
|
format!("fts index {} is not a directory", fts_dir.display()),
|
|
)
|
|
.into());
|
|
}
|
|
|
|
let db_index = DbFtsIndex::new(&fts_dir, &self.config);
|
|
let results = db_index.lookup(query, Arc::new(TopDocs::with_limit(top_n)))?;
|
|
return Ok(results);
|
|
}
|
|
|
|
/// Time at which fts was created
|
|
pub fn ts_fts(&self) -> Result<Option<DateTime<Utc>>, std::io::Error> {
|
|
let fts_dir = self.path_workdir.join("fts");
|
|
|
|
if !fts_dir.exists() {
|
|
return Ok(None);
|
|
}
|
|
if !fts_dir.is_dir() {
|
|
return Err(std::io::Error::new(
|
|
ErrorKind::NotADirectory,
|
|
format!("fts index {} is not a directory", fts_dir.display()),
|
|
));
|
|
}
|
|
|
|
return path_ts_earliest(&fts_dir);
|
|
}
|
|
|
|
/// Time at which data was last modified
|
|
pub fn ts_data(&self) -> Result<Option<DateTime<Utc>>, std::io::Error> {
|
|
let mut ts: Option<DateTime<Utc>> = None;
|
|
|
|
for (label, source) in &self.config.dataset.source {
|
|
match source {
|
|
Source::Flac { path } => {
|
|
let s = DirDataSource::new(label, path.clone().to_vec());
|
|
match (ts, s.latest_change()?) {
|
|
(_, None) => continue,
|
|
(None, Some(new)) => ts = Some(new),
|
|
(Some(old), Some(new)) => ts = Some(old.max(new)),
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
return Ok(ts);
|
|
}
|
|
|
|
/// Returns true if we do not have an fts index,
|
|
/// or if our fts index is older than our data.
|
|
pub fn needs_fts(&self) -> Result<bool, std::io::Error> {
|
|
let start = Instant::now();
|
|
let ts_fts = self.ts_fts()?;
|
|
let ts_data = self.ts_data()?;
|
|
|
|
let result = match (ts_fts, ts_data) {
|
|
(None, Some(_)) => true,
|
|
(None, None) | (Some(_), None) => {
|
|
warn!("Could not determine data age");
|
|
false
|
|
}
|
|
|
|
(Some(ts_fts), Some(ts_data)) => ts_data > ts_fts,
|
|
};
|
|
|
|
debug!(
|
|
message = "Ran needs_fts",
|
|
time_ms = start.elapsed().as_millis(),
|
|
?result
|
|
);
|
|
|
|
return Ok(result);
|
|
}
|
|
}
|
|
|
|
//
|
|
// MARK: read_task
|
|
//
|
|
|
|
fn start_read_task(
|
|
config: &ConfigToml,
|
|
batch_size: usize,
|
|
) -> (
|
|
JoinHandle<()>,
|
|
Receiver<Result<Vec<(PathBuf, Box<dyn Item<Key = PathBuf>>)>, DatasetError>>,
|
|
) {
|
|
let config = config.clone();
|
|
let (read_tx, read_rx) = std::sync::mpsc::sync_channel(2);
|
|
|
|
let read_task = std::thread::spawn(move || {
|
|
let mut batch = Vec::with_capacity(batch_size);
|
|
for (name, source) in &config.dataset.source {
|
|
info!("Loading source {name}");
|
|
|
|
match source {
|
|
Source::Flac { path: dir } => {
|
|
let source = DirDataSource::new(name, dir.clone().to_vec());
|
|
for i in source.iter() {
|
|
match i {
|
|
Ok(x) => batch.push(x),
|
|
Err(err) => {
|
|
let err = Err(DatasetError::from(err));
|
|
let _ = read_tx.send(err);
|
|
return;
|
|
}
|
|
}
|
|
|
|
if batch.len() >= batch_size {
|
|
let b = std::mem::replace(&mut batch, Vec::with_capacity(batch_size));
|
|
|
|
match read_tx.send(Ok(b)) {
|
|
Ok(()) => {}
|
|
Err(_) => return,
|
|
};
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if !batch.is_empty() {
|
|
match read_tx.send(Ok(batch)) {
|
|
Ok(()) => {}
|
|
Err(_) => return,
|
|
};
|
|
}
|
|
});
|
|
|
|
return (read_task, read_rx);
|
|
}
|