use std::cell::Cell;
use std::fmt;
use std::io;
use std::os::windows::io::{AsHandle, AsRawHandle, BorrowedHandle, OwnedHandle};
use std::ptr;
use std::sync::Arc;
use windows_overlapped_io_sys::{
Issued, Operation, OperationId, OperationRegistry, OperationState, Submitted,
UnassociatedEndpoint, reclaim_overlapped,
};
use windows_sys::Win32::Foundation::{FALSE, HANDLE, NO_ERROR};
use windows_sys::Win32::System::IO::{CancelIoEx, OVERLAPPED};
use windows_sys::Win32::System::Threading::{
CancelThreadpoolIo, CloseThreadpoolIo, CreateThreadpoolIo, PTP_CALLBACK_INSTANCE, PTP_IO,
StartThreadpoolIo, WaitForThreadpoolIoCallbacks,
};
use crate::callback_env::CallbackEnviron;
struct IoContext {
live: Arc<OperationRegistry>,
callback: Box<dyn Fn(&IoCompletion) + Send + Sync + 'static>,
}
unsafe extern "system" fn io_trampoline(
_instance: PTP_CALLBACK_INSTANCE,
context: *mut core::ffi::c_void,
overlapped: *mut core::ffi::c_void,
io_result: u32,
bytes_transferred: usize,
_io: PTP_IO,
) {
let ctx = unsafe { &*(context as *const IoContext) };
let overlapped = overlapped.cast::<OVERLAPPED>();
let id = ctx.live.remove(overlapped);
let completion = IoCompletion {
overlapped,
id,
io_result,
bytes_transferred,
claimed: Cell::new(false),
};
(ctx.callback)(&completion);
}
pub struct ThreadpoolIo {
tp_io: PTP_IO,
handle: OwnedHandle,
context: *mut IoContext,
live: Arc<OperationRegistry>,
}
unsafe impl Send for ThreadpoolIo {}
unsafe impl Sync for ThreadpoolIo {}
impl ThreadpoolIo {
pub fn new<F>(
endpoint: UnassociatedEndpoint,
callback: F,
env: Option<&mut CallbackEnviron>,
) -> io::Result<Self>
where
F: Fn(&IoCompletion) + Send + Sync + 'static,
{
let handle = endpoint.into_handle();
let live = Arc::new(OperationRegistry::new());
let context = Box::into_raw(Box::new(IoContext {
live: Arc::clone(&live),
callback: Box::new(callback),
}));
let env_ptr = env.map_or(ptr::null_mut(), |e| e.as_mut_ptr());
let tp_io = unsafe {
CreateThreadpoolIo(
handle.as_raw_handle(),
Some(io_trampoline),
context.cast(),
env_ptr.cast_const(),
)
};
if tp_io == 0 {
let error = io::Error::last_os_error();
unsafe { drop(Box::from_raw(context)) };
return Err(error);
}
Ok(Self {
tp_io,
handle,
context,
live,
})
}
#[must_use]
pub fn handle(&self) -> BorrowedHandle<'_> {
self.handle.as_handle()
}
#[must_use]
pub fn outstanding(&self) -> usize {
self.live.len()
}
pub unsafe fn submit<P, F>(&self, operation: Operation<P>, issue: F) -> Submitted<P>
where
P: Send + 'static,
F: FnOnce(BorrowedHandle<'_>, *mut OVERLAPPED) -> io::Result<Issued>,
{
let overlapped = operation.into_overlapped();
let id = OperationId::mint(overlapped);
self.live.insert(id);
unsafe { StartThreadpoolIo(self.tp_io) };
match issue(self.handle(), overlapped) {
Ok(Issued::Pending) => Submitted::Pending(id),
Ok(Issued::Completed { bytes_transferred }) => {
unsafe { CancelThreadpoolIo(self.tp_io) };
self.live.remove(overlapped);
let mut operation = unsafe { Operation::<P>::from_overlapped(overlapped) };
operation.set_state(OperationState::Completed);
Submitted::Completed {
operation,
bytes_transferred,
}
}
Err(error) => {
unsafe { CancelThreadpoolIo(self.tp_io) };
self.live.remove(overlapped);
let mut operation = unsafe { Operation::<P>::from_overlapped(overlapped) };
operation.set_state(OperationState::Idle);
Submitted::Failed { operation, error }
}
}
}
pub fn cancel(&self, id: OperationId) -> io::Result<()> {
self.live.cancel_if_live(id, || {
let ok = unsafe { CancelIoEx(self.raw_handle(), id.as_ptr()) };
if ok == 0 {
return Err(io::Error::last_os_error());
}
Ok(())
})
}
pub fn cancel_all(&self) -> io::Result<()> {
let ok = unsafe { CancelIoEx(self.raw_handle(), ptr::null()) };
if ok == 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
pub fn run_down(&self) {
self.live.wait_until_empty();
self.wait();
}
pub fn wait(&self) {
unsafe { WaitForThreadpoolIoCallbacks(self.tp_io, FALSE) };
}
fn raw_handle(&self) -> HANDLE {
self.handle.as_raw_handle()
}
}
impl fmt::Debug for ThreadpoolIo {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ThreadpoolIo")
.field("outstanding", &self.outstanding())
.finish_non_exhaustive()
}
}
impl Drop for ThreadpoolIo {
fn drop(&mut self) {
let count = self.outstanding();
if count > 0 {
eprintln!(
"windows-threadpool-sys: ThreadpoolIo dropped with {count} operation(s) still \
outstanding; call cancel_all() and run_down() before dropping to control when \
this blocks."
);
let _ = self.cancel_all();
self.live.wait_until_empty();
}
unsafe {
WaitForThreadpoolIoCallbacks(self.tp_io, FALSE);
CloseThreadpoolIo(self.tp_io);
}
unsafe { drop(Box::from_raw(self.context)) };
}
}
pub struct IoCompletion {
overlapped: *mut OVERLAPPED,
id: Option<OperationId>,
io_result: u32,
bytes_transferred: usize,
claimed: Cell<bool>,
}
impl fmt::Debug for IoCompletion {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("IoCompletion")
.field("overlapped", &self.overlapped)
.field("id", &self.id)
.field("io_result", &self.io_result)
.field("bytes_transferred", &self.bytes_transferred)
.finish_non_exhaustive()
}
}
impl IoCompletion {
#[must_use]
pub fn bytes_transferred(&self) -> usize {
self.bytes_transferred
}
#[must_use]
pub fn io_result(&self) -> u32 {
self.io_result
}
#[must_use]
pub fn error(&self) -> Option<io::Error> {
if self.io_result == NO_ERROR {
return None;
}
Some(io::Error::from_raw_os_error(self.io_result as i32))
}
#[must_use]
pub fn id(&self) -> Option<OperationId> {
self.id
}
#[must_use]
pub fn overlapped_ptr(&self) -> *mut OVERLAPPED {
self.overlapped
}
pub unsafe fn claim<P>(&self) -> Operation<P> {
self.claimed.set(true);
let mut operation = unsafe { Operation::<P>::from_overlapped(self.overlapped) };
operation.set_state(OperationState::Completed);
operation
}
}
impl Drop for IoCompletion {
fn drop(&mut self) {
if self.claimed.get() || self.overlapped.is_null() {
return;
}
unsafe { reclaim_overlapped(self.overlapped) };
}
}
#[cfg(test)]
mod tests;