use std::io;
use std::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, FromRawHandle, OwnedHandle};
use std::ptr;
use std::sync::Mutex;
use std::sync::atomic::{AtomicIsize, Ordering};
use std::time::Duration;
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) {
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,
suppress_rearm: Mutex<u32>,
callback: Box<dyn Fn(&WaitActivation<'_>) + Send + Sync + 'static>,
}
impl WaitContext {
fn suppression(&self) -> std::sync::MutexGuard<'_, u32> {
self.suppress_rearm
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
fn suppress_and_disarm(&self) {
let mut suppressed = self.suppression();
*suppressed = suppressed.saturating_add(1);
let wait = self.wait.load(Ordering::Acquire);
if wait != 0 {
unsafe { disarm_raw(wait) };
}
}
fn release_suppression(&self) {
let mut suppressed = self.suppression();
*suppressed = suppressed.saturating_sub(1);
}
}
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 suppressed = self.ctx.suppression();
if *suppressed > 0 {
return false;
}
let wait = self.ctx.wait.load(Ordering::Acquire);
debug_assert_ne!(
wait, 0,
"the wait object must be published before callbacks"
);
unsafe { arm_raw(wait, self.ctx.handle, timeout) };
drop(suppressed);
true
}
}
pub(crate) unsafe fn disarm_raw(wait: PTP_WAIT) {
unsafe { SetThreadpoolWait(wait, ptr::null_mut(), ptr::null()) };
}
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>) {
match timeout {
Some(timeout) => {
let filetime = relative_filetime(timeout);
unsafe { SetThreadpoolWait(wait, handle, &filetime) };
}
None => unsafe { SetThreadpoolWait(wait, handle, ptr::null()) },
}
}
unsafe extern "system" fn wait_trampoline(
_instance: PTP_CALLBACK_INSTANCE,
context: *mut core::ffi::c_void,
_wait: PTP_WAIT,
wait_result: u32,
) {
let ctx = unsafe { &*(context as *const WaitContext) };
let activation = WaitActivation {
result: WaitResult::from_raw(wait_result),
ctx,
};
(ctx.callback)(&activation);
}
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 context = Box::into_raw(Box::new(WaitContext {
wait: AtomicIsize::new(0),
handle: target.raw(),
suppress_rearm: Mutex::new(0),
callback: Box::new(callback),
}));
let env_ptr = env.map_or(ptr::null_mut(), |e| e.as_mut_ptr());
let wait = 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) };
Ok(Self {
wait,
target,
context,
})
}
#[must_use]
pub fn handle(&self) -> BorrowedHandle<'_> {
self.target.borrow()
}
pub fn arm(&self, timeout: Option<Duration>) {
unsafe { arm_raw(self.wait, self.target.raw(), timeout) };
}
pub fn disarm(&self) {
unsafe { SetThreadpoolWait(self.wait, ptr::null_mut(), ptr::null()) };
}
pub fn wait(&self) {
unsafe { WaitForThreadpoolWaitCallbacks(self.wait, FALSE) };
}
pub fn cancel_pending(&self) {
unsafe { WaitForThreadpoolWaitCallbacks(self.wait, TRUE) };
}
pub fn stop_and_drain(&self) {
let ctx = unsafe { &*self.context };
ctx.suppress_and_disarm();
self.cancel_pending();
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>()) });
}
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 };
ctx.suppress_and_disarm();
self.cancel_pending();
unsafe {
CloseThreadpoolWait(self.wait);
drop(Box::from_raw(self.context));
}
}
}
#[cfg(test)]
mod tests;