use std::ffi::c_void;
use std::io;
use std::mem::ManuallyDrop;
use std::os::windows::io::{AsRawHandle, OwnedHandle};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use windows_sys::Win32::Foundation::HANDLE;
use windows_sys::Win32::Storage::FileSystem::{
BuildIoRingCancelRequest, BuildIoRingFlushFile, BuildIoRingReadFile,
BuildIoRingRegisterBuffers, BuildIoRingRegisterFileHandles, BuildIoRingWriteFile,
FILE_FLUSH_DEFAULT, FILE_WRITE_FLAGS_NONE, IORING_BUFFER_INFO, IORING_BUFFER_REF,
IORING_BUFFER_REF_0, IORING_HANDLE_REF, IORING_HANDLE_REF_0, IORING_REF_RAW,
IORING_REF_REGISTERED, IORING_REGISTERED_BUFFER, IORING_SQE_FLAGS,
IOSQE_FLAGS_DRAIN_PRECEDING_OPS, IOSQE_FLAGS_NONE, SubmitIoRing,
};
use crate::buf::{IoBuf, IoBufMut};
use crate::error::check;
use crate::ring::{Completion, IoRing, Op, RingId};
use crate::token::Token;
#[derive(Clone, Copy, Debug, Default)]
#[must_use]
pub struct PushOptions {
drain_preceding: bool,
}
impl PushOptions {
pub fn new() -> Self {
Self::default()
}
pub fn drain_preceding(mut self, drain: bool) -> Self {
self.drain_preceding = drain;
self
}
fn sqe_flags(self) -> IORING_SQE_FLAGS {
if self.drain_preceding {
IOSQE_FLAGS_DRAIN_PRECEDING_OPS
} else {
IOSQE_FLAGS_NONE
}
}
}
fn handle_ref(file: FileRef, ring_id: RingId) -> io::Result<IORING_HANDLE_REF> {
match file {
FileRef::Raw(handle) => Ok(IORING_HANDLE_REF {
Kind: IORING_REF_RAW,
Handle: IORING_HANDLE_REF_0 { Handle: handle },
}),
FileRef::Registered(file) => {
if file.ring_id != ring_id {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"this RegisteredFile was registered on a different IoRing",
));
}
Ok(IORING_HANDLE_REF {
Kind: IORING_REF_REGISTERED,
Handle: IORING_HANDLE_REF_0 {
Index: file.index(),
},
})
}
}
}
fn raw_buffer_ref(address: *mut c_void) -> IORING_BUFFER_REF {
IORING_BUFFER_REF {
Kind: IORING_REF_RAW,
Buffer: IORING_BUFFER_REF_0 { Address: address },
}
}
fn registered_buffer_ref(index: u32, offset: u32) -> IORING_BUFFER_REF {
IORING_BUFFER_REF {
Kind: IORING_REF_REGISTERED,
Buffer: IORING_BUFFER_REF_0 {
IndexAndOffset: IORING_REGISTERED_BUFFER {
BufferIndex: index,
Offset: offset,
},
},
}
}
fn checked_len(len: usize) -> io::Result<u32> {
u32::try_from(len).map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
format!("an IoRing buffer is limited to u32::MAX bytes; {len} does not fit"),
)
})
}
#[derive(Clone, Copy, Debug)]
pub enum FileRef {
Raw(HANDLE),
Registered(RegisteredFile),
}
impl From<HANDLE> for FileRef {
fn from(handle: HANDLE) -> Self {
FileRef::Raw(handle)
}
}
impl From<RegisteredFile> for FileRef {
fn from(file: RegisteredFile) -> Self {
FileRef::Registered(file)
}
}
#[derive(Clone, Debug)]
pub struct SharedFile(Arc<OwnedHandle>);
impl SharedFile {
#[must_use]
pub fn new(handle: OwnedHandle) -> Self {
Self(Arc::new(handle))
}
fn raw_handle(&self) -> HANDLE {
self.0.as_raw_handle()
}
}
impl From<OwnedHandle> for SharedFile {
fn from(handle: OwnedHandle) -> Self {
Self::new(handle)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RegisteredFile {
index: u32,
ring_id: RingId,
}
impl RegisteredFile {
#[must_use]
pub fn index(self) -> u32 {
self.index
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RegisteredFiles {
base_index: u32,
count: u32,
ring_id: RingId,
}
impl RegisteredFiles {
#[must_use]
pub fn len(&self) -> u32 {
self.count
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.count == 0
}
#[must_use]
pub fn get(&self, i: u32) -> Option<RegisteredFile> {
(i < self.count).then(|| RegisteredFile {
index: self.base_index + i,
ring_id: self.ring_id,
})
}
}
#[derive(Debug)]
pub struct PendingFileRegistration {
user_data: usize,
base_index: u32,
count: u32,
ring_id: RingId,
}
impl PendingFileRegistration {
#[must_use]
pub fn user_data(&self) -> usize {
self.user_data
}
pub fn claim_if(self, completion: &Completion) -> Result<io::Result<RegisteredFiles>, Self> {
if completion.user_data() != self.user_data || completion.ring_id() != self.ring_id {
return Err(self);
}
Ok(completion.result().map(|_| RegisteredFiles {
base_index: self.base_index,
count: self.count,
ring_id: self.ring_id,
}))
}
}
pub struct RegisteredUse(Arc<AtomicUsize>);
impl Drop for RegisteredUse {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::SeqCst);
}
}
pub struct RegisteredBuffers<B: IoBufMut> {
buffers: ManuallyDrop<Vec<B>>,
base_index: u32,
outstanding: Arc<AtomicUsize>,
ring_id: RingId,
}
impl<B: IoBufMut> RegisteredBuffers<B> {
#[must_use]
pub fn len(&self) -> u32 {
u32::try_from(self.buffers.len()).unwrap_or(u32::MAX)
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.buffers.is_empty()
}
#[must_use]
pub fn get(&self, i: u32) -> Option<&B> {
self.buffers.get(i as usize)
}
fn checked_index(&self, i: u32) -> io::Result<u32> {
if i < self.len() {
Ok(self.base_index + i)
} else {
Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("buffer index {i} is out of range for this registration"),
))
}
}
fn checked_span(&self, span: RegisteredSpan) -> io::Result<u32> {
let index = self.checked_index(span.buffer_index)?;
let buffer = self
.buffers
.get(span.buffer_index as usize)
.expect("checked_index just validated this index");
let bytes_len = buffer.bytes_len();
let fits = u64::from(span.offset)
.checked_add(u64::from(span.len))
.is_some_and(|end| end <= bytes_len as u64);
if !fits {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"span offset {} + len {} does not fit inside buffer {} ({bytes_len} bytes)",
span.offset, span.len, span.buffer_index
),
));
}
Ok(index)
}
}
impl<B: IoBufMut> Drop for RegisteredBuffers<B> {
fn drop(&mut self) {
if self.outstanding.load(Ordering::SeqCst) > 0 {
debug_assert!(
false,
"RegisteredBuffers dropped while an operation still references it"
);
return;
}
unsafe { ManuallyDrop::drop(&mut self.buffers) };
}
}
#[derive(Clone, Copy, Debug)]
pub struct RegisteredSpan {
pub buffer_index: u32,
pub offset: u32,
pub len: u32,
}
pub struct PendingBufferRegistration<B: IoBufMut> {
user_data: usize,
base_index: u32,
buffers: ManuallyDrop<Vec<B>>,
ring_id: RingId,
}
impl<B: IoBufMut> PendingBufferRegistration<B> {
#[must_use]
pub fn user_data(&self) -> usize {
self.user_data
}
pub fn claim_if(
mut self,
completion: &Completion,
) -> Result<io::Result<RegisteredBuffers<B>>, Self> {
if completion.user_data() != self.user_data || completion.ring_id() != self.ring_id {
return Err(self);
}
let buffers = unsafe { ManuallyDrop::take(&mut self.buffers) };
let base_index = self.base_index;
let ring_id = self.ring_id;
Ok(completion.result().map(move |_| RegisteredBuffers {
buffers: ManuallyDrop::new(buffers),
base_index,
outstanding: Arc::new(AtomicUsize::new(0)),
ring_id,
}))
}
}
impl<B: IoBufMut> Drop for PendingBufferRegistration<B> {
fn drop(&mut self) {
}
}
impl<B: IoBufMut> std::fmt::Debug for PendingBufferRegistration<B> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PendingBufferRegistration")
.field("user_data", &self.user_data)
.field("base_index", &self.base_index)
.finish_non_exhaustive()
}
}
pub struct Batch<'ring> {
ring: &'ring mut IoRing,
submitted: bool,
}
impl<'ring> Batch<'ring> {
pub fn new(ring: &'ring mut IoRing) -> Self {
Self {
ring,
submitted: false,
}
}
fn require(&self, op: Op) -> io::Result<()> {
if self.ring.supports(op) {
Ok(())
} else {
Err(io::Error::new(
io::ErrorKind::Unsupported,
format!("this ring does not support {op:?}"),
))
}
}
pub unsafe fn read_raw<B: IoBufMut>(
&mut self,
file: impl Into<FileRef>,
mut buffer: B,
offset: u64,
options: PushOptions,
) -> io::Result<Token<B>> {
self.require(Op::Read)?;
let len = checked_len(buffer.bytes_len())?;
let address = buffer.stable_mut_ptr().cast::<c_void>();
let target = handle_ref(file.into(), self.ring.ring_id())?;
let token = Token::new(self.ring, buffer)?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingReadFile(
self.ring.raw_handle(),
target,
raw_buffer_ref(address),
len,
offset,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub fn read<B: IoBufMut>(
&mut self,
file: &SharedFile,
mut buffer: B,
offset: u64,
options: PushOptions,
) -> io::Result<Token<(B, SharedFile)>> {
self.require(Op::Read)?;
let len = checked_len(buffer.bytes_len())?;
let address = buffer.stable_mut_ptr().cast::<c_void>();
let raw = file.raw_handle();
let target = handle_ref(FileRef::Raw(raw), self.ring.ring_id())?;
let token = Token::new(self.ring, (buffer, file.clone()))?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingReadFile(
self.ring.raw_handle(),
target,
raw_buffer_ref(address),
len,
offset,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub unsafe fn write_raw<B: IoBuf>(
&mut self,
file: impl Into<FileRef>,
buffer: B,
offset: u64,
options: PushOptions,
) -> io::Result<Token<B>> {
self.require(Op::Write)?;
let len = checked_len(buffer.bytes_len())?;
let address = buffer.stable_ptr().cast_mut().cast::<c_void>();
let target = handle_ref(file.into(), self.ring.ring_id())?;
let token = Token::new(self.ring, buffer)?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingWriteFile(
self.ring.raw_handle(),
target,
raw_buffer_ref(address),
len,
offset,
FILE_WRITE_FLAGS_NONE,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub fn write<B: IoBuf>(
&mut self,
file: &SharedFile,
buffer: B,
offset: u64,
options: PushOptions,
) -> io::Result<Token<(B, SharedFile)>> {
self.require(Op::Write)?;
let len = checked_len(buffer.bytes_len())?;
let address = buffer.stable_ptr().cast_mut().cast::<c_void>();
let raw = file.raw_handle();
let target = handle_ref(FileRef::Raw(raw), self.ring.ring_id())?;
let token = Token::new(self.ring, (buffer, file.clone()))?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingWriteFile(
self.ring.raw_handle(),
target,
raw_buffer_ref(address),
len,
offset,
FILE_WRITE_FLAGS_NONE,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
fn check_registration_ring<B: IoBufMut>(
&self,
registration: &RegisteredBuffers<B>,
) -> io::Result<()> {
if registration.ring_id == self.ring.ring_id() {
Ok(())
} else {
Err(io::Error::new(
io::ErrorKind::InvalidInput,
"this RegisteredBuffers was registered on a different IoRing",
))
}
}
fn finish_push<T: Send + 'static>(
&mut self,
hr: windows_sys::core::HRESULT,
token: Token<T>,
) -> io::Result<Token<T>> {
match check(hr) {
Ok(()) => Ok(token),
Err(error) => {
let _ = token.claim();
self.ring.cancel_reservation();
Err(error)
}
}
}
pub unsafe fn flush_raw(
&mut self,
file: impl Into<FileRef>,
options: PushOptions,
) -> io::Result<usize> {
self.require(Op::Flush)?;
let target = handle_ref(file.into(), self.ring.ring_id())?;
let user_data = self.ring.reserve_user_data()?;
let hr = unsafe {
BuildIoRingFlushFile(
self.ring.raw_handle(),
target,
FILE_FLUSH_DEFAULT,
user_data,
options.sqe_flags(),
)
};
if let Err(error) = check(hr) {
self.ring.cancel_reservation();
return Err(error);
}
Ok(user_data)
}
pub fn flush(
&mut self,
file: &SharedFile,
options: PushOptions,
) -> io::Result<Token<SharedFile>> {
self.require(Op::Flush)?;
let raw = file.raw_handle();
let target = handle_ref(FileRef::Raw(raw), self.ring.ring_id())?;
let token = Token::new(self.ring, file.clone())?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingFlushFile(
self.ring.raw_handle(),
target,
FILE_FLUSH_DEFAULT,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub unsafe fn cancel_raw(
&mut self,
file: impl Into<FileRef>,
target: usize,
) -> io::Result<usize> {
self.require(Op::Cancel)?;
let handle = handle_ref(file.into(), self.ring.ring_id())?;
let user_data = self.ring.reserve_user_data()?;
let hr =
unsafe { BuildIoRingCancelRequest(self.ring.raw_handle(), handle, target, user_data) };
if let Err(error) = check(hr) {
self.ring.cancel_reservation();
return Err(error);
}
Ok(user_data)
}
pub fn cancel(&mut self, file: &SharedFile, target: usize) -> io::Result<Token<SharedFile>> {
self.require(Op::Cancel)?;
let raw = file.raw_handle();
let handle = handle_ref(FileRef::Raw(raw), self.ring.ring_id())?;
let token = Token::new(self.ring, file.clone())?;
let user_data = token.id();
let hr =
unsafe { BuildIoRingCancelRequest(self.ring.raw_handle(), handle, target, user_data) };
self.finish_push(hr, token)
}
pub unsafe fn register_files(
&mut self,
handles: &[HANDLE],
) -> io::Result<PendingFileRegistration> {
self.require(Op::RegisterFiles)?;
if self.ring.registered_file_count() > 0 {
return Err(io::Error::new(
io::ErrorKind::AlreadyExists,
"this ring already has a file-handle registration; BuildIoRingRegisterFileHandles \
replaces the whole table, so a second call would invalidate every RegisteredFile \
index already handed out",
));
}
let count = checked_len(handles.len())?;
let base_index = self.ring.registered_file_count();
let user_data = self.ring.reserve_user_data()?;
let hr = unsafe {
BuildIoRingRegisterFileHandles(
self.ring.raw_handle(),
count,
handles.as_ptr(),
user_data,
)
};
if let Err(error) = check(hr) {
self.ring.cancel_reservation();
return Err(error);
}
self.ring.reserve_registered_files(count);
Ok(PendingFileRegistration {
user_data,
base_index,
count,
ring_id: self.ring.ring_id(),
})
}
pub fn register_buffers<B: IoBufMut>(
&mut self,
mut buffers: Vec<B>,
) -> io::Result<PendingBufferRegistration<B>> {
self.require(Op::RegisterBuffers)?;
if self.ring.registered_buffer_count() > 0 {
return Err(io::Error::new(
io::ErrorKind::AlreadyExists,
"this ring already has a buffer registration; BuildIoRingRegisterBuffers \
replaces the whole table, so a second call would invalidate every buffer \
index already handed out",
));
}
let count = checked_len(buffers.len())?;
let base_index = self.ring.registered_buffer_count();
let mut infos = Vec::with_capacity(buffers.len());
for buffer in &mut buffers {
let length = checked_len(buffer.bytes_len())?;
infos.push(IORING_BUFFER_INFO {
Address: buffer.stable_mut_ptr().cast::<c_void>(),
Length: length,
});
}
let user_data = self.ring.reserve_user_data()?;
let hr = unsafe {
BuildIoRingRegisterBuffers(self.ring.raw_handle(), count, infos.as_ptr(), user_data)
};
if let Err(error) = check(hr) {
self.ring.cancel_reservation();
return Err(error);
}
self.ring.reserve_registered_buffers(count);
Ok(PendingBufferRegistration {
user_data,
base_index,
buffers: ManuallyDrop::new(buffers),
ring_id: self.ring.ring_id(),
})
}
pub unsafe fn read_registered_raw<B: IoBufMut>(
&mut self,
file: impl Into<FileRef>,
registration: &RegisteredBuffers<B>,
span: RegisteredSpan,
file_offset: u64,
options: PushOptions,
) -> io::Result<Token<RegisteredUse>> {
self.require(Op::Read)?;
self.check_registration_ring(registration)?;
let target = handle_ref(file.into(), self.ring.ring_id())?;
let index = registration.checked_span(span)?;
registration.outstanding.fetch_add(1, Ordering::SeqCst);
let token = Token::new(
self.ring,
RegisteredUse(Arc::clone(®istration.outstanding)),
)?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingReadFile(
self.ring.raw_handle(),
target,
registered_buffer_ref(index, span.offset),
span.len,
file_offset,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub fn read_registered<B: IoBufMut>(
&mut self,
file: &SharedFile,
registration: &RegisteredBuffers<B>,
span: RegisteredSpan,
file_offset: u64,
options: PushOptions,
) -> io::Result<Token<(RegisteredUse, SharedFile)>> {
self.require(Op::Read)?;
self.check_registration_ring(registration)?;
let raw = file.raw_handle();
let target = handle_ref(FileRef::Raw(raw), self.ring.ring_id())?;
let index = registration.checked_span(span)?;
registration.outstanding.fetch_add(1, Ordering::SeqCst);
let token = Token::new(
self.ring,
(
RegisteredUse(Arc::clone(®istration.outstanding)),
file.clone(),
),
)?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingReadFile(
self.ring.raw_handle(),
target,
registered_buffer_ref(index, span.offset),
span.len,
file_offset,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub unsafe fn write_registered_raw<B: IoBufMut>(
&mut self,
file: impl Into<FileRef>,
registration: &RegisteredBuffers<B>,
span: RegisteredSpan,
file_offset: u64,
options: PushOptions,
) -> io::Result<Token<RegisteredUse>> {
self.require(Op::Write)?;
self.check_registration_ring(registration)?;
let target = handle_ref(file.into(), self.ring.ring_id())?;
let index = registration.checked_span(span)?;
registration.outstanding.fetch_add(1, Ordering::SeqCst);
let token = Token::new(
self.ring,
RegisteredUse(Arc::clone(®istration.outstanding)),
)?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingWriteFile(
self.ring.raw_handle(),
target,
registered_buffer_ref(index, span.offset),
span.len,
file_offset,
FILE_WRITE_FLAGS_NONE,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub fn write_registered<B: IoBufMut>(
&mut self,
file: &SharedFile,
registration: &RegisteredBuffers<B>,
span: RegisteredSpan,
file_offset: u64,
options: PushOptions,
) -> io::Result<Token<(RegisteredUse, SharedFile)>> {
self.require(Op::Write)?;
self.check_registration_ring(registration)?;
let raw = file.raw_handle();
let target = handle_ref(FileRef::Raw(raw), self.ring.ring_id())?;
let index = registration.checked_span(span)?;
registration.outstanding.fetch_add(1, Ordering::SeqCst);
let token = Token::new(
self.ring,
(
RegisteredUse(Arc::clone(®istration.outstanding)),
file.clone(),
),
)?;
let user_data = token.id();
let hr = unsafe {
BuildIoRingWriteFile(
self.ring.raw_handle(),
target,
registered_buffer_ref(index, span.offset),
span.len,
file_offset,
FILE_WRITE_FLAGS_NONE,
user_data,
options.sqe_flags(),
)
};
self.finish_push(hr, token)
}
pub fn submit(self) -> io::Result<u32> {
self.submit_and_wait(0, 0)
}
pub fn submit_and_wait(mut self, wait_operations: u32, timeout_ms: u32) -> io::Result<u32> {
self.do_submit(wait_operations, timeout_ms)
}
fn do_submit(&mut self, wait_operations: u32, timeout_ms: u32) -> io::Result<u32> {
let mut submitted = 0_u32;
let hr = unsafe {
SubmitIoRing(
self.ring.raw_handle(),
wait_operations,
timeout_ms,
&raw mut submitted,
)
};
self.submitted = true;
check(hr)?;
Ok(submitted)
}
}
impl Drop for Batch<'_> {
fn drop(&mut self) {
if !self.submitted {
let _ = self.do_submit(0, 0);
}
}
}
#[cfg(test)]
mod tests;