1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

#[cfg(feature = "once_cell_try")]
use std::sync::OnceLock;
use std::{io, marker::PhantomData};

use compio_buf::IntoInner;
use compio_driver::AsRawFd;
#[cfg(not(feature = "once_cell_try"))]
use once_cell::sync::OnceCell as OnceLock;

use crate::Runtime;

/// Attach a handle to the driver of current thread.
///
/// A handle can and only can attach once to one driver. However, the handle
/// itself is Send & Sync. We mark it !Send & !Sync to warn users, making them
/// ensure that they are using it in the correct thread.
#[derive(Debug, Clone)]
pub struct Attacher {
    // Make it thread safe.
    once: OnceLock<usize>,
    // Make it !Send & !Sync.
    _p: PhantomData<*mut ()>,
}

impl Attacher {
    /// Create [`Attacher`].
    pub const fn new() -> Self {
        Self {
            once: OnceLock::new(),
            _p: PhantomData,
        }
    }

    /// Attach the source. This method could be called many times, but if the
    /// action fails, the error will only return once.
    ///
    /// You should always call this method before accessing the runtime. It
    /// ensures that the current runtime is the exact runtime attached before.
    pub fn attach(&self, source: &impl AsRawFd) -> io::Result<()> {
        let r = Runtime::current();
        let inner = r.inner();
        let id = self.once.get_or_try_init(|| {
            inner.attach(source.as_raw_fd())?;
            io::Result::Ok(inner.id())
        })?;
        if id != &inner.id() {
            Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "the current runtime is not the attached runtime",
            ))
        } else {
            Ok(())
        }
    }

    /// Check if [`attach`] has been called.
    pub fn is_attached(&self) -> bool {
        self.once.get().is_some()
    }

    /// Try clone self with the cloned source. The attach state will be
    /// reserved.
    ///
    /// ## Platform specific
    /// * io-uring/polling: it will try to attach in the current thread if
    ///   needed.
    pub fn try_clone(&self, source: &impl AsRawFd) -> io::Result<Self> {
        if cfg!(windows) {
            Ok(self.clone())
        } else {
            let new_self = Self::new();
            if self.is_attached() {
                new_self.attach(source)?;
            }
            Ok(new_self)
        }
    }
}

/// Represents an attachable resource to driver.
pub trait Attachable {
    /// Attach self to the global driver.
    fn attach(&self) -> io::Result<()>;

    /// Check if [`Attachable::attach`] has been called.
    fn is_attached(&self) -> bool;
}

/// A [`Send`] wrapper for attachable resource that has not been attached. The
/// resource should be able to send to another thread before attaching.
pub struct Unattached<T: Attachable>(T);

impl<T: Attachable> Unattached<T> {
    /// Create the [`Unattached`] wrapper, or fail if the resource has already
    /// been attached.
    pub fn new(a: T) -> Result<Self, T> {
        if a.is_attached() { Err(a) } else { Ok(Self(a)) }
    }

    /// Create [`Unattached`] without checking.
    ///
    /// # Safety
    ///
    /// The caller should ensure that the resource has not been attached.
    pub unsafe fn new_unchecked(a: T) -> Self {
        Self(a)
    }
}

impl<T: Attachable> IntoInner for Unattached<T> {
    type Inner = T;

    fn into_inner(self) -> Self::Inner {
        self.0
    }
}

unsafe impl<T: Attachable> Send for Unattached<T> {}
unsafe impl<T: Attachable> Sync for Unattached<T> {}

#[macro_export]
#[doc(hidden)]
macro_rules! impl_attachable {
    ($t:ty, $inner:ident) => {
        impl $crate::Attachable for $t {
            fn attach(&self) -> ::std::io::Result<()> {
                self.$inner.attach()
            }

            fn is_attached(&self) -> bool {
                self.$inner.is_attached()
            }
        }
    };
}