use std::{
io,
num::NonZeroUsize,
sync::{Arc, Mutex},
};
use super::{
bundle::DropReservation,
error::{DropAdmissionError, DropCapacityError, DropStartError},
executor::{CORE_WORKER_COUNT, DropExecutor, WorkerSpawner, system_spawner},
};
struct DropDomainInner {
worker_count: usize,
capacity: Option<NonZeroUsize>,
spawner: Arc<WorkerSpawner>,
executor: Mutex<Option<Arc<DropExecutor>>>,
}
#[derive(Clone)]
pub(crate) struct DropDomain(
Arc<DropDomainInner>,
);
impl DropDomain {
pub(crate) fn unstarted(capacity: Option<NonZeroUsize>) -> Self {
Self::unstarted_with_limit(CORE_WORKER_COUNT, capacity, system_spawner())
.expect("the production drop domain configuration must be valid")
}
#[cfg(test)]
pub(crate) fn try_start(capacity: usize) -> Result<Self, DropStartError> {
Self::try_start_with(CORE_WORKER_COUNT, capacity, system_spawner())
}
#[cfg(test)]
pub(super) fn try_start_with(
worker_count: usize,
capacity: usize,
spawner: Arc<WorkerSpawner>,
) -> Result<Self, DropStartError> {
let domain = Self::unstarted_with(worker_count, capacity, spawner)?;
domain.executor()?;
Ok(domain)
}
#[cfg(test)]
pub(super) fn unstarted_with(
worker_count: usize,
capacity: usize,
spawner: Arc<WorkerSpawner>,
) -> Result<Self, DropStartError> {
let capacity = NonZeroUsize::new(capacity).ok_or_else(|| {
DropStartError::new(
0,
io::Error::new(
io::ErrorKind::InvalidInput,
"worker count and capacity must be positive",
),
)
})?;
Self::unstarted_with_limit(worker_count, Some(capacity), spawner)
}
fn unstarted_with_limit(
worker_count: usize,
capacity: Option<NonZeroUsize>,
spawner: Arc<WorkerSpawner>,
) -> Result<Self, DropStartError> {
if worker_count == 0 {
return Err(DropStartError::new(
0,
io::Error::new(io::ErrorKind::InvalidInput, "worker count must be positive"),
));
}
Ok(Self(Arc::new(DropDomainInner {
worker_count,
capacity,
spawner,
executor: Mutex::new(None),
})))
}
fn executor(&self) -> Result<Arc<DropExecutor>, DropStartError> {
let mut executor = self
.0
.executor
.lock()
.unwrap_or_else(|error| error.into_inner());
if let Some(executor) = executor.as_ref() {
return Ok(Arc::clone(executor));
}
let started = DropExecutor::try_start_with(
self.0.worker_count,
self.0.capacity,
Arc::clone(&self.0.spawner),
)?;
*executor = Some(Arc::clone(&started));
Ok(started)
}
pub(crate) fn capacity(&self) -> Option<NonZeroUsize> {
self.0.capacity
}
pub(crate) async fn reserve(&self) -> Result<DropReservation, DropAdmissionError> {
self.executor()
.map_err(DropAdmissionError::Start)?
.reserve()
.await
.map_err(DropAdmissionError::Capacity)
}
pub(crate) fn try_reserve(&self) -> Result<DropReservation, DropAdmissionError> {
self.executor()
.map_err(DropAdmissionError::Start)?
.try_reserve()
.map_err(DropAdmissionError::Capacity)
}
pub(crate) fn try_reserve_many(
&self,
count: usize,
) -> Result<Vec<DropReservation>, DropAdmissionError> {
if count == 0 {
return Ok(Vec::new());
}
if self
.0
.capacity
.is_some_and(|capacity| count > capacity.get())
{
return Err(DropAdmissionError::Capacity(DropCapacityError::new(
self.0.capacity,
)));
}
self.executor()
.map_err(DropAdmissionError::Start)?
.try_reserve_many(count)
.map_err(DropAdmissionError::Capacity)
}
#[cfg(test)]
pub(crate) fn is_started(&self) -> bool {
self.0
.executor
.lock()
.unwrap_or_else(|error| error.into_inner())
.is_some()
}
#[cfg(test)]
pub(super) fn started_executor(&self) -> Arc<DropExecutor> {
self.0
.executor
.lock()
.unwrap_or_else(|error| error.into_inner())
.as_ref()
.map(Arc::clone)
.expect("the test domain must already be started")
}
}