use std::fmt;
use std::num::NonZeroUsize;
use std::path::PathBuf;
use crate::tdata::TData;
use crate::{Error, Result};
pub const MAX_DEFAULT_THREADS: NonZeroUsize = NonZeroUsize::new(16).unwrap();
#[derive(Clone, PartialEq, Eq)]
pub struct LoadRequest {
pub dir: PathBuf,
pub passcode: Vec<u8>,
pub key_file: Option<String>,
}
impl LoadRequest {
#[must_use]
pub fn new(
dir: impl Into<PathBuf>,
passcode: impl Into<Vec<u8>>,
key_file: Option<String>,
) -> Self {
Self {
dir: dir.into(),
passcode: passcode.into(),
key_file,
}
}
}
impl fmt::Debug for LoadRequest {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("LoadRequest")
.field("dir", &self.dir)
.field("passcode_len", &self.passcode.len())
.field("key_file", &self.key_file)
.finish()
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct BatchOptions {
pub threads: Option<NonZeroUsize>,
}
impl BatchOptions {
pub fn from_threads<T: TryInto<usize>>(threads: Option<T>) -> Result<Self> {
let threads = threads
.map(|count| {
count
.try_into()
.ok()
.and_then(NonZeroUsize::new)
.ok_or_else(|| Error::InvalidInput {
what: "threads",
reason: "threads must be at least 1".to_owned(),
})
})
.transpose()?;
Ok(Self { threads })
}
#[must_use]
pub fn thread_count(&self) -> NonZeroUsize {
self.threads.unwrap_or_else(|| {
std::thread::available_parallelism()
.unwrap_or(NonZeroUsize::MIN)
.min(MAX_DEFAULT_THREADS)
})
}
}
fn load_one(request: &LoadRequest) -> Result<TData> {
TData::load(&request.dir, &request.passcode, request.key_file.as_deref())
}
#[must_use]
pub fn load_many(requests: &[LoadRequest], options: &BatchOptions) -> Vec<Result<TData>> {
#[cfg(feature = "parallel")]
{
if requests.is_empty() {
return Vec::new();
}
match pool::build(options.thread_count()) {
Ok(pool) => pool::map(&pool, requests, load_one),
Err(error) => requests
.iter()
.map(|_| Err(pool::to_error(&error)))
.collect(),
}
}
#[cfg(not(feature = "parallel"))]
{
let _ = options;
requests.iter().map(load_one).collect()
}
}
pub fn map_many<T, R>(
items: &[T],
options: &BatchOptions,
f: impl Fn(&T) -> R + Sync,
) -> Result<Vec<R>>
where
T: Sync,
R: Send,
{
#[cfg(feature = "parallel")]
{
if items.is_empty() {
return Ok(Vec::new());
}
let pool = pool::build(options.thread_count()).map_err(|error| pool::to_error(&error))?;
Ok(pool::map(&pool, items, f))
}
#[cfg(not(feature = "parallel"))]
{
let _ = options;
Ok(items.iter().map(f).collect())
}
}
#[cfg(feature = "parallel")]
mod pool {
use std::num::NonZeroUsize;
use rayon::prelude::*;
use rayon::{ThreadPool, ThreadPoolBuildError, ThreadPoolBuilder};
use crate::Error;
pub(super) fn build(threads: NonZeroUsize) -> Result<ThreadPool, ThreadPoolBuildError> {
ThreadPoolBuilder::new()
.num_threads(threads.get())
.thread_name(|index| format!("opentele2-batch-{index}"))
.build()
}
pub(super) fn to_error(error: &ThreadPoolBuildError) -> Error {
Error::Io(std::io::Error::other(error.to_string()))
}
pub(super) fn map<T, R>(pool: &ThreadPool, items: &[T], f: impl Fn(&T) -> R + Sync) -> Vec<R>
where
T: Sync,
R: Send,
{
pool.install(|| items.par_iter().map(&f).collect())
}
}