compio_driver/
key.rs

1use std::{io, marker::PhantomData, mem::MaybeUninit, pin::Pin, task::Waker};
2
3use compio_buf::BufResult;
4
5use crate::{OpCode, Overlapped, PushEntry, RawFd};
6
7/// An operation with other needed information. It should be allocated on the
8/// heap. The pointer to this struct is used as `user_data`, and on Windows, it
9/// is used as the pointer to `OVERLAPPED`.
10///
11/// `*const RawOp<dyn OpCode>` can be obtained from any `Key<T: OpCode>` by
12/// first casting `Key::user_data` to `*const RawOp<()>`, then upcasted with
13/// `upcast_fn`. It is done in [`Key::as_op_pin`].
14#[repr(C)]
15pub(crate) struct RawOp<T: ?Sized> {
16    header: Overlapped,
17    // The cancelled flag and the result here are manual reference counting. The driver holds the
18    // strong ref until it completes; the runtime holds the strong ref until the future is
19    // dropped.
20    cancelled: bool,
21    // The metadata in `*mut RawOp<dyn OpCode>`
22    metadata: usize,
23    result: PushEntry<Option<Waker>, io::Result<usize>>,
24    flags: u32,
25    op: T,
26}
27
28#[repr(C)]
29union OpCodePtrRepr {
30    ptr: *mut RawOp<dyn OpCode>,
31    components: OpCodePtrComponents,
32}
33
34#[repr(C)]
35#[derive(Clone, Copy)]
36struct OpCodePtrComponents {
37    data_pointer: *mut (),
38    metadata: usize,
39}
40
41fn opcode_metadata<T: OpCode + 'static>() -> usize {
42    let mut op = MaybeUninit::<RawOp<T>>::uninit();
43    // SAFETY: same as `core::ptr::metadata`.
44    unsafe {
45        OpCodePtrRepr {
46            ptr: op.as_mut_ptr(),
47        }
48        .components
49        .metadata
50    }
51}
52
53const unsafe fn opcode_dyn_mut(ptr: *mut (), metadata: usize) -> *mut RawOp<dyn OpCode> {
54    // SAFETY: same as `core::ptr::from_raw_parts_mut`.
55    unsafe {
56        OpCodePtrRepr {
57            components: OpCodePtrComponents {
58                data_pointer: ptr,
59                metadata,
60            },
61        }
62        .ptr
63    }
64}
65
66/// A typed wrapper for key of Ops submitted into driver. It doesn't free the
67/// inner on dropping. Instead, the memory is managed by the proactor. The inner
68/// is only freed when:
69///
70/// 1. The op is completed and the future asks the result. `into_inner` will be
71///    called by the proactor.
72/// 2. The op is completed and the future cancels it. `into_box` will be called
73///    by the proactor.
74#[derive(PartialEq, Eq, Hash)]
75pub struct Key<T: ?Sized> {
76    user_data: *mut (),
77    _p: PhantomData<Box<RawOp<T>>>,
78}
79
80impl<T: ?Sized> Unpin for Key<T> {}
81
82impl<T: OpCode + 'static> Key<T> {
83    /// Create [`RawOp`] and get the [`Key`] to it.
84    pub(crate) fn new(driver: RawFd, op: T) -> Self {
85        let header = Overlapped::new(driver);
86        let raw_op = Box::new(RawOp {
87            header,
88            cancelled: false,
89            metadata: opcode_metadata::<T>(),
90            result: PushEntry::Pending(None),
91            flags: 0,
92            op,
93        });
94        unsafe { Self::new_unchecked(Box::into_raw(raw_op) as _) }
95    }
96}
97
98impl<T: ?Sized> Key<T> {
99    /// Create a new `Key` with the given user data.
100    ///
101    /// # Safety
102    ///
103    /// Caller needs to ensure that `T` does correspond to `user_data` in driver
104    /// this `Key` is created with. In most cases, it is enough to let `T` be
105    /// `dyn OpCode`.
106    pub unsafe fn new_unchecked(user_data: usize) -> Self {
107        Self {
108            user_data: user_data as _,
109            _p: PhantomData,
110        }
111    }
112
113    /// Get the unique user-defined data.
114    pub fn user_data(&self) -> usize {
115        self.user_data as _
116    }
117
118    fn as_opaque(&self) -> &RawOp<()> {
119        // SAFETY: user_data is unique and RawOp is repr(C).
120        unsafe { &*(self.user_data as *const RawOp<()>) }
121    }
122
123    fn as_opaque_mut(&mut self) -> &mut RawOp<()> {
124        // SAFETY: see `as_opaque`.
125        unsafe { &mut *(self.user_data as *mut RawOp<()>) }
126    }
127
128    fn as_dyn_mut_ptr(&mut self) -> *mut RawOp<dyn OpCode> {
129        let user_data = self.user_data;
130        let this = self.as_opaque_mut();
131        // SAFETY: metadata from `Key::new`.
132        unsafe { opcode_dyn_mut(user_data, this.metadata) }
133    }
134
135    /// A pointer to OVERLAPPED.
136    #[cfg(windows)]
137    pub(crate) fn as_mut_ptr(&mut self) -> *mut Overlapped {
138        &mut self.as_opaque_mut().header
139    }
140
141    /// Cancel the op, decrease the ref count. The return value indicates if the
142    /// op is completed. If so, the op should be dropped because it is
143    /// useless.
144    pub(crate) fn set_cancelled(&mut self) -> bool {
145        self.as_opaque_mut().cancelled = true;
146        self.has_result()
147    }
148
149    /// Complete the op, decrease the ref count. Wake the future if a waker is
150    /// set. The return value indicates if the op is cancelled. If so, the
151    /// op should be dropped because it is useless.
152    pub(crate) fn set_result(&mut self, res: io::Result<usize>) -> bool {
153        let this = unsafe { &mut *self.as_dyn_mut_ptr() };
154        #[cfg(io_uring)]
155        if let Ok(res) = res {
156            unsafe {
157                Pin::new_unchecked(&mut this.op).set_result(res);
158            }
159        }
160        if let PushEntry::Pending(Some(w)) =
161            std::mem::replace(&mut this.result, PushEntry::Ready(res))
162        {
163            w.wake();
164        }
165        this.cancelled
166    }
167
168    pub(crate) fn set_flags(&mut self, flags: u32) {
169        self.as_opaque_mut().flags = flags;
170    }
171
172    pub(crate) fn flags(&self) -> u32 {
173        self.as_opaque().flags
174    }
175
176    /// Whether the op is completed.
177    pub(crate) fn has_result(&self) -> bool {
178        self.as_opaque().result.is_ready()
179    }
180
181    /// Set waker of the current future.
182    pub(crate) fn set_waker(&mut self, waker: Waker) {
183        if let PushEntry::Pending(w) = &mut self.as_opaque_mut().result {
184            *w = Some(waker)
185        }
186    }
187
188    /// Get the inner [`RawOp`]. It is usually used to drop the inner
189    /// immediately, without knowing about the inner `T`.
190    ///
191    /// # Safety
192    ///
193    /// Call it only when the op is cancelled and completed, which is the case
194    /// when the ref count becomes zero. See doc of [`Key::set_cancelled`]
195    /// and [`Key::set_result`].
196    pub(crate) unsafe fn into_box(mut self) -> Box<RawOp<dyn OpCode>> {
197        unsafe { Box::from_raw(self.as_dyn_mut_ptr()) }
198    }
199}
200
201impl<T> Key<T> {
202    /// Get the inner result if it is completed.
203    ///
204    /// # Safety
205    ///
206    /// Call it only when the op is completed, otherwise it is UB.
207    pub(crate) unsafe fn into_inner(self) -> BufResult<usize, T> {
208        let op = unsafe { Box::from_raw(self.user_data as *mut RawOp<T>) };
209        BufResult(unsafe { op.result.take_ready().unwrap_unchecked() }, op.op)
210    }
211}
212
213impl<T: OpCode + ?Sized> Key<T> {
214    /// Pin the inner op.
215    pub(crate) fn as_op_pin(&mut self) -> Pin<&mut dyn OpCode> {
216        // SAFETY: the inner won't be moved.
217        unsafe {
218            let this = &mut *self.as_dyn_mut_ptr();
219            Pin::new_unchecked(&mut this.op)
220        }
221    }
222
223    /// Call [`OpCode::operate`] and assume that it is not an overlapped op,
224    /// which means it never returns [`Poll::Pending`].
225    ///
226    /// [`Poll::Pending`]: std::task::Poll::Pending
227    #[cfg(windows)]
228    pub(crate) fn operate_blocking(&mut self) -> io::Result<usize> {
229        use std::task::Poll;
230
231        let optr = self.as_mut_ptr();
232        let op = self.as_op_pin();
233        let res = unsafe { op.operate(optr.cast()) };
234        match res {
235            Poll::Pending => unreachable!("this operation is not overlapped"),
236            Poll::Ready(res) => res,
237        }
238    }
239}
240
241impl<T: ?Sized> std::fmt::Debug for Key<T> {
242    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
243        write!(f, "Key({})", self.user_data())
244    }
245}
compio_driver/key.rs

compio_driver/
key.rs
