Windows：IOCP 与 AFD

pub struct Selector {
    #[cfg(debug_assertions)]
    id: usize,
    pub(super) inner: Arc<SelectorInner>,
}

pub struct SelectorInner {
    pub(super) cp: Arc<CompletionPort>,
    update_queue: Mutex<VecDeque<Pin<Arc<Mutex<SockState>>>>>,
    afd_group: AfdGroup,
    is_polling: AtomicBool,
}

    const POLL_GROUP_MAX_GROUP_SIZE: usize = 32;

    impl AfdGroup {
        pub fn acquire(&self) -> io::Result<Arc<Afd>> {
            let mut afd_group = self.afd_group.lock().unwrap();
            if afd_group.len() == 0 {
                self._alloc_afd_group(&mut afd_group)?;
            } else {
                // + 1 reference in Vec
                if Arc::strong_count(afd_group.last().unwrap()) > POLL_GROUP_MAX_GROUP_SIZE  {
                    self._alloc_afd_group(&mut afd_group)?;
                }
            }
            // ...

                let fd = File::from_raw_handle(afd_helper_handle as RawHandle);
                // Increment by 2 to reserve space for other types of handles.
                // Non-AFD types (currently only NamedPipe), use odd numbered
                // tokens. This allows the selector to differentiate between them
                // and dispatch events accordingly.
                let token = NEXT_TOKEN.fetch_add(2, Ordering::Relaxed) + 2;
                let afd = Afd { fd };
                cp.add_handle(token, &afd.fd)?;
                match SetFileCompletionNotificationModes(
                    afd_helper_handle,
                    FILE_SKIP_SET_EVENT_ON_HANDLE as u8 // This is just 2, so fits in u8
                ) {

#[repr(C)]
pub struct AfdPollInfo {
    pub timeout: i64,
    // Can have only value 1.
    pub number_of_handles: u32,
    pub exclusive: u32,
    pub handles: [AfdPollHandleInfo; 1],
}

        let info_ptr = info as *mut _ as *mut c_void;
        (*iosb).Anonymous.Status = STATUS_PENDING;
        let status = NtDeviceIoControlFile(
            self.fd.as_raw_handle() as HANDLE,
            std::ptr::null_mut(),
            None,
            overlapped,
            iosb,
            IOCTL_AFD_POLL,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
        );
        match status {
            STATUS_SUCCESS => Ok(true),
            STATUS_PENDING => Ok(false),
            _ => Err(io::Error::from_raw_os_error(
                RtlNtStatusToDosError(status) as i32
            )),
        }

    pub fn select(&self, events: &mut Events, timeout: Option<Duration>) -> io::Result<()> {
        events.clear();

        if timeout.is_none() {
            loop {
                let len = self.select2(&mut events.statuses, &mut events.events, None)?;
                if len == 0 {
                    continue;
                }
                break Ok(());
            }
        } else {
            self.select2(&mut events.statuses, &mut events.events, timeout)?;
            Ok(())
        }
    }

    pub fn select2(
        &self,
        statuses: &mut [CompletionStatus],
        events: &mut Vec<Event>,
        timeout: Option<Duration>,
    ) -> io::Result<usize> {
        assert!(!self.is_polling.swap(true, Ordering::AcqRel));

        unsafe { self.update_sockets_events() }?;

        let result = self.cp.get_many(statuses, timeout);

        self.is_polling.store(false, Ordering::Relaxed);

        match result {
            Ok(iocp_events) => Ok(unsafe { self.feed_events(events, iocp_events) }),
            Err(ref e) if e.raw_os_error() == Some(WAIT_TIMEOUT as i32) => Ok(0),
            Err(e) => Err(e),
        }
    }

        unsafe fn update_sockets_events_if_polling(&self) -> io::Result<()> {
            if self.is_polling.load(Ordering::Acquire) {
                self.update_sockets_events()
            } else {
                Ok(())
            }
        }

pub struct SockState {
    iosb: IoStatusBlock,
    poll_info: AfdPollInfo,
    afd: Arc<Afd>,

    base_socket: RawSocket,

    user_evts: u32,
    pending_evts: u32,

    user_data: u64,

    poll_status: SockPollStatus,
    delete_pending: bool,

    // last raw os error
    error: Option<i32>,

    _pinned: PhantomPinned,
}

        if let SockPollStatus::Pending = self.poll_status {
            if (self.user_evts & afd::KNOWN_EVENTS & !self.pending_evts) == 0 {
                /* All the events the user is interested in are already being monitored by
                 * the pending poll operation. It might spuriously complete because of an
                 * event that we're no longer interested in; when that happens we'll submit
                 * a new poll operation with the updated event mask. */
            } else {
                /* A poll operation is already pending, but it's not monitoring for all the
                 * events that the user is interested in. Therefore, cancel the pending
                 * poll operation; when we receive it's completion package, a new poll
                 * operation will be submitted with the correct event mask. */
                if let Err(e) = self.cancel() {
                    self.error = e.raw_os_error();
                    return Err(e);
                }
                return Ok(());
            }
        } else if let SockPollStatus::Cancelled = self.poll_status {
            /* The poll operation has already been cancelled, we're still waiting for
             * it to return. For now, there's nothing that needs to be done. */
        } else if let SockPollStatus::Idle = self.poll_status {
            /* No poll operation is pending; start one. */
            self.poll_info.exclusive = 0;
            self.poll_info.number_of_handles = 1;
            self.poll_info.timeout = i64::MAX;
            self.poll_info.handles[0].handle = self.base_socket as HANDLE;
            self.poll_info.handles[0].status = 0;
            self.poll_info.handles[0].events = self.user_evts | afd::POLL_LOCAL_CLOSE;

            // Increase the ref count as the memory will be used by the kernel.
            let overlapped_ptr = into_overlapped(self_arc.clone());
            // ...
            self.poll_status = SockPollStatus::Pending;
            self.pending_evts = self.user_evts;
        }

        for iocp_event in iocp_events.iter() {
            if iocp_event.overlapped().is_null() {
                events.push(Event::from_completion_status(iocp_event));
                n += 1;
                continue;
            } else if iocp_event.token() % 2 == 1 {
                // Handle is a named pipe. This could be extended to be any non-AFD event.
                let callback = (*(iocp_event.overlapped() as *mut super::Overlapped)).callback;

                let len = events.len();
                callback(iocp_event.entry(), Some(events));
                n += events.len() - len;
                continue;
            }

            let sock_state = from_overlapped(iocp_event.overlapped());
            let mut sock_guard = sock_state.lock().unwrap();
            if let Some(e) = sock_guard.feed_event() {
                events.push(e);
                n += 1;
            }

            if !sock_guard.is_pending_deletion() {
                update_queue.push_back(sock_state.clone());
            }
        }

    fn feed_event(&mut self) -> Option<Event> {
        self.poll_status = SockPollStatus::Idle;
        self.pending_evts = 0;

        let mut afd_events = 0;
        // We use the status info in IO_STATUS_BLOCK to determine the socket poll status. It is unsafe to use a pointer of IO_STATUS_BLOCK.
        unsafe {
            if self.delete_pending {
                return None;
            } else if self.iosb.Anonymous.Status == STATUS_CANCELLED {
                /* The poll request was cancelled by CancelIoEx. */
            } else if self.iosb.Anonymous.Status < 0 {
                /* The overlapped request itself failed in an unexpected way. */
                afd_events = afd::POLL_CONNECT_FAIL;
            } else if self.poll_info.number_of_handles < 1 {
                /* This poll operation succeeded but didn't report any socket events. */
            } else if self.poll_info.handles[0].events & afd::POLL_LOCAL_CLOSE != 0 {
                /* The poll operation reported that the socket was closed. */
                self.mark_delete();
                return None;
            } else {
                afd_events = self.poll_info.handles[0].events;
            }
        }

        afd_events &= self.user_evts;

        if afd_events == 0 {
            return None;
        }

        // In mio, we have to simulate Edge-triggered behavior to match API usage.
        // The strategy here is to intercept all read/write from user that could cause WouldBlock usage,
        // then reregister the socket to reset the interests.
        self.user_evts &= !afd_events;

        Some(Event {
            data: self.user_data,
            flags: afd_events,
        })
    }

        pub fn do_io<T, F, R>(&self, f: F, io: &T) -> io::Result<R>
        where
            F: FnOnce(&T) -> io::Result<R>,
        {
            let result = f(io);
            if let Err(ref e) = result {
                if e.kind() == io::ErrorKind::WouldBlock {
                    self.inner.as_ref().map_or(Ok(()), |state| {
                        state
                            .selector
                            .reregister(state.sock_state.clone(), state.token, state.interests)
                    })?;
                }
            }
            result
        }

pub struct Waker {
    token: Token,
    port: Arc<CompletionPort>,
}

impl Waker {
    pub fn new(selector: &Selector, token: Token) -> io::Result<Waker> {
        Ok(Waker {
            token,
            port: selector.clone_port(),
        })
    }

    pub fn wake(&self) -> io::Result<()> {
        let mut ev = Event::new(self.token);
        ev.set_readable();

        self.port.post(ev.to_completion_status())
    }
}

/// Note, IOCP is a *completion* based model whereas mio is a *readiness* based
/// model. To bridge this, `NamedPipe` performs internal buffering. Writes are
/// written to an internal buffer and the buffer is submitted to IOCP. IOCP
/// reads are submitted using internal buffers and `NamedPipe::read` reads from
/// this internal buffer.

#[repr(C)]
struct Inner {
    // NOTE: careful modifying the order of these three fields, the `ptr_from_*`
    // methods depend on the layout!
    connect: Overlapped,
    read: Overlapped,
    write: Overlapped,
    event: Overlapped,
    // END NOTE.
    handle: Handle,
    connecting: AtomicBool,
    io: Mutex<Io>,
    pool: Mutex<BufferPool>,
}

        // If `connect_overlapped` does not complete immediately then we need
        // to pass a reference to `self.inner` to the completion handler.
        // We need to increase the reference count here because `connect_done`
        // could be called as soon as we call `connect_overlapped`.
        let inner_for_completion = self.inner.clone();

        if res == 0 {
            let err = io::Error::last_os_error();
            if err.raw_os_error() != Some(ERROR_IO_PENDING as i32) {
                return Err(err);
            }
        }

        let mut bytes = 0;
        let res = GetOverlappedResult(self.handle.raw(), overlapped, &mut bytes, 0);
        if res == 0 {
            let err = io::Error::last_os_error();
            if err.raw_os_error() == Some(ERROR_IO_INCOMPLETE as i32) {
                Ok(None)
            } else {
                Err(err)
            }
        } else {
            Ok(Some(bytes as usize))
        }

        if res == 0 {
            let err = io::Error::last_os_error();
            // `ERROR_MORE_DATA` is only a warning (the message didn't fit), the
            // read still completes through its packet.
            if err.raw_os_error() != Some(ERROR_IO_PENDING as i32)
                && err.raw_os_error() != Some(ERROR_MORE_DATA as i32)
            {
                return Err(err);
            }
            // Submitted: a completion packet will be queued, so this must not return Err.
            return Ok(None);
        }

        // Completed, but the completion still delivers the outcome, so this must not return Err.
        let mut bytes = 0;
        if GetOverlappedResult(self.handle.raw(), overlapped, &mut bytes, 0) == 0 {
            Ok(None)
        } else {
            Ok(Some(bytes as usize))
        }

    let flags = WSA_FLAG_OVERLAPPED | WSA_FLAG_NO_HANDLE_INHERIT;
    let socket = syscall!(
        WSASocketW(domain as i32, socket_type, 0, std::ptr::null(), 0, flags),
        PartialEq::eq,
        INVALID_SOCKET
    )?;

use std::io::{self, Read};
use std::thread;

use mio::net::{TcpListener, TcpStream};
use mio::{Events, Interest, Poll, Token, Waker};

const LISTENER: Token = Token(0);
const CLIENT: Token = Token(1);
const WAKER: Token = Token(2);

fn main() -> io::Result<()> {
  let mut poll = Poll::new()?;
  let mut events = Events::with_capacity(16);

  let mut listener = TcpListener::bind("127.0.0.1:0".parse().unwrap())?;
  poll.registry().register(&mut listener, LISTENER, Interest::READABLE)?;
  let mut client = TcpStream::connect(listener.local_addr()?)?;
  poll
    .registry()
    .register(&mut client, CLIENT, Interest::READABLE | Interest::WRITABLE)?;

  let waker = Waker::new(poll.registry(), WAKER)?;
  let handle = thread::spawn(move || waker.wake());

  #[cfg(windows)]
  {
    // Windows-only type from os-ext: the one mio type that buffers internally.
    let mut pipe = mio::windows::NamedPipe::new(r"\\.\pipe\mio-windows-check")?;
    poll.registry().register(&mut pipe, Token(3), Interest::READABLE)?;
  }

  poll.poll(&mut events, None)?;
  for event in events.iter() {
    println!("{:?} r={} w={}", event.token(), event.is_readable(), event.is_writable());
  }
  let mut buf = [0u8; 16];
  match client.read(&mut buf) {
    Err(e) if e.kind() == io::ErrorKind::WouldBlock => println!("read: WouldBlock"),
    other => println!("read: {:?}", other.map(|n| n)),
  }
  handle.join().unwrap()?;
  Ok(())
}

    Checking exp-mio-windows v0.1.0 (exp-mio-windows)
    Finished `release` profile [optimized] target(s) in 0.07s

   Compiling exp-mio-windows v0.1.0 (exp-mio-windows)
error: linker `link.exe` not found
  |
  = note: No such file or directory (os error 2)

note: the msvc targets depend on the msvc linker but `link.exe` was not found

note: please ensure that Visual Studio 2017 or later, or Build Tools for Visual Studio were installed with the Visual C++ option

note: VS Code is a different product, and is not sufficient

error: could not compile `exp-mio-windows` (bin "wincheck") due to 1 previous error

    Checking exp-mio-windows v0.1.0 (exp-mio-windows)
error[E0432]: unresolved import `mio::unix`
  --> src/bin/sourcefd.rs:1:10
   |
 1 | use mio::unix::SourceFd;
   |          ^^^^ could not find `unix` in `mio`
   |
note: found an item that was configured out
  --> ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/mio-1.2.4/src/lib.rs:77:9
   |
75 | #[cfg(all(unix, feature = "os-ext"))]
   |           ---- the item is gated here
76 | #[cfg_attr(docsrs, doc(cfg(all(unix, feature = "os-ext"))))]
77 | pub mod unix {
   |         ^^^^

For more information about this error, try `rustc --explain E0432`.
error: could not compile `exp-mio-windows` (bin "sourcefd") due to 1 previous error