use std::io;
use std::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, OwnedHandle};
use std::ptr;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::time::Duration;
use crate::rearm::RearmSuppression;
use windows_sys::Win32::Foundation::{FALSE, FILETIME, HANDLE, TRUE, WAIT_TIMEOUT};
use windows_sys::Win32::System::Threading::{
CloseThreadpoolWait, CreateEventW, CreateThreadpoolWait, PTP_CALLBACK_INSTANCE, PTP_WAIT,
SetThreadpoolWait, WaitForThreadpoolWaitCallbacks,
};
use windows_sys::core::BOOL;
use crate::callback_env::CallbackEnviron;
mod wait_status {
pub const SIGNALLED: u32 = 0;
}
const FILETIME_TICKS_PER_SECOND: u64 = 10_000_000;
const FILETIME_NANOS_PER_TICK: u32 = 100;
fn relative_filetime(timeout: Duration) -> FILETIME {
let ticks = timeout
.as_secs()
.saturating_mul(FILETIME_TICKS_PER_SECOND)
.saturating_add(u64::from(timeout.subsec_nanos() / FILETIME_NANOS_PER_TICK));
let ticks = i64::try_from(ticks).unwrap_or(i64::MAX);
let bits = (-ticks) as u64;
FILETIME {
dwLowDateTime: bits as u32,
dwHighDateTime: (bits >> 32) as u32,
}
}
pub type WaitCloseFn = unsafe extern "system" fn(HANDLE) -> BOOL;
pub(crate) struct CustomClose {
raw: HANDLE,
close: WaitCloseFn,
}
impl Drop for CustomClose {
fn drop(&mut self) {
crate::trace_call!("WaitCloseFn", self.raw as usize, self.close as usize, {
unsafe { (self.close)(self.raw) };
});
}
}
pub(crate) enum WaitTarget {
Owned(OwnedHandle),
Custom(CustomClose),
}
unsafe impl Send for WaitTarget {}
unsafe impl Sync for WaitTarget {}
impl WaitTarget {
pub(crate) fn raw(&self) -> HANDLE {
match self {
WaitTarget::Owned(handle) => handle.as_raw_handle(),
WaitTarget::Custom(custom) => custom.raw,
}
}
pub(crate) fn borrow(&self) -> BorrowedHandle<'_> {
match self {
WaitTarget::Owned(handle) => handle.as_handle(),
WaitTarget::Custom(custom) => unsafe { BorrowedHandle::borrow_raw(custom.raw) },
}
}
}
impl std::fmt::Debug for WaitTarget {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WaitTarget")
.field("raw", &self.raw())
.field(
"close",
&match self {
WaitTarget::Owned(_) => "CloseHandle",
WaitTarget::Custom(_) => "custom",
},
)
.finish()
}
}
#[derive(Debug)]
pub struct WaitableHandle {
target: WaitTarget,
}
impl WaitableHandle {
pub fn event(manual_reset: bool, initially_signalled: bool) -> io::Result<Self> {
let raw = unsafe {
CreateEventW(
ptr::null(),
if manual_reset { TRUE } else { FALSE },
if initially_signalled { TRUE } else { FALSE },
ptr::null(),
)
};
if raw.is_null() {
return Err(io::Error::last_os_error());
}
Ok(Self {
target: WaitTarget::Owned(unsafe { OwnedHandle::from_raw_handle(raw) }),
})
}
#[must_use]
pub unsafe fn assume_waitable(handle: OwnedHandle) -> Self {
Self {
target: WaitTarget::Owned(handle),
}
}
#[must_use]
pub unsafe fn assume_waitable_with(handle: HANDLE, close: WaitCloseFn) -> Self {
Self {
target: WaitTarget::Custom(CustomClose { raw: handle, close }),
}
}
#[must_use]
pub fn handle(&self) -> BorrowedHandle<'_> {
self.target.borrow()
}
pub fn into_handle(self) -> Result<OwnedHandle, Self> {
match self.target {
WaitTarget::Owned(handle) => Ok(handle),
target @ WaitTarget::Custom(_) => Err(Self { target }),
}
}
pub(crate) fn into_target(self) -> WaitTarget {
self.target
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WaitResult {
Signalled,
TimedOut,
Other(u32),
}
impl WaitResult {
fn from_raw(value: u32) -> Self {
match value {
wait_status::SIGNALLED => Self::Signalled,
WAIT_TIMEOUT => Self::TimedOut,
other => Self::Other(other),
}
}
}
struct WaitContext {
wait: AtomicIsize,
handle: HANDLE,
rearm: RearmSuppression,
obligation: crate::obligation::CloseObligation,
registration: crate::heal::Registration,
callback: Box<dyn Fn(&WaitActivation<'_>) + Send + Sync + 'static>,
}
impl WaitContext {
fn suppression(&self) -> std::sync::MutexGuard<'_, u32> {
self.rearm.lock()
}
fn suppress_and_disarm(&self) {
self.rearm.suppress_and(|suppressed| {
let wait = self.wait.load(Ordering::Acquire);
crate::trace_record!("wait", "suppress-and-disarm", wait, suppressed);
if wait == 0 {
return;
}
unsafe { disarm_raw(wait) };
});
}
fn release_suppression(&self) {
self.rearm.release();
}
}
unsafe impl Send for WaitContext {}
unsafe impl Sync for WaitContext {}
pub struct WaitActivation<'ctx> {
result: WaitResult,
ctx: &'ctx WaitContext,
}
impl std::fmt::Debug for WaitActivation<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WaitActivation")
.field("result", &self.result)
.finish_non_exhaustive()
}
}
impl WaitActivation<'_> {
#[must_use]
pub fn result(&self) -> WaitResult {
self.result
}
#[must_use]
pub fn is_signalled(&self) -> bool {
self.result == WaitResult::Signalled
}
#[must_use]
pub fn handle(&self) -> BorrowedHandle<'_> {
unsafe { BorrowedHandle::borrow_raw(self.ctx.handle) }
}
pub fn rearm(&self, timeout: Option<Duration>) {
let _ = self.rearm_reporting(timeout);
}
pub(crate) fn rearm_reporting(&self, timeout: Option<Duration>) -> bool {
let wait = self.ctx.wait.load(Ordering::Acquire);
crate::trace_record!("wait", "rearm-entered", wait);
let suppressed = self.ctx.suppression();
if *suppressed > 0 {
crate::trace_record!("wait", "rearm-suppressed", wait, *suppressed);
return false;
}
debug_assert_ne!(
wait, 0,
"the wait object must be published before callbacks"
);
self.ctx.obligation.record_live_before(|| {
unsafe { arm_raw(wait, self.ctx.handle, timeout) };
});
drop(suppressed);
crate::trace_record!("wait", "rearm-left", wait, self.ctx.handle as usize);
true
}
}
pub(crate) unsafe fn disarm_raw(wait: PTP_WAIT) {
crate::trace_call!("SetThreadpoolWait(disarm)", wait, 0, {
unsafe { SetThreadpoolWait(wait, ptr::null_mut(), ptr::null()) };
});
crate::trace_record!("wait", "disarmed", wait);
}
pub(crate) unsafe fn arm_member(wait: PTP_WAIT, target: &WaitTarget, timeout: Option<Duration>) {
unsafe { arm_raw(wait, target.raw(), timeout) };
}
unsafe fn arm_raw(wait: PTP_WAIT, handle: HANDLE, timeout: Option<Duration>) {
crate::trace_call!("SetThreadpoolWait", wait, handle as usize, {
match timeout {
Some(timeout) => {
let filetime = relative_filetime(timeout);
unsafe { SetThreadpoolWait(wait, handle, &filetime) };
}
None => unsafe { SetThreadpoolWait(wait, handle, ptr::null()) },
}
});
crate::trace_record!("wait", "armed", wait, handle as usize);
}
unsafe extern "system" fn wait_trampoline(
_instance: PTP_CALLBACK_INSTANCE,
context: *mut core::ffi::c_void,
_wait: PTP_WAIT,
wait_result: u32,
) {
crate::trace_record!("wait", "trampoline-entered", _wait, wait_result);
let ctx = unsafe { &*(context as *const WaitContext) };
ctx.obligation.record_settled();
let activation = WaitActivation {
result: WaitResult::from_raw(wait_result),
ctx,
};
(ctx.callback)(&activation);
crate::trace_record!("wait", "trampoline-left", _wait, wait_result);
}
pub struct ThreadpoolWait {
wait: PTP_WAIT,
target: WaitTarget,
context: *mut WaitContext,
}
unsafe impl Send for ThreadpoolWait {}
unsafe impl Sync for ThreadpoolWait {}
impl ThreadpoolWait {
pub fn new<F>(
handle: WaitableHandle,
callback: F,
env: Option<&mut CallbackEnviron<'_>>,
) -> io::Result<Self>
where
F: Fn(&WaitActivation<'_>) + Send + Sync + 'static,
{
let target = handle.into_target();
let registration = crate::heal::register(crate::heal::key_of(env.as_deref()));
let context = Box::into_raw(Box::new(WaitContext {
wait: AtomicIsize::new(0),
handle: target.raw(),
rearm: RearmSuppression::new(),
obligation: crate::obligation::CloseObligation::new(),
registration,
callback: Box::new(callback),
}));
let env_ptr = env.map_or(ptr::null_mut(), |e| e.as_mut_ptr());
let wait = crate::trace_call!("CreateThreadpoolWait", 0, 0, {
unsafe {
CreateThreadpoolWait(Some(wait_trampoline), context.cast(), env_ptr.cast_const())
}
});
if wait == 0 {
let error = io::Error::last_os_error();
unsafe { drop(Box::from_raw(context)) };
return Err(error);
}
unsafe { (*context).wait.store(wait, Ordering::Release) };
crate::trace_record!("wait", "created", wait, target.raw() as usize);
Ok(Self {
wait,
target,
context,
})
}
#[must_use]
pub fn handle(&self) -> BorrowedHandle<'_> {
self.target.borrow()
}
pub fn arm(&self, timeout: Option<Duration>) {
unsafe { &*self.context }.obligation.record_live_before(|| {
unsafe { arm_raw(self.wait, self.target.raw(), timeout) };
});
}
pub fn disarm(&self) {
crate::trace_call!("SetThreadpoolWait(disarm)", self.wait, 0, {
unsafe { SetThreadpoolWait(self.wait, ptr::null_mut(), ptr::null()) };
});
crate::trace_record!("wait", "disarmed", self.wait);
}
pub fn wait(&self) {
crate::trace_call!("WaitForThreadpoolWaitCallbacks", self.wait, 0, {
unsafe { WaitForThreadpoolWaitCallbacks(self.wait, FALSE) };
});
}
#[cfg(feature = "self-heal")]
pub fn try_cancel_pending(&self) {
unsafe { self.try_cancel_pending_no_heal_tracking() };
let tracked = unsafe { &*self.context }.registration.owe_repair();
crate::obligation::fail_fast_if_untracked(tracked, "ThreadpoolWait");
}
pub unsafe fn try_cancel_pending_no_heal_tracking(&self) {
crate::trace_call!("WaitForThreadpoolWaitCallbacks(cancel)", self.wait, 1, {
unsafe { WaitForThreadpoolWaitCallbacks(self.wait, TRUE) };
});
}
pub(crate) unsafe fn recover_repair(
context: *mut core::ffi::c_void,
) -> Option<crate::heal::Registration> {
let ctx = unsafe { &*context.cast::<WaitContext>() };
Some(ctx.registration.reclaim())
}
#[cfg(feature = "self-heal")]
#[must_use]
pub(crate) unsafe fn owe_repair(context: *mut core::ffi::c_void) -> bool {
let ctx = unsafe { &*context.cast::<WaitContext>() };
ctx.registration.owe_repair()
}
pub fn stop_and_drain(&self) {
unsafe { Self::stop_and_drain_parts(self.context.cast(), self.wait) };
}
pub(crate) unsafe fn stop_and_drain_parts(context: *mut core::ffi::c_void, wait: PTP_WAIT) {
let ctx = unsafe { &*context.cast::<WaitContext>() };
ctx.suppress_and_disarm();
crate::trace_call!("WaitForThreadpoolWaitCallbacks", wait, 0, {
unsafe { WaitForThreadpoolWaitCallbacks(wait, FALSE) };
});
ctx.obligation.record_settled();
ctx.release_suppression();
}
pub(crate) fn into_parts(self) -> (PTP_WAIT, *mut core::ffi::c_void, WaitTarget) {
let this = std::mem::ManuallyDrop::new(self);
let target = unsafe { ptr::read(&this.target) };
(this.wait, this.context.cast(), target)
}
pub(crate) unsafe fn drop_context(context: *mut core::ffi::c_void) {
drop(unsafe { Box::from_raw(context.cast::<WaitContext>()) });
}
#[cfg(test)]
pub(crate) fn obligation_owed(&self) -> bool {
unsafe { &*self.context }.obligation.is_owed()
}
pub(crate) unsafe fn prepare_shutdown(context: *mut core::ffi::c_void) {
let ctx = unsafe { &*context.cast::<WaitContext>() };
ctx.suppress_and_disarm();
}
}
impl Drop for ThreadpoolWait {
fn drop(&mut self) {
let ctx = unsafe { &*self.context };
crate::trace_record!("wait", "drop-begin", self.wait, self.target.raw() as usize);
let owed = ctx.obligation.is_owed();
if owed {
crate::trace_record!("wait", crate::obligation::DROP_OBLIGATION_OWED, self.wait);
}
ctx.suppress_and_disarm();
self.wait();
crate::trace_record!("wait", "drop-drained", self.wait);
crate::trace_call!("CloseThreadpoolWait", self.wait, 0, {
unsafe { CloseThreadpoolWait(self.wait) };
});
unsafe {
drop(Box::from_raw(self.context));
}
crate::trace_record!("wait", "drop-closed", self.wait, self.target.raw() as usize);
crate::obligation::fail_fast_if_owed(owed, "ThreadpoolWait", "stop_and_drain");
}
}
#[cfg(test)]
mod tests;