Skip to main content

async_selector/selector/
id.rs

1use std::{
2    cmp::Ordering,
3    fmt,
4    hash::{Hash, Hasher},
5    mem::ManuallyDrop,
6    sync::{Arc, Weak},
7};
8
9use crate::{
10    mpsc::{Queue, WeakSender},
11    task::Task,
12};
13
14/// A unique id of a task stored in a [`Selector`](crate::selector::Selector).
15///
16/// See [example](https://github.com/Razz4780/async-selector/blob/main/examples/map.rs)
17/// of how it can be leverage to use the selector like a map.
18///
19/// Mind that this keeping this id alive prevents the selector
20/// from deallocating memory used to store the task.
21pub struct Id {
22    /// Type erased pointer to the selector's queue.
23    sender_ptr: *const (),
24    /// Type erased pointer to the task.
25    task_ptr: *const (),
26    /// Virtual methods table that allows us to clone and drop this id,
27    /// even though the types were erased.
28    vtable: &'static IdVTable,
29}
30
31impl Id {
32    pub(super) fn new<P>(task: Weak<Task<P>>, sender: WeakSender<Task<P>>) -> Self {
33        let sender_ptr = sender.into_raw().cast();
34        let task_ptr = task.into_raw().cast();
35        Self {
36            sender_ptr,
37            task_ptr,
38            vtable: &Task::<P>::ID_VTABLE,
39        }
40    }
41
42    pub(super) fn sender_ptr(&self) -> *const () {
43        self.sender_ptr
44    }
45
46    /// Recovers the task, if it's still alive.
47    ///
48    /// # Safety
49    ///
50    /// Caller must ensure that the type of the task matches.
51    pub(super) unsafe fn task<P>(&self) -> Option<Arc<Task<P>>> {
52        let weak = unsafe { Weak::from_raw(self.task_ptr.cast::<Task<P>>()) };
53        let strong = weak.upgrade();
54        let _ = ManuallyDrop::new(weak);
55        strong
56    }
57}
58
59impl Clone for Id {
60    fn clone(&self) -> Self {
61        unsafe { (self.vtable.clone_raw)(self.sender_ptr, self.task_ptr) }
62    }
63}
64
65impl Drop for Id {
66    fn drop(&mut self) {
67        unsafe { (self.vtable.drop_raw)(self.sender_ptr, self.task_ptr) }
68    }
69}
70
71impl fmt::Debug for Id {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        f.debug_struct("Id")
74            .field("task", &self.task_ptr)
75            .field("selector", &self.sender_ptr)
76            .finish_non_exhaustive()
77    }
78}
79
80impl PartialEq for Id {
81    fn eq(&self, other: &Self) -> bool {
82        std::ptr::eq(self.task_ptr, other.task_ptr)
83    }
84}
85
86impl Eq for Id {}
87
88impl Hash for Id {
89    fn hash<H: Hasher>(&self, state: &mut H) {
90        std::ptr::hash(self.task_ptr, state);
91    }
92}
93
94impl PartialOrd for Id {
95    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
96        Some(self.cmp(other))
97    }
98}
99
100impl Ord for Id {
101    fn cmp(&self, other: &Self) -> Ordering {
102        self.task_ptr.cmp(&other.task_ptr)
103    }
104}
105
106unsafe impl Send for Id {}
107unsafe impl Sync for Id {}
108
109impl<P> Task<P> {
110    const ID_VTABLE: IdVTable = IdVTable {
111        clone_raw: Self::clone_raw,
112        drop_raw: Self::drop_raw,
113    };
114
115    unsafe fn clone_raw(sender_ptr: *const (), task_ptr: *const ()) -> Id {
116        let sender = unsafe { WeakSender::from_raw(sender_ptr.cast::<Queue<Self>>()) };
117        let cloned = sender.clone();
118        let _ = ManuallyDrop::new(sender);
119        let _ = ManuallyDrop::new(cloned);
120
121        let weak = unsafe { Weak::from_raw(task_ptr.cast::<Self>()) };
122        let cloned = weak.clone();
123        let _ = ManuallyDrop::new(weak);
124        let _ = ManuallyDrop::new(cloned);
125
126        Id {
127            sender_ptr,
128            task_ptr,
129            vtable: &Self::ID_VTABLE,
130        }
131    }
132
133    unsafe fn drop_raw(sender_ptr: *const (), task_ptr: *const ()) {
134        let sender = unsafe { WeakSender::from_raw(sender_ptr.cast::<Queue<Self>>()) };
135        drop(sender);
136        let task = unsafe { Weak::from_raw(task_ptr.cast::<Self>()) };
137        drop(task);
138    }
139}
140
141/// Allows for cloning and dropping type-erased [`Id`]s.
142struct IdVTable {
143    clone_raw: unsafe fn(*const (), *const ()) -> Id,
144    drop_raw: unsafe fn(*const (), *const ()),
145}