use std::io;
use std::os::windows::io::{BorrowedHandle, OwnedHandle};
#[cfg(test)]
use std::sync::atomic::AtomicBool;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, Weak};
use std::time::Duration;
use windows_impersonation_token_sys::ImpersonationToken;
use windows_threadpool_sys::callback_env::CallbackEnviron;
use windows_threadpool_sys::work::ThreadpoolWork;
use crate::admission::{self, EnumerationHandle};
use crate::completion::{Completion, EnumerationId, TerminalOutcome};
use crate::completion_ring::{CompletionRing, MINIMUM_COMPLETION_CAPACITY};
use crate::engine::{self, EngineState};
use crate::error::{BeginError, SessionError, SessionFailure};
use crate::registry::{EnumerationState, Registry};
use crate::request::EnumerationRequest;
use crate::submission_ring::{
AbandonSlot, ControlMessage, PushOutcome, SubmissionRing, release_retire_slot,
};
pub const MINIMUM_SUBMISSION_CAPACITY: usize = 4;
pub const MINIMUM_COMPLETION_RING_CAPACITY: usize = MINIMUM_COMPLETION_CAPACITY;
#[derive(Debug)]
pub(crate) enum QuantumOutcome {
#[allow(
dead_code,
reason = "only the cfg(test) state-machine model constructs this; a plain lib build never does"
)]
Idle,
Yielded,
Parked,
Finished(TerminalOutcome),
}
struct SessionWork {
servicer: ThreadpoolWork,
engine: ThreadpoolWork,
#[cfg(test)]
suppressed: AtomicBool,
}
impl SessionWork {
#[cfg(test)]
fn is_suppressed(&self) -> bool {
self.suppressed.load(Ordering::Acquire)
}
#[cfg(not(test))]
fn is_suppressed(&self) -> bool {
false
}
fn submit_servicer(&self) {
if !self.is_suppressed() {
self.servicer.submit();
}
}
fn submit_engine(&self) {
if !self.is_suppressed() {
self.engine.submit();
}
}
}
pub(crate) struct SessionShared {
pub(crate) completions: Arc<CompletionRing>,
pub(crate) submissions: SubmissionRing,
registry: Mutex<Registry>,
next_id: AtomicU64,
work: Mutex<Option<SessionWork>>,
handles: AtomicUsize,
#[cfg(test)]
scripted: Mutex<std::collections::VecDeque<QuantumOutcome>>,
}
impl SessionShared {
fn registry(&self) -> MutexGuard<'_, Registry> {
self.registry
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
fn work(&self) -> MutexGuard<'_, Option<SessionWork>> {
self.work
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
pub(crate) fn next_enumeration_id(&self) -> EnumerationId {
EnumerationId::from_raw(self.next_id.fetch_add(1, Ordering::Relaxed))
}
#[cfg(test)]
pub(crate) fn contains(&self, enumeration: EnumerationId) -> bool {
self.registry().contains(enumeration)
}
pub(crate) fn registered(&self) -> usize {
self.registry().len()
}
#[cfg(test)]
pub(crate) fn ready(&self) -> usize {
self.registry().ready_len()
}
pub(crate) fn ring_servicer(&self, outcome: PushOutcome) {
if outcome != PushOutcome::RingDoorbell {
return;
}
if let Some(work) = self.work().as_ref() {
work.submit_servicer();
}
}
fn acquire_handle(&self) {
self.handles.fetch_add(1, Ordering::AcqRel);
}
fn release_handle(&self) {
if self.handles.fetch_sub(1, Ordering::AcqRel) != 1 {
return;
}
let work = self.work().take();
drop(work);
}
pub(crate) fn drain_submissions(&self) {
while let Some(message) = self.submissions.take_for_service() {
match message {
ControlMessage::Begin(begin) => self.service_begin(*begin),
ControlMessage::Cancel(enumeration) => self.service_cancel(enumeration),
ControlMessage::Retire(enumeration) => self.service_retire(enumeration),
ControlMessage::Abandon => self.service_abandon(),
}
}
}
fn service_begin(&self, begin: crate::submission_ring::BeginMessage) {
let enumeration = begin.enumeration;
{
let mut registry = self.registry();
if !registry.is_accepting() {
drop(registry);
self.retire_state(
EnumerationState::new(begin.engine, begin.terminal, begin.retire),
None,
);
return;
}
registry.insert(
enumeration,
EnumerationState::new(begin.engine, begin.terminal, begin.retire),
);
}
self.schedule(enumeration);
}
fn service_cancel(&self, enumeration: EnumerationId) {
let finished = {
let mut registry = self.registry();
let Some(state) = registry.get_mut(enumeration) else {
return;
};
state.cancelled = true;
state.parked = false;
if state.is_quiescent() {
registry.remove(enumeration)
} else {
None
}
};
if let Some(state) = finished {
self.retire_state(state, Some(TerminalOutcome::Cancelled));
}
}
fn service_retire(&self, enumeration: EnumerationId) {
let state = self.registry().remove(enumeration);
if let Some(state) = state {
self.retire_state(state, None);
}
}
fn service_abandon(&self) {
let abandoned = {
let mut registry = self.registry();
registry.stop_accepting();
registry.drain_all()
};
for (_, state) in abandoned {
self.retire_state(state, None);
}
}
fn retire_state(&self, mut state: EnumerationState, outcome: Option<TerminalOutcome>) {
if let Some(retire) = state.retire.take() {
release_retire_slot(&self.submissions, retire);
}
match (state.terminal.take(), outcome) {
(Some(terminal), Some(outcome)) => terminal.send(outcome),
(slot, _) => drop(slot),
}
}
pub(crate) fn schedule(&self, enumeration: EnumerationId) {
self.registry().mark_ready(enumeration);
if let Some(work) = self.work().as_ref() {
work.submit_engine();
}
}
pub(crate) fn run_engine_quantum(&self) {
let Some((enumeration, mut engine)) = self.claim_next() else {
return;
};
let outcome = self.advance(enumeration, &mut engine);
self.report_quantum(enumeration, engine, outcome);
}
pub(crate) fn claim_next(&self) -> Option<(EnumerationId, EngineState)> {
self.registry().claim_next()
}
fn advance(&self, enumeration: EnumerationId, engine: &mut EngineState) -> QuantumOutcome {
#[cfg(test)]
if let Some(scripted) = self
.scripted
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.pop_front()
{
return scripted;
}
engine::advance(engine, enumeration, &self.completions)
}
pub(crate) fn report_quantum(
&self,
enumeration: EnumerationId,
engine: EngineState,
outcome: QuantumOutcome,
) {
let mut resume = false;
let finish = {
let mut registry = self.registry();
let Some(state) = registry.get_mut(enumeration) else {
return;
};
state.running = false;
state.engine = Some(engine);
match outcome {
QuantumOutcome::Finished(outcome) => Some(outcome),
_ if state.cancelled => Some(TerminalOutcome::Cancelled),
QuantumOutcome::Yielded => {
resume = true;
None
}
QuantumOutcome::Parked => {
if self.completions.has_data_room() {
resume = true;
} else {
state.parked = true;
}
None
}
QuantumOutcome::Idle => None,
}
};
if let Some(outcome) = finish {
self.finish_from_worker(enumeration, outcome);
} else if resume {
self.schedule(enumeration);
}
}
fn finish_from_worker(&self, enumeration: EnumerationId, outcome: TerminalOutcome) {
let (terminal, retire) = {
let mut registry = self.registry();
match registry.get_mut(enumeration) {
Some(state) => (state.terminal.take(), state.retire.take()),
None => return,
}
};
if let Some(terminal) = terminal {
terminal.send(outcome);
}
if let Some(retire) = retire {
let pushed = self.submissions.push_retire(retire, enumeration);
self.ring_servicer(pushed);
}
}
pub(crate) fn resume_parked(&self) {
let parked = {
let registry = self.registry();
registry.parked()
};
for enumeration in parked {
self.schedule(enumeration);
}
}
#[cfg(test)]
pub(crate) fn script_quantum(&self, outcome: QuantumOutcome) {
self.scripted
.lock()
.unwrap_or_else(|poison| poison.into_inner())
.push_back(outcome);
}
}
pub struct Session {
pub(crate) shared: Arc<SessionShared>,
}
impl Session {
pub fn new(
submission_capacity: usize,
completion_capacity: usize,
) -> Result<(Session, Receiver), SessionError> {
if submission_capacity < MINIMUM_SUBMISSION_CAPACITY {
return Err(SessionError::new(
SessionFailure::SubmissionCapacityTooSmall,
));
}
if completion_capacity < MINIMUM_COMPLETION_RING_CAPACITY {
return Err(SessionError::new(
SessionFailure::CompletionCapacityTooSmall,
));
}
let shared = Arc::new(SessionShared {
completions: Arc::new(CompletionRing::new(completion_capacity)),
submissions: SubmissionRing::new(submission_capacity),
registry: Mutex::new(Registry::new()),
next_id: AtomicU64::new(1),
work: Mutex::new(None),
handles: AtomicUsize::new(2),
#[cfg(test)]
scripted: Mutex::new(std::collections::VecDeque::new()),
});
let servicer = {
let weak: Weak<SessionShared> = Arc::downgrade(&shared);
ThreadpoolWork::new(
move || {
if let Some(shared) = weak.upgrade() {
shared.drain_submissions();
}
},
None,
)
.map_err(|error| SessionError::with_source(SessionFailure::WorkObject, error))?
};
let engine = {
let weak: Weak<SessionShared> = Arc::downgrade(&shared);
let mut environment = CallbackEnviron::new();
environment.set_runs_long();
ThreadpoolWork::new(
move || {
if let Some(shared) = weak.upgrade() {
shared.run_engine_quantum();
}
},
Some(&mut environment),
)
.map_err(|error| SessionError::with_source(SessionFailure::WorkObject, error))?
};
*shared.work() = Some(SessionWork {
servicer,
engine,
#[cfg(test)]
suppressed: AtomicBool::new(false),
});
let abandon = shared
.submissions
.reserve_abandon()
.expect("a fresh submission ring always has room for the abandon slot");
let receiver = Receiver {
shared: Arc::clone(&shared),
abandon: Some(abandon),
};
Ok((Session { shared }, receiver))
}
#[must_use]
pub fn submission_capacity(&self) -> usize {
self.shared.submissions.capacity()
}
#[must_use]
pub fn completion_capacity(&self) -> usize {
self.shared.completions.capacity()
}
#[must_use]
pub fn enumerations(&self) -> usize {
self.shared.registered()
}
#[must_use]
pub fn is_abandoned(&self) -> bool {
self.shared.submissions.is_abandoned()
}
#[cfg(test)]
pub(crate) fn suppress_pool(&self) {
if let Some(work) = self.shared.work().as_ref() {
work.suppressed.store(true, Ordering::Release);
}
}
pub fn try_begin(&self, request: EnumerationRequest) -> Result<EnumerationHandle, BeginError> {
admission::try_begin(&self.shared, request)
}
pub fn try_begin_with_token(
&self,
request: EnumerationRequest,
token: ImpersonationToken,
) -> Result<EnumerationHandle, BeginError> {
admission::try_begin_with_token(&self.shared, request, token)
}
}
impl Clone for Session {
fn clone(&self) -> Self {
self.shared.completions.add_session();
self.shared.acquire_handle();
Self {
shared: Arc::clone(&self.shared),
}
}
}
impl Drop for Session {
fn drop(&mut self) {
self.shared.completions.remove_session();
self.shared.release_handle();
}
}
impl std::fmt::Debug for Session {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Session")
.field("submission_capacity", &self.submission_capacity())
.field("completion_capacity", &self.completion_capacity())
.field("enumerations", &self.enumerations())
.finish_non_exhaustive()
}
}
pub struct Receiver {
shared: Arc<SessionShared>,
abandon: Option<AbandonSlot>,
}
impl Receiver {
#[must_use]
pub fn try_recv(&self) -> Option<Completion> {
let record = self.shared.completions.try_take();
if record.is_some() {
self.shared.resume_parked();
}
record
}
#[must_use]
pub fn recv(&self) -> Option<Completion> {
let record = self.shared.completions.take_blocking(None);
if record.is_some() {
self.shared.resume_parked();
}
record
}
#[must_use]
pub fn recv_timeout(&self, timeout: Duration) -> Option<Completion> {
let record = self.shared.completions.take_blocking(Some(timeout));
if record.is_some() {
self.shared.resume_parked();
}
record
}
pub fn doorbell(&self) -> io::Result<BorrowedHandle<'_>> {
self.shared.completions.doorbell()
}
pub fn doorbell_owned(&self) -> io::Result<OwnedHandle> {
self.shared.completions.doorbell_owned()
}
#[must_use]
pub fn is_disconnected(&self) -> bool {
self.shared.completions.is_closed()
}
#[must_use]
pub fn len(&self) -> usize {
self.shared.completions.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn capacity(&self) -> usize {
self.shared.completions.capacity()
}
}
impl Drop for Receiver {
fn drop(&mut self) {
if let Some(slot) = self.abandon.take() {
let pushed = self.shared.submissions.push_abandon(slot);
self.shared.ring_servicer(pushed);
}
self.shared.release_handle();
}
}
impl std::fmt::Debug for Receiver {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Receiver")
.field("queued", &self.len())
.field("capacity", &self.capacity())
.field("disconnected", &self.is_disconnected())
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests;