<commit_before>#![allow(missing_docs, nonstandard_style)]

use crate::ffi::{OsStr, OsString};
use crate::io::ErrorKind;
use crate::os::windows::ffi::{OsStrExt, OsStringExt};
use crate::path::PathBuf;
use crate::ptr;
use crate::time::Duration;

pub use self::rand::hashmap_random_keys;
pub use libc::strlen;

#[macro_use]
pub mod compat;

pub mod alloc;
pub mod args;
pub mod c;
pub mod cmath;
pub mod condvar;
pub mod env;
pub mod ext;
pub mod fs;
pub mod handle;
pub mod io;
pub mod memchr;
pub mod mutex;
pub mod net;
pub mod os;
pub mod os_str;
pub mod path;
pub mod pipe;
pub mod process;
pub mod rand;
pub mod rwlock;
pub mod thread;
pub mod thread_local_dtor;
pub mod thread_local_key;
pub mod time;
cfg_if::cfg_if! {
    if #[cfg(not(target_vendor = "uwp"))] {
        pub mod stdio;
        pub mod stack_overflow;
    } else {
        pub mod stdio_uwp;
        pub mod stack_overflow_uwp;
        pub use self::stdio_uwp as stdio;
        pub use self::stack_overflow_uwp as stack_overflow;
    }
}

#[cfg(not(test))]
pub fn init() {}

pub fn decode_error_kind(errno: i32) -> ErrorKind {
    match errno as c::DWORD {
        c::ERROR_ACCESS_DENIED => return ErrorKind::PermissionDenied,
        c::ERROR_ALREADY_EXISTS => return ErrorKind::AlreadyExists,
        c::ERROR_FILE_EXISTS => return ErrorKind::AlreadyExists,
        c::ERROR_BROKEN_PIPE => return ErrorKind::BrokenPipe,
        c::ERROR_FILE_NOT_FOUND => return ErrorKind::NotFound,
        c::ERROR_PATH_NOT_FOUND => return ErrorKind::NotFound,
        c::ERROR_NO_DATA => return ErrorKind::BrokenPipe,
        c::ERROR_INVALID_PARAMETER => return ErrorKind::InvalidInput,
        c::ERROR_SEM_TIMEOUT
        | c::WAIT_TIMEOUT
        | c::ERROR_DRIVER_CANCEL_TIMEOUT
        | c::ERROR_OPERATION_ABORTED
        | c::ERROR_SERVICE_REQUEST_TIMEOUT
        | c::ERROR_COUNTER_TIMEOUT
        | c::ERROR_TIMEOUT
        | c::ERROR_RESOURCE_CALL_TIMED_OUT
        | c::ERROR_CTX_MODEM_RESPONSE_TIMEOUT
        | c::ERROR_CTX_CLIENT_QUERY_TIMEOUT
        | c::FRS_ERR_SYSVOL_POPULATE_TIMEOUT
        | c::ERROR_DS_TIMELIMIT_EXCEEDED
        | c::DNS_ERROR_RECORD_TIMED_OUT
        | c::ERROR_IPSEC_IKE_TIMED_OUT
        | c::ERROR_RUNLEVEL_SWITCH_TIMEOUT
        | c::ERROR_RUNLEVEL_SWITCH_AGENT_TIMEOUT => return ErrorKind::TimedOut,
        _ => {}
    }

    match errno {
        c::WSAEACCES => ErrorKind::PermissionDenied,
        c::WSAEADDRINUSE => ErrorKind::AddrInUse,
        c::WSAEADDRNOTAVAIL => ErrorKind::AddrNotAvailable,
        c::WSAECONNABORTED => ErrorKind::ConnectionAborted,
        c::WSAECONNREFUSED => ErrorKind::ConnectionRefused,
        c::WSAECONNRESET => ErrorKind::ConnectionReset,
        c::WSAEINVAL => ErrorKind::InvalidInput,
        c::WSAENOTCONN => ErrorKind::NotConnected,
        c::WSAEWOULDBLOCK => ErrorKind::WouldBlock,
        c::WSAETIMEDOUT => ErrorKind::TimedOut,

        _ => ErrorKind::Other,
    }
}

pub fn unrolled_find_u16s(needle: u16, haystack: &[u16]) -> Option<usize> {
    let ptr = haystack.as_ptr();
    let mut len = haystack.len();
    let mut start = &haystack[..];

    // For performance reasons unfold the loop eight times.
    while len >= 8 {
        if start[0] == needle {
            return Some((start.as_ptr() as usize - ptr as usize) / 2);
        }
        if start[1] == needle {
            return Some((start[1..].as_ptr() as usize - ptr as usize) / 2);
        }
        if start[2] == needle {
            return Some((start[2..].as_ptr() as usize - ptr as usize) / 2);
        }
        if start[3] == needle {
            return Some((start[3..].as_ptr() as usize - ptr as usize) / 2);
        }
        if start[4] == needle {
            return Some((start[4..].as_ptr() as usize - ptr as usize) / 2);
        }
        if start[5] == needle {
            return Some((start[5..].as_ptr() as usize - ptr as usize) / 2);
        }
        if start[6] == needle {
            return Some((start[6..].as_ptr() as usize - ptr as usize) / 2);
        }
        if start[7] == needle {
            return Some((start[7..].as_ptr() as usize - ptr as usize) / 2);
        }

        start = &start[8..];
        len -= 8;
    }

    for (i, c) in start.iter().enumerate() {
        if *c == needle {
            return Some((start.as_ptr() as usize - ptr as usize) / 2 + i);
        }
    }
    None
}

pub fn to_u16s<S: AsRef<OsStr>>(s: S) -> crate::io::Result<Vec<u16>> {
    fn inner(s: &OsStr) -> crate::io::Result<Vec<u16>> {
        let mut maybe_result: Vec<u16> = s.encode_wide().collect();
        if unrolled_find_u16s(0, &maybe_result).is_some() {
            return Err(crate::io::Error::new(
                ErrorKind::InvalidInput,
                "strings passed to WinAPI cannot contain NULs",
            ));
        }
        maybe_result.push(0);
        Ok(maybe_result)
    }
    inner(s.as_ref())
}

// Many Windows APIs follow a pattern of where we hand a buffer and then they
// will report back to us how large the buffer should be or how many bytes
// currently reside in the buffer. This function is an abstraction over these
// functions by making them easier to call.
//
// The first callback, `f1`, is yielded a (pointer, len) pair which can be
// passed to a syscall. The `ptr` is valid for `len` items (u16 in this case).
// The closure is expected to return what the syscall returns which will be
// interpreted by this function to determine if the syscall needs to be invoked
// again (with more buffer space).
//
// Once the syscall has completed (errors bail out early) the second closure is
// yielded the data which has been read from the syscall. The return value
// from this closure is then the return value of the function.
fn fill_utf16_buf<F1, F2, T>(mut f1: F1, f2: F2) -> crate::io::Result<T>
where
    F1: FnMut(*mut u16, c::DWORD) -> c::DWORD,
    F2: FnOnce(&[u16]) -> T,
{
    // Start off with a stack buf but then spill over to the heap if we end up
    // needing more space.
    let mut stack_buf = [0u16; 512];
    let mut heap_buf = Vec::new();
    unsafe {
        let mut n = stack_buf.len();
        loop {
            let buf = if n <= stack_buf.len() {
                &mut stack_buf[..]
            } else {
                let extra = n - heap_buf.len();
                heap_buf.reserve(extra);
                heap_buf.set_len(n);
                &mut heap_buf[..]
            };

            // This function is typically called on windows API functions which
            // will return the correct length of the string, but these functions
            // also return the `0` on error. In some cases, however, the
            // returned "correct length" may actually be 0!
            //
            // To handle this case we call `SetLastError` to reset it to 0 and
            // then check it again if we get the "0 error value". If the "last
            // error" is still 0 then we interpret it as a 0 length buffer and
            // not an actual error.
            c::SetLastError(0);
            let k = match f1(buf.as_mut_ptr(), n as c::DWORD) {
                0 if c::GetLastError() == 0 => 0,
                0 => return Err(crate::io::Error::last_os_error()),
                n => n,
            } as usize;
            if k == n && c::GetLastError() == c::ERROR_INSUFFICIENT_BUFFER {
                n *= 2;
            } else if k >= n {
                n = k;
            } else {
                return Ok(f2(&buf[..k]));
            }
        }
    }
}

fn os2path(s: &[u16]) -> PathBuf {
    PathBuf::from(OsString::from_wide(s))
}

#[allow(dead_code)] // Only used in backtrace::gnu::get_executable_filename()
fn wide_char_to_multi_byte(
    code_page: u32,
    flags: u32,
    s: &[u16],
    no_default_char: bool,
) -> crate::io::Result<Vec<i8>> {
    unsafe {
        let mut size = c::WideCharToMultiByte(
            code_page,
            flags,
            s.as_ptr(),
            s.len() as i32,
            ptr::null_mut(),
            0,
            ptr::null(),
            ptr::null_mut(),
        );
        if size == 0 {
            return Err(crate::io::Error::last_os_error());
        }

        let mut buf = Vec::with_capacity(size as usize);
        buf.set_len(size as usize);

        let mut used_default_char = c::FALSE;
        size = c::WideCharToMultiByte(
            code_page,
            flags,
            s.as_ptr(),
            s.len() as i32,
            buf.as_mut_ptr(),
            buf.len() as i32,
            ptr::null(),
            if no_default_char { &mut used_default_char } else { ptr::null_mut() },
        );
        if size == 0 {
            return Err(crate::io::Error::last_os_error());
        }
        if no_default_char && used_default_char == c::TRUE {
            return Err(crate::io::Error::new(
                crate::io::ErrorKind::InvalidData,
                "string cannot be converted to requested code page",
            ));
        }

        buf.set_len(size as usize);

        Ok(buf)
    }
}

pub fn truncate_utf16_at_nul(v: &[u16]) -> &[u16] {
    match unrolled_find_u16s(0, v) {
        // don't include the 0
        Some(i) => &v[..i],
        None => v,
    }
}

pub trait IsZero {
    fn is_zero(&self) -> bool;
}

macro_rules! impl_is_zero {
    ($($t:ident)*) => ($(impl IsZero for $t {
        fn is_zero(&self) -> bool {
            *self == 0
        }
    })*)
}

impl_is_zero! { i8 i16 i32 i64 isize u8 u16 u32 u64 usize }

pub fn cvt<I: IsZero>(i: I) -> crate::io::Result<I> {
    if i.is_zero() { Err(crate::io::Error::last_os_error()) } else { Ok(i) }
}

pub fn dur2timeout(dur: Duration) -> c::DWORD {
    // Note that a duration is a (u64, u32) (seconds, nanoseconds) pair, and the
    // timeouts in windows APIs are typically u32 milliseconds. To translate, we
    // have two pieces to take care of:
    //
    // * Nanosecond precision is rounded up
    // * Greater than u32::MAX milliseconds (50 days) is rounded up to INFINITE
    //   (never time out).
    dur.as_secs()
        .checked_mul(1000)
        .and_then(|ms| ms.checked_add((dur.subsec_nanos() as u64) / 1_000_000))
        .and_then(|ms| ms.checked_add(if dur.subsec_nanos() % 1_000_000 > 0 { 1 } else { 0 }))
        .map(|ms| if ms > <c::DWORD>::MAX as u64 { c::INFINITE } else { ms as c::DWORD })
        .unwrap_or(c::INFINITE)
}

// On Windows, use the processor-specific __fastfail mechanism.  In Windows 8
// and later, this will terminate the process immediately without running any
// in-process exception handlers.  In earlier versions of Windows, this
// sequence of instructions will be treated as an access violation,
// terminating the process but without necessarily bypassing all exception
// handlers.
//
// https://docs.microsoft.com/en-us/cpp/intrinsics/fastfail
#[allow(unreachable_code)]
pub fn abort_internal() -> ! {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    unsafe {
        llvm_asm!("int $$0x29" :: "{ecx}"(7) ::: volatile); // 7 is FAST_FAIL_FATAL_APP_EXIT
        crate::intrinsics::unreachable();
    }
    crate::intrinsics::abort();
}
<commit_msg>Auto merge of #74289 - lzutao:unroll, r=LukasKalbertodt<commit_after>#![allow(missing_docs, nonstandard_style)]

use crate::ffi::{OsStr, OsString};
use crate::io::ErrorKind;
use crate::os::windows::ffi::{OsStrExt, OsStringExt};
use crate::path::PathBuf;
use crate::ptr;
use crate::time::Duration;

pub use self::rand::hashmap_random_keys;
pub use libc::strlen;

#[macro_use]
pub mod compat;

pub mod alloc;
pub mod args;
pub mod c;
pub mod cmath;
pub mod condvar;
pub mod env;
pub mod ext;
pub mod fs;
pub mod handle;
pub mod io;
pub mod memchr;
pub mod mutex;
pub mod net;
pub mod os;
pub mod os_str;
pub mod path;
pub mod pipe;
pub mod process;
pub mod rand;
pub mod rwlock;
pub mod thread;
pub mod thread_local_dtor;
pub mod thread_local_key;
pub mod time;
cfg_if::cfg_if! {
    if #[cfg(not(target_vendor = "uwp"))] {
        pub mod stdio;
        pub mod stack_overflow;
    } else {
        pub mod stdio_uwp;
        pub mod stack_overflow_uwp;
        pub use self::stdio_uwp as stdio;
        pub use self::stack_overflow_uwp as stack_overflow;
    }
}

#[cfg(not(test))]
pub fn init() {}

pub fn decode_error_kind(errno: i32) -> ErrorKind {
    match errno as c::DWORD {
        c::ERROR_ACCESS_DENIED => return ErrorKind::PermissionDenied,
        c::ERROR_ALREADY_EXISTS => return ErrorKind::AlreadyExists,
        c::ERROR_FILE_EXISTS => return ErrorKind::AlreadyExists,
        c::ERROR_BROKEN_PIPE => return ErrorKind::BrokenPipe,
        c::ERROR_FILE_NOT_FOUND => return ErrorKind::NotFound,
        c::ERROR_PATH_NOT_FOUND => return ErrorKind::NotFound,
        c::ERROR_NO_DATA => return ErrorKind::BrokenPipe,
        c::ERROR_INVALID_PARAMETER => return ErrorKind::InvalidInput,
        c::ERROR_SEM_TIMEOUT
        | c::WAIT_TIMEOUT
        | c::ERROR_DRIVER_CANCEL_TIMEOUT
        | c::ERROR_OPERATION_ABORTED
        | c::ERROR_SERVICE_REQUEST_TIMEOUT
        | c::ERROR_COUNTER_TIMEOUT
        | c::ERROR_TIMEOUT
        | c::ERROR_RESOURCE_CALL_TIMED_OUT
        | c::ERROR_CTX_MODEM_RESPONSE_TIMEOUT
        | c::ERROR_CTX_CLIENT_QUERY_TIMEOUT
        | c::FRS_ERR_SYSVOL_POPULATE_TIMEOUT
        | c::ERROR_DS_TIMELIMIT_EXCEEDED
        | c::DNS_ERROR_RECORD_TIMED_OUT
        | c::ERROR_IPSEC_IKE_TIMED_OUT
        | c::ERROR_RUNLEVEL_SWITCH_TIMEOUT
        | c::ERROR_RUNLEVEL_SWITCH_AGENT_TIMEOUT => return ErrorKind::TimedOut,
        _ => {}
    }

    match errno {
        c::WSAEACCES => ErrorKind::PermissionDenied,
        c::WSAEADDRINUSE => ErrorKind::AddrInUse,
        c::WSAEADDRNOTAVAIL => ErrorKind::AddrNotAvailable,
        c::WSAECONNABORTED => ErrorKind::ConnectionAborted,
        c::WSAECONNREFUSED => ErrorKind::ConnectionRefused,
        c::WSAECONNRESET => ErrorKind::ConnectionReset,
        c::WSAEINVAL => ErrorKind::InvalidInput,
        c::WSAENOTCONN => ErrorKind::NotConnected,
        c::WSAEWOULDBLOCK => ErrorKind::WouldBlock,
        c::WSAETIMEDOUT => ErrorKind::TimedOut,

        _ => ErrorKind::Other,
    }
}

pub fn unrolled_find_u16s(needle: u16, haystack: &[u16]) -> Option<usize> {
    let ptr = haystack.as_ptr();
    let mut start = &haystack[..];

    // For performance reasons unfold the loop eight times.
    while start.len() >= 8 {
        macro_rules! if_return {
            ($($n:literal,)+) => {
                $(
                    if start[$n] == needle {
                        return Some((&start[$n] as *const u16 as usize - ptr as usize) / 2);
                    }
                )+
            }
        }

        if_return!(0, 1, 2, 3, 4, 5, 6, 7,);

        start = &start[8..];
    }

    for c in start {
        if *c == needle {
            return Some((c as *const u16 as usize - ptr as usize) / 2);
        }
    }
    None
}

pub fn to_u16s<S: AsRef<OsStr>>(s: S) -> crate::io::Result<Vec<u16>> {
    fn inner(s: &OsStr) -> crate::io::Result<Vec<u16>> {
        let mut maybe_result: Vec<u16> = s.encode_wide().collect();
        if unrolled_find_u16s(0, &maybe_result).is_some() {
            return Err(crate::io::Error::new(
                ErrorKind::InvalidInput,
                "strings passed to WinAPI cannot contain NULs",
            ));
        }
        maybe_result.push(0);
        Ok(maybe_result)
    }
    inner(s.as_ref())
}

// Many Windows APIs follow a pattern of where we hand a buffer and then they
// will report back to us how large the buffer should be or how many bytes
// currently reside in the buffer. This function is an abstraction over these
// functions by making them easier to call.
//
// The first callback, `f1`, is yielded a (pointer, len) pair which can be
// passed to a syscall. The `ptr` is valid for `len` items (u16 in this case).
// The closure is expected to return what the syscall returns which will be
// interpreted by this function to determine if the syscall needs to be invoked
// again (with more buffer space).
//
// Once the syscall has completed (errors bail out early) the second closure is
// yielded the data which has been read from the syscall. The return value
// from this closure is then the return value of the function.
fn fill_utf16_buf<F1, F2, T>(mut f1: F1, f2: F2) -> crate::io::Result<T>
where
    F1: FnMut(*mut u16, c::DWORD) -> c::DWORD,
    F2: FnOnce(&[u16]) -> T,
{
    // Start off with a stack buf but then spill over to the heap if we end up
    // needing more space.
    let mut stack_buf = [0u16; 512];
    let mut heap_buf = Vec::new();
    unsafe {
        let mut n = stack_buf.len();
        loop {
            let buf = if n <= stack_buf.len() {
                &mut stack_buf[..]
            } else {
                let extra = n - heap_buf.len();
                heap_buf.reserve(extra);
                heap_buf.set_len(n);
                &mut heap_buf[..]
            };

            // This function is typically called on windows API functions which
            // will return the correct length of the string, but these functions
            // also return the `0` on error. In some cases, however, the
            // returned "correct length" may actually be 0!
            //
            // To handle this case we call `SetLastError` to reset it to 0 and
            // then check it again if we get the "0 error value". If the "last
            // error" is still 0 then we interpret it as a 0 length buffer and
            // not an actual error.
            c::SetLastError(0);
            let k = match f1(buf.as_mut_ptr(), n as c::DWORD) {
                0 if c::GetLastError() == 0 => 0,
                0 => return Err(crate::io::Error::last_os_error()),
                n => n,
            } as usize;
            if k == n && c::GetLastError() == c::ERROR_INSUFFICIENT_BUFFER {
                n *= 2;
            } else if k >= n {
                n = k;
            } else {
                return Ok(f2(&buf[..k]));
            }
        }
    }
}

fn os2path(s: &[u16]) -> PathBuf {
    PathBuf::from(OsString::from_wide(s))
}

#[allow(dead_code)] // Only used in backtrace::gnu::get_executable_filename()
fn wide_char_to_multi_byte(
    code_page: u32,
    flags: u32,
    s: &[u16],
    no_default_char: bool,
) -> crate::io::Result<Vec<i8>> {
    unsafe {
        let mut size = c::WideCharToMultiByte(
            code_page,
            flags,
            s.as_ptr(),
            s.len() as i32,
            ptr::null_mut(),
            0,
            ptr::null(),
            ptr::null_mut(),
        );
        if size == 0 {
            return Err(crate::io::Error::last_os_error());
        }

        let mut buf = Vec::with_capacity(size as usize);
        buf.set_len(size as usize);

        let mut used_default_char = c::FALSE;
        size = c::WideCharToMultiByte(
            code_page,
            flags,
            s.as_ptr(),
            s.len() as i32,
            buf.as_mut_ptr(),
            buf.len() as i32,
            ptr::null(),
            if no_default_char { &mut used_default_char } else { ptr::null_mut() },
        );
        if size == 0 {
            return Err(crate::io::Error::last_os_error());
        }
        if no_default_char && used_default_char == c::TRUE {
            return Err(crate::io::Error::new(
                crate::io::ErrorKind::InvalidData,
                "string cannot be converted to requested code page",
            ));
        }

        buf.set_len(size as usize);

        Ok(buf)
    }
}

pub fn truncate_utf16_at_nul(v: &[u16]) -> &[u16] {
    match unrolled_find_u16s(0, v) {
        // don't include the 0
        Some(i) => &v[..i],
        None => v,
    }
}

pub trait IsZero {
    fn is_zero(&self) -> bool;
}

macro_rules! impl_is_zero {
    ($($t:ident)*) => ($(impl IsZero for $t {
        fn is_zero(&self) -> bool {
            *self == 0
        }
    })*)
}

impl_is_zero! { i8 i16 i32 i64 isize u8 u16 u32 u64 usize }

pub fn cvt<I: IsZero>(i: I) -> crate::io::Result<I> {
    if i.is_zero() { Err(crate::io::Error::last_os_error()) } else { Ok(i) }
}

pub fn dur2timeout(dur: Duration) -> c::DWORD {
    // Note that a duration is a (u64, u32) (seconds, nanoseconds) pair, and the
    // timeouts in windows APIs are typically u32 milliseconds. To translate, we
    // have two pieces to take care of:
    //
    // * Nanosecond precision is rounded up
    // * Greater than u32::MAX milliseconds (50 days) is rounded up to INFINITE
    //   (never time out).
    dur.as_secs()
        .checked_mul(1000)
        .and_then(|ms| ms.checked_add((dur.subsec_nanos() as u64) / 1_000_000))
        .and_then(|ms| ms.checked_add(if dur.subsec_nanos() % 1_000_000 > 0 { 1 } else { 0 }))
        .map(|ms| if ms > <c::DWORD>::MAX as u64 { c::INFINITE } else { ms as c::DWORD })
        .unwrap_or(c::INFINITE)
}

// On Windows, use the processor-specific __fastfail mechanism.  In Windows 8
// and later, this will terminate the process immediately without running any
// in-process exception handlers.  In earlier versions of Windows, this
// sequence of instructions will be treated as an access violation,
// terminating the process but without necessarily bypassing all exception
// handlers.
//
// https://docs.microsoft.com/en-us/cpp/intrinsics/fastfail
#[allow(unreachable_code)]
pub fn abort_internal() -> ! {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    unsafe {
        llvm_asm!("int $$0x29" :: "{ecx}"(7) ::: volatile); // 7 is FAST_FAIL_FATAL_APP_EXIT
        crate::intrinsics::unreachable();
    }
    crate::intrinsics::abort();
}
<|endoftext|>


<commit_before>//! The implementation is based on Dmitry Vyukov's bounded MPMC queue.
//!
//! Source:
//!   - http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue
//!
//! Copyright & License:
//!   - Copyright (c) 2010-2011 Dmitry Vyukov
//!   - Simplified BSD License and Apache License, Version 2.0
//!   - http://www.1024cores.net/home/code-license

use alloc::vec::Vec;
use core::cell::UnsafeCell;
use core::fmt;
use core::marker::PhantomData;
use core::mem;
use core::ptr;
use core::sync::atomic::{self, AtomicUsize, Ordering};

use crossbeam_utils::{Backoff, CachePadded};

use err::{PopError, PushError};

/// A slot in a queue.
struct Slot<T> {
    /// The current stamp.
    ///
    /// If the stamp equals the tail, this node will be next written to. If it equals head + 1,
    /// this node will be next read from.
    stamp: AtomicUsize,

    /// The value in this slot.
    value: UnsafeCell<T>,
}

/// A bounded multi-producer multi-consumer queue.
///
/// This queue allocates a fixed-capacity buffer on construction, which is used to store pushed
/// elements. The queue cannot hold more elements than the buffer allows. Attempting to push an
/// element into a full queue will fail. Having a buffer allocated upfront makes this queue a bit
/// faster than [`SegQueue`].
///
/// [`SegQueue`]: struct.SegQueue.html
///
/// # Examples
///
/// ```
/// use crossbeam_queue::{ArrayQueue, PushError};
///
/// let q = ArrayQueue::new(2);
///
/// assert_eq!(q.push('a'), Ok(()));
/// assert_eq!(q.push('b'), Ok(()));
/// assert_eq!(q.push('c'), Err(PushError('c')));
/// assert_eq!(q.pop(), Ok('a'));
/// ```
pub struct ArrayQueue<T> {
    /// The head of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are popped from the head of the queue.
    head: CachePadded<AtomicUsize>,

    /// The tail of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are pushed into the tail of the queue.
    tail: CachePadded<AtomicUsize>,

    /// The buffer holding slots.
    buffer: *mut Slot<T>,

    /// The queue capacity.
    cap: usize,

    /// A stamp with the value of `{ lap: 1, index: 0 }`.
    one_lap: usize,

    /// Indicates that dropping an `ArrayQueue<T>` may drop elements of type `T`.
    _marker: PhantomData<T>,
}

unsafe impl<T: Send> Sync for ArrayQueue<T> {}
unsafe impl<T: Send> Send for ArrayQueue<T> {}

impl<T> ArrayQueue<T> {
    /// Creates a new bounded queue with the given capacity.
    ///
    /// # Panics
    ///
    /// Panics if the capacity is zero.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::ArrayQueue;
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    /// ```
    pub fn new(cap: usize) -> ArrayQueue<T> {
        assert!(cap > 0, "capacity must be non-zero");

        // Head is initialized to `{ lap: 0, index: 0 }`.
        // Tail is initialized to `{ lap: 0, index: 0 }`.
        let head = 0;
        let tail = 0;

        // Allocate a buffer of `cap` slots.
        let buffer = {
            let mut v = Vec::<Slot<T>>::with_capacity(cap);
            let ptr = v.as_mut_ptr();
            mem::forget(v);
            ptr
        };

        // Initialize stamps in the slots.
        for i in 0..cap {
            unsafe {
                // Set the stamp to `{ lap: 0, index: i }`.
                let slot = buffer.add(i);
                ptr::write(&mut (*slot).stamp, AtomicUsize::new(i));
            }
        }

        // One lap is the smallest power of two greater than `cap`.
        let one_lap = (cap + 1).next_power_of_two();

        ArrayQueue {
            buffer,
            cap,
            one_lap,
            head: CachePadded::new(AtomicUsize::new(head)),
            tail: CachePadded::new(AtomicUsize::new(tail)),
            _marker: PhantomData,
        }
    }

    /// Attempts to push an element into the queue.
    ///
    /// If the queue is full, the element is returned back as an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PushError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert_eq!(q.push(10), Ok(()));
    /// assert_eq!(q.push(20), Err(PushError(20)));
    /// ```
    pub fn push(&self, value: T) -> Result<(), PushError<T>> {
        let backoff = Backoff::new();
        let mut tail = self.tail.load(Ordering::Relaxed);

        loop {
            // Deconstruct the tail.
            let index = tail & (self.one_lap - 1);
            let lap = tail & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the tail and the stamp match, we may attempt to push.
            if tail == stamp {
                let new_tail = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    tail + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the tail.
                match self.tail.compare_exchange_weak(
                    tail,
                    new_tail,
                    Ordering::SeqCst,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => {
                        // Write the value into the slot and update the stamp.
                        unsafe {
                            slot.value.get().write(value);
                        }
                        slot.stamp.store(tail + 1, Ordering::Release);
                        return Ok(());
                    }
                    Err(t) => {
                        tail = t;
                        backoff.spin();
                    }
                }
            } else if stamp.wrapping_add(self.one_lap) == tail + 1 {
                atomic::fence(Ordering::SeqCst);
                let head = self.head.load(Ordering::Relaxed);

                // If the head lags one lap behind the tail as well...
                if head.wrapping_add(self.one_lap) == tail {
                    // ...then the queue is full.
                    return Err(PushError(value));
                }

                backoff.spin();
                tail = self.tail.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                tail = self.tail.load(Ordering::Relaxed);
            }
        }
    }

    /// Attempts to pop an element from the queue.
    ///
    /// If the queue is empty, an error is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    /// assert_eq!(q.push(10), Ok(()));
    ///
    /// assert_eq!(q.pop(), Ok(10));
    /// assert_eq!(q.pop(), Err(PopError));
    /// ```
    pub fn pop(&self) -> Result<T, PopError> {
        let backoff = Backoff::new();
        let mut head = self.head.load(Ordering::Relaxed);

        loop {
            // Deconstruct the head.
            let index = head & (self.one_lap - 1);
            let lap = head & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the the stamp is ahead of the head by 1, we may attempt to pop.
            if head + 1 == stamp {
                let new = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    head + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the head.
                match self.head.compare_exchange_weak(
                    head,
                    new,
                    Ordering::SeqCst,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => {
                        // Read the value from the slot and update the stamp.
                        let msg = unsafe { slot.value.get().read() };
                        slot.stamp
                            .store(head.wrapping_add(self.one_lap), Ordering::Release);
                        return Ok(msg);
                    }
                    Err(h) => {
                        head = h;
                        backoff.spin();
                    }
                }
            } else if stamp == head {
                atomic::fence(Ordering::SeqCst);
                let tail = self.tail.load(Ordering::Relaxed);

                // If the tail equals the head, that means the channel is empty.
                if tail == head {
                    return Err(PopError);
                }

                backoff.spin();
                head = self.head.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                head = self.head.load(Ordering::Relaxed);
            }
        }
    }

    /// Returns the capacity of the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    ///
    /// assert_eq!(q.capacity(), 100);
    /// ```
    pub fn capacity(&self) -> usize {
        self.cap
    }

    /// Returns `true` if the queue is empty.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    ///
    /// assert!(q.is_empty());
    /// q.push(1).unwrap();
    /// assert!(!q.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::SeqCst);
        let tail = self.tail.load(Ordering::SeqCst);

        // Is the tail lagging one lap behind head?
        // Is the tail equal to the head?
        //
        // Note: If the head changes just before we load the tail, that means there was a moment
        // when the channel was not empty, so it is safe to just return `false`.
        tail == head
    }

    /// Returns `true` if the queue is full.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert!(!q.is_full());
    /// q.push(1).unwrap();
    /// assert!(q.is_full());
    /// ```
    pub fn is_full(&self) -> bool {
        let tail = self.tail.load(Ordering::SeqCst);
        let head = self.head.load(Ordering::SeqCst);

        // Is the head lagging one lap behind tail?
        //
        // Note: If the tail changes just before we load the head, that means there was a moment
        // when the queue was not full, so it is safe to just return `false`.
        head.wrapping_add(self.one_lap) == tail
    }

    /// Returns the number of elements in the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    /// assert_eq!(q.len(), 0);
    ///
    /// q.push(10).unwrap();
    /// assert_eq!(q.len(), 1);
    ///
    /// q.push(20).unwrap();
    /// assert_eq!(q.len(), 2);
    /// ```
    pub fn len(&self) -> usize {
        loop {
            // Load the tail, then load the head.
            let tail = self.tail.load(Ordering::SeqCst);
            let head = self.head.load(Ordering::SeqCst);

            // If the tail didn't change, we've got consistent values to work with.
            if self.tail.load(Ordering::SeqCst) == tail {
                let hix = head & (self.one_lap - 1);
                let tix = tail & (self.one_lap - 1);

                return if hix < tix {
                    tix - hix
                } else if hix > tix {
                    self.cap - hix + tix
                } else if tail == head {
                    0
                } else {
                    self.cap
                };
            }
        }
    }
}

impl<T> Drop for ArrayQueue<T> {
    fn drop(&mut self) {
        // Get the index of the head.
        let hix = self.head.load(Ordering::Relaxed) & (self.one_lap - 1);

        // Loop over all slots that hold a message and drop them.
        for i in 0..self.len() {
            // Compute the index of the next slot holding a message.
            let index = if hix + i < self.cap {
                hix + i
            } else {
                hix + i - self.cap
            };

            unsafe {
                self.buffer.add(index).drop_in_place();
            }
        }

        // Finally, deallocate the buffer, but don't run any destructors.
        unsafe {
            Vec::from_raw_parts(self.buffer, 0, self.cap);
        }
    }
}

impl<T> fmt::Debug for ArrayQueue<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("ArrayQueue { .. }")
    }
}
<commit_msg>queue: fix unsoundness issue in array_queue (introduce MaybeUninit)<commit_after>//! The implementation is based on Dmitry Vyukov's bounded MPMC queue.
//!
//! Source:
//!   - http://www.1024cores.net/home/lock-free-algorithms/queues/bounded-mpmc-queue
//!
//! Copyright & License:
//!   - Copyright (c) 2010-2011 Dmitry Vyukov
//!   - Simplified BSD License and Apache License, Version 2.0
//!   - http://www.1024cores.net/home/code-license

use alloc::vec::Vec;
use core::cell::UnsafeCell;
use core::fmt;
use core::marker::PhantomData;
use core::mem::{self, MaybeUninit};
use core::ptr;
use core::sync::atomic::{self, AtomicUsize, Ordering};

use crossbeam_utils::{Backoff, CachePadded};

use err::{PopError, PushError};

/// A slot in a queue.
struct Slot<T> {
    /// The current stamp.
    ///
    /// If the stamp equals the tail, this node will be next written to. If it equals head + 1,
    /// this node will be next read from.
    stamp: AtomicUsize,

    /// The value in this slot.
    value: UnsafeCell<MaybeUninit<T>>,
}

/// A bounded multi-producer multi-consumer queue.
///
/// This queue allocates a fixed-capacity buffer on construction, which is used to store pushed
/// elements. The queue cannot hold more elements than the buffer allows. Attempting to push an
/// element into a full queue will fail. Having a buffer allocated upfront makes this queue a bit
/// faster than [`SegQueue`].
///
/// [`SegQueue`]: struct.SegQueue.html
///
/// # Examples
///
/// ```
/// use crossbeam_queue::{ArrayQueue, PushError};
///
/// let q = ArrayQueue::new(2);
///
/// assert_eq!(q.push('a'), Ok(()));
/// assert_eq!(q.push('b'), Ok(()));
/// assert_eq!(q.push('c'), Err(PushError('c')));
/// assert_eq!(q.pop(), Ok('a'));
/// ```
pub struct ArrayQueue<T> {
    /// The head of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are popped from the head of the queue.
    head: CachePadded<AtomicUsize>,

    /// The tail of the queue.
    ///
    /// This value is a "stamp" consisting of an index into the buffer and a lap, but packed into a
    /// single `usize`. The lower bits represent the index, while the upper bits represent the lap.
    ///
    /// Elements are pushed into the tail of the queue.
    tail: CachePadded<AtomicUsize>,

    /// The buffer holding slots.
    buffer: *mut Slot<T>,

    /// The queue capacity.
    cap: usize,

    /// A stamp with the value of `{ lap: 1, index: 0 }`.
    one_lap: usize,

    /// Indicates that dropping an `ArrayQueue<T>` may drop elements of type `T`.
    _marker: PhantomData<T>,
}

unsafe impl<T: Send> Sync for ArrayQueue<T> {}
unsafe impl<T: Send> Send for ArrayQueue<T> {}

impl<T> ArrayQueue<T> {
    /// Creates a new bounded queue with the given capacity.
    ///
    /// # Panics
    ///
    /// Panics if the capacity is zero.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::ArrayQueue;
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    /// ```
    pub fn new(cap: usize) -> ArrayQueue<T> {
        assert!(cap > 0, "capacity must be non-zero");

        // Head is initialized to `{ lap: 0, index: 0 }`.
        // Tail is initialized to `{ lap: 0, index: 0 }`.
        let head = 0;
        let tail = 0;

        // Allocate a buffer of `cap` slots.
        let buffer = {
            let mut v = Vec::<Slot<T>>::with_capacity(cap);
            let ptr = v.as_mut_ptr();
            mem::forget(v);
            ptr
        };

        // Initialize stamps in the slots.
        for i in 0..cap {
            unsafe {
                // Set the stamp to `{ lap: 0, index: i }`.
                let slot = buffer.add(i);
                ptr::write(&mut (*slot).stamp, AtomicUsize::new(i));
            }
        }

        // One lap is the smallest power of two greater than `cap`.
        let one_lap = (cap + 1).next_power_of_two();

        ArrayQueue {
            buffer,
            cap,
            one_lap,
            head: CachePadded::new(AtomicUsize::new(head)),
            tail: CachePadded::new(AtomicUsize::new(tail)),
            _marker: PhantomData,
        }
    }

    /// Attempts to push an element into the queue.
    ///
    /// If the queue is full, the element is returned back as an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PushError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert_eq!(q.push(10), Ok(()));
    /// assert_eq!(q.push(20), Err(PushError(20)));
    /// ```
    pub fn push(&self, value: T) -> Result<(), PushError<T>> {
        let backoff = Backoff::new();
        let mut tail = self.tail.load(Ordering::Relaxed);

        loop {
            // Deconstruct the tail.
            let index = tail & (self.one_lap - 1);
            let lap = tail & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the tail and the stamp match, we may attempt to push.
            if tail == stamp {
                let new_tail = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    tail + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the tail.
                match self.tail.compare_exchange_weak(
                    tail,
                    new_tail,
                    Ordering::SeqCst,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => {
                        // Write the value into the slot and update the stamp.
                        unsafe { slot.value.get().write(MaybeUninit::new(value)) }
                        slot.stamp.store(tail + 1, Ordering::Release);
                        return Ok(());
                    }
                    Err(t) => {
                        tail = t;
                        backoff.spin();
                    }
                }
            } else if stamp.wrapping_add(self.one_lap) == tail + 1 {
                atomic::fence(Ordering::SeqCst);
                let head = self.head.load(Ordering::Relaxed);

                // If the head lags one lap behind the tail as well...
                if head.wrapping_add(self.one_lap) == tail {
                    // ...then the queue is full.
                    return Err(PushError(value));
                }

                backoff.spin();
                tail = self.tail.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                tail = self.tail.load(Ordering::Relaxed);
            }
        }
    }

    /// Attempts to pop an element from the queue.
    ///
    /// If the queue is empty, an error is returned.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    /// assert_eq!(q.push(10), Ok(()));
    ///
    /// assert_eq!(q.pop(), Ok(10));
    /// assert_eq!(q.pop(), Err(PopError));
    /// ```
    pub fn pop(&self) -> Result<T, PopError> {
        let backoff = Backoff::new();
        let mut head = self.head.load(Ordering::Relaxed);

        loop {
            // Deconstruct the head.
            let index = head & (self.one_lap - 1);
            let lap = head & !(self.one_lap - 1);

            // Inspect the corresponding slot.
            let slot = unsafe { &*self.buffer.add(index) };
            let stamp = slot.stamp.load(Ordering::Acquire);

            // If the the stamp is ahead of the head by 1, we may attempt to pop.
            if head + 1 == stamp {
                let new = if index + 1 < self.cap {
                    // Same lap, incremented index.
                    // Set to `{ lap: lap, index: index + 1 }`.
                    head + 1
                } else {
                    // One lap forward, index wraps around to zero.
                    // Set to `{ lap: lap.wrapping_add(1), index: 0 }`.
                    lap.wrapping_add(self.one_lap)
                };

                // Try moving the head.
                match self.head.compare_exchange_weak(
                    head,
                    new,
                    Ordering::SeqCst,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => {
                        // Read the value from the slot and update the stamp.
                        let msg = unsafe { slot.value.get().read().assume_init() };
                        slot.stamp
                            .store(head.wrapping_add(self.one_lap), Ordering::Release);
                        return Ok(msg);
                    }
                    Err(h) => {
                        head = h;
                        backoff.spin();
                    }
                }
            } else if stamp == head {
                atomic::fence(Ordering::SeqCst);
                let tail = self.tail.load(Ordering::Relaxed);

                // If the tail equals the head, that means the channel is empty.
                if tail == head {
                    return Err(PopError);
                }

                backoff.spin();
                head = self.head.load(Ordering::Relaxed);
            } else {
                // Snooze because we need to wait for the stamp to get updated.
                backoff.snooze();
                head = self.head.load(Ordering::Relaxed);
            }
        }
    }

    /// Returns the capacity of the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::<i32>::new(100);
    ///
    /// assert_eq!(q.capacity(), 100);
    /// ```
    pub fn capacity(&self) -> usize {
        self.cap
    }

    /// Returns `true` if the queue is empty.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    ///
    /// assert!(q.is_empty());
    /// q.push(1).unwrap();
    /// assert!(!q.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        let head = self.head.load(Ordering::SeqCst);
        let tail = self.tail.load(Ordering::SeqCst);

        // Is the tail lagging one lap behind head?
        // Is the tail equal to the head?
        //
        // Note: If the head changes just before we load the tail, that means there was a moment
        // when the channel was not empty, so it is safe to just return `false`.
        tail == head
    }

    /// Returns `true` if the queue is full.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(1);
    ///
    /// assert!(!q.is_full());
    /// q.push(1).unwrap();
    /// assert!(q.is_full());
    /// ```
    pub fn is_full(&self) -> bool {
        let tail = self.tail.load(Ordering::SeqCst);
        let head = self.head.load(Ordering::SeqCst);

        // Is the head lagging one lap behind tail?
        //
        // Note: If the tail changes just before we load the head, that means there was a moment
        // when the queue was not full, so it is safe to just return `false`.
        head.wrapping_add(self.one_lap) == tail
    }

    /// Returns the number of elements in the queue.
    ///
    /// # Examples
    ///
    /// ```
    /// use crossbeam_queue::{ArrayQueue, PopError};
    ///
    /// let q = ArrayQueue::new(100);
    /// assert_eq!(q.len(), 0);
    ///
    /// q.push(10).unwrap();
    /// assert_eq!(q.len(), 1);
    ///
    /// q.push(20).unwrap();
    /// assert_eq!(q.len(), 2);
    /// ```
    pub fn len(&self) -> usize {
        loop {
            // Load the tail, then load the head.
            let tail = self.tail.load(Ordering::SeqCst);
            let head = self.head.load(Ordering::SeqCst);

            // If the tail didn't change, we've got consistent values to work with.
            if self.tail.load(Ordering::SeqCst) == tail {
                let hix = head & (self.one_lap - 1);
                let tix = tail & (self.one_lap - 1);

                return if hix < tix {
                    tix - hix
                } else if hix > tix {
                    self.cap - hix + tix
                } else if tail == head {
                    0
                } else {
                    self.cap
                };
            }
        }
    }
}

impl<T> Drop for ArrayQueue<T> {
    fn drop(&mut self) {
        // Get the index of the head.
        let hix = self.head.load(Ordering::Relaxed) & (self.one_lap - 1);

        // Loop over all slots that hold a message and drop them.
        for i in 0..self.len() {
            // Compute the index of the next slot holding a message.
            let index = if hix + i < self.cap {
                hix + i
            } else {
                hix + i - self.cap
            };

            unsafe {
                self.buffer.add(index).drop_in_place();
            }
        }

        // Finally, deallocate the buffer, but don't run any destructors.
        unsafe {
            Vec::from_raw_parts(self.buffer, 0, self.cap);
        }
    }
}

impl<T> fmt::Debug for ArrayQueue<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.pad("ArrayQueue { .. }")
    }
}
<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use crate::transport::ti50emulator::Ti50Emulator;
use crate::transport::Transport;
use anyhow::Result;
use std::path::PathBuf;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub struct Ti50EmulatorOpts {
    #[structopt(long)]
    executable_directory: String,

    #[structopt(long, default_value = "host_emulation")]
    executable: String,

    #[structopt(long)]
    instance_prefix: String,
}

pub fn create(args: &Ti50EmulatorOpts) -> Result<Box<dyn Transport>> {
    Ok(Box::new(Ti50Emulator::open(
        PathBuf::from(&args.executable_directory),
        &args.executable,
        &args.instance_prefix,
    )?))
}
<commit_msg>[opentitanlib] Add auto-default values for the ti50 emulator backend<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use crate::transport::ti50emulator::Ti50Emulator;
use crate::transport::Transport;
use anyhow::Result;
use std::path::PathBuf;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub struct Ti50EmulatorOpts {
    #[structopt(long, default_value)]
    executable_directory: String,

    #[structopt(long, default_value = "host_emulation")]
    executable: String,

    #[structopt(long, default_value = "ti50")]
    instance_prefix: String,
}

pub fn create(args: &Ti50EmulatorOpts) -> Result<Box<dyn Transport>> {
    Ok(Box::new(Ti50Emulator::open(
        PathBuf::from(&args.executable_directory),
        &args.executable,
        &args.instance_prefix,
    )?))
}
<|endoftext|>


<commit_before>use std::ops::Range;

use lsp_types::Position;

pub trait RangeToUtf8ByteOffset {
    fn to_utf8_byte_offset(&self, text: &str) -> Range<usize>;
}
impl RangeToUtf8ByteOffset for lsp_types::Range {
    fn to_utf8_byte_offset(&self, text: &str) -> Range<usize> {
        let start = self.start.to_utf8_byte_offset(text);
        let end = self.end.to_utf8_byte_offset(text);
        start..end
    }
}

pub trait PositionToUtf8ByteOffset {
    fn to_utf8_byte_offset(&self, text: &str) -> usize;
}
impl PositionToUtf8ByteOffset for Position {
    fn to_utf8_byte_offset(&self, text: &str) -> usize {
        let mut line_index = 0;
        let mut line_offset = 0;
        while line_index < self.line {
            match text.bytes().nth(line_offset).unwrap() {
                b'\n' => {
                    line_index += 1;
                    line_offset += 1;
                }
                _ => {
                    line_offset += 1;
                }
            }
        }

        let mut line_length_bytes = 0;
        loop {
            match text.bytes().nth(line_offset + line_length_bytes) {
                Some(b'\r' | b'\n') | None => break,
                Some(_) => line_length_bytes += 1,
            }
        }

        let line = &text[line_offset..line_offset + line_length_bytes];

        let words = line.encode_utf16().collect::<Vec<_>>();
        let char_offset = if self.character as usize >= words.len() {
            line_length_bytes
        } else {
            String::from_utf16(&words[0..self.character as usize])
                .unwrap()
                .len()
        };

        line_offset + char_offset
    }
}
<commit_msg>feat(compiler): add conversion from byte offset to LSP Position/Range<commit_after>use std::ops::Range;

use lsp_types::Position;

pub trait RangeToUtf8ByteOffset {
    fn to_utf8_byte_offset(&self, text: &str) -> Range<usize>;
}
impl RangeToUtf8ByteOffset for lsp_types::Range {
    fn to_utf8_byte_offset(&self, text: &str) -> Range<usize> {
        let start = self.start.to_utf8_byte_offset(text);
        let end = self.end.to_utf8_byte_offset(text);
        start..end
    }
}

pub trait PositionToUtf8ByteOffset {
    fn to_utf8_byte_offset(&self, text: &str) -> usize;
}
impl PositionToUtf8ByteOffset for Position {
    fn to_utf8_byte_offset(&self, text: &str) -> usize {
        let mut line_index = 0;
        let mut line_offset = 0;
        while line_index < self.line {
            match text.bytes().nth(line_offset).unwrap() {
                b'\n' => {
                    line_index += 1;
                    line_offset += 1;
                }
                _ => {
                    line_offset += 1;
                }
            }
        }

        let mut line_length_bytes = 0;
        loop {
            match text.bytes().nth(line_offset + line_length_bytes) {
                Some(b'\r' | b'\n') | None => break,
                Some(_) => line_length_bytes += 1,
            }
        }

        let line = &text[line_offset..line_offset + line_length_bytes];

        let words = line.encode_utf16().collect::<Vec<_>>();
        let char_offset = if self.character as usize >= words.len() {
            line_length_bytes
        } else {
            String::from_utf16(&words[0..self.character as usize])
                .unwrap()
                .len()
        };

        line_offset + char_offset
    }
}

pub trait RangeToLsp {
    fn to_lsp(&self, text: &str) -> lsp_types::Range;
}
impl RangeToLsp for Range<usize> {
    fn to_lsp(&self, text: &str) -> lsp_types::Range {
        let start = self.start.utf8_byte_offset_to_lsp(text);
        let end = self.end.utf8_byte_offset_to_lsp(text);
        lsp_types::Range { start, end }
    }
}

pub trait Utf8ByteOffsetToLsp {
    fn utf8_byte_offset_to_lsp(&self, text: &str) -> Position;
}
impl Utf8ByteOffsetToLsp for usize {
    fn utf8_byte_offset_to_lsp(&self, text: &str) -> Position {
        let mut index = 0;
        let mut line_index = 0;
        let mut line_start = 0;
        let mut character_byte_offset = 0;
        while &index < self {
            match text.bytes().nth(index).unwrap() {
                b'\n' => {
                    line_index += 1;
                    line_start = index;
                    character_byte_offset = 0;
                }
                _ => {
                    character_byte_offset += 1;
                }
            }
            index += 1;
        }

        let utf16_offset = text[line_start..line_start + character_byte_offset]
            .encode_utf16()
            .count();
        Position {
            line: line_index,
            character: utf16_offset as u32,
        }
    }
}
<|endoftext|>


<commit_before>#![deny(warnings, rust_2018_idioms)]

use pin_utils::pin_mut;
use std::future::Future;
use tokio_sync::lock::Lock;
use tokio_test::task::MockTask;
use tokio_test::{assert_pending, assert_ready};

#[test]
fn straight_execution() {
    let mut task = MockTask::new();
    let mut l = Lock::new(100);

    // We can immediately acquire the lock and take the value
    task.enter(|cx| {
        let mut g = assert_ready!(l.poll_lock(cx));
        assert_eq!(&*g, &100);
        *g = 99;
        drop(g);

        let mut g = assert_ready!(l.poll_lock(cx));
        assert_eq!(&*g, &99);
        *g = 98;
        drop(g);

        let mut g = assert_ready!(l.poll_lock(cx));
        assert_eq!(&*g, &98);

        // We can continue to access the guard even if the lock is dropped
        drop(l);
        *g = 97;
        assert_eq!(&*g, &97);
    });
}

#[test]
fn readiness() {
    let mut t1 = MockTask::new();
    let mut t2 = MockTask::new();

    let mut l = Lock::new(100);

    let g = assert_ready!(t1.enter(|cx| l.poll_lock(cx)));

    // We can't now acquire the lease since it's already held in g
    assert_pending!(t2.enter(|cx| l.poll_lock(cx)));

    // But once g unlocks, we can acquire it
    drop(g);
    assert!(t2.is_woken());
    assert_ready!(t2.enter(|cx| l.poll_lock(cx)));
}

#[test]
fn lock() {
    let mut lock = Lock::new(false);

    let mut lock2 = lock.clone();
    std::thread::spawn(move || {
        let l = lock2.lock();
        pin_mut!(l);

        let mut task = MockTask::new();
        let mut g = task.enter(|cx| assert_ready!(l.poll(cx)));
        std::thread::sleep(std::time::Duration::from_millis(500));
        *g = true;
        drop(g);
    });

    std::thread::sleep(std::time::Duration::from_millis(50));
    let mut task = MockTask::new();
    let l = lock.lock();
    pin_mut!(l);

    task.enter(|cx| assert_pending!(l.poll(cx)));
    std::thread::sleep(std::time::Duration::from_millis(500));
    assert!(task.is_woken());
}
<commit_msg>sync: Fix lock test to actually test the inner lock value (#1197)<commit_after>#![deny(warnings, rust_2018_idioms)]

use pin_utils::pin_mut;
use tokio_sync::lock::Lock;
use tokio_test::task::MockTask;
use tokio_test::{assert_pending, assert_ready};

#[test]
fn straight_execution() {
    let mut task = MockTask::new();
    let mut l = Lock::new(100);

    // We can immediately acquire the lock and take the value
    task.enter(|cx| {
        let mut g = assert_ready!(l.poll_lock(cx));
        assert_eq!(&*g, &100);
        *g = 99;
        drop(g);

        let mut g = assert_ready!(l.poll_lock(cx));
        assert_eq!(&*g, &99);
        *g = 98;
        drop(g);

        let mut g = assert_ready!(l.poll_lock(cx));
        assert_eq!(&*g, &98);

        // We can continue to access the guard even if the lock is dropped
        drop(l);
        *g = 97;
        assert_eq!(&*g, &97);
    });
}

#[test]
fn readiness() {
    let mut t1 = MockTask::new();
    let mut t2 = MockTask::new();

    let mut l = Lock::new(100);

    let g = assert_ready!(t1.enter(|cx| l.poll_lock(cx)));

    // We can't now acquire the lease since it's already held in g
    assert_pending!(t2.enter(|cx| l.poll_lock(cx)));

    // But once g unlocks, we can acquire it
    drop(g);
    assert!(t2.is_woken());
    assert_ready!(t2.enter(|cx| l.poll_lock(cx)));
}

#[test]
fn lock() {
    let mut lock = Lock::new(false);

    let mut lock2 = lock.clone();
    std::thread::spawn(move || {
        let l = lock2.lock();
        pin_mut!(l);

        let mut task = MockTask::new();
        let mut g = assert_ready!(task.poll(&mut l));
        std::thread::sleep(std::time::Duration::from_millis(500));
        *g = true;
        drop(g);
    });

    std::thread::sleep(std::time::Duration::from_millis(50));
    let mut task = MockTask::new();
    let l = lock.lock();
    pin_mut!(l);

    assert_pending!(task.poll(&mut l));

    std::thread::sleep(std::time::Duration::from_millis(500));
    assert!(task.is_woken());
    let result = assert_ready!(task.poll(&mut l));
    assert!(*result);
}
<|endoftext|>


<commit_before>//! Unix-specific types for signal handling.
//!
//! This module is only defined on Unix platforms and contains the primary
//! `Signal` type for receiving notifications of signals.

#![cfg(unix)]

use crate::signal::registry::{globals, EventId, EventInfo, Globals, Init, Storage};
use crate::sync::mpsc::error::TryRecvError;
use crate::sync::mpsc::{channel, Receiver};

use libc::c_int;
use mio::net::UnixStream;
use std::io::{self, Error, ErrorKind, Write};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Once;
use std::task::{Context, Poll};

pub(crate) mod driver;
use self::driver::Handle;

pub(crate) type OsStorage = Vec<SignalInfo>;

// Number of different unix signals
// (FreeBSD has 33)
const SIGNUM: usize = 33;

impl Init for OsStorage {
    fn init() -> Self {
        (0..SIGNUM).map(|_| SignalInfo::default()).collect()
    }
}

impl Storage for OsStorage {
    fn event_info(&self, id: EventId) -> Option<&EventInfo> {
        self.get(id).map(|si| &si.event_info)
    }

    fn for_each<'a, F>(&'a self, f: F)
    where
        F: FnMut(&'a EventInfo),
    {
        self.iter().map(|si| &si.event_info).for_each(f)
    }
}

#[derive(Debug)]
pub(crate) struct OsExtraData {
    sender: UnixStream,
    receiver: UnixStream,
}

impl Init for OsExtraData {
    fn init() -> Self {
        let (receiver, sender) = UnixStream::pair().expect("failed to create UnixStream");

        Self { sender, receiver }
    }
}

/// Represents the specific kind of signal to listen for.
#[derive(Debug, Clone, Copy)]
pub struct SignalKind(c_int);

impl SignalKind {
    /// Allows for listening to any valid OS signal.
    ///
    /// For example, this can be used for listening for platform-specific
    /// signals.
    /// ```rust,no_run
    /// # use tokio::signal::unix::SignalKind;
    /// # let signum = -1;
    /// // let signum = libc::OS_SPECIFIC_SIGNAL;
    /// let kind = SignalKind::from_raw(signum);
    /// ```
    pub fn from_raw(signum: c_int) -> Self {
        Self(signum)
    }

    /// Represents the SIGALRM signal.
    ///
    /// On Unix systems this signal is sent when a real-time timer has expired.
    /// By default, the process is terminated by this signal.
    pub fn alarm() -> Self {
        Self(libc::SIGALRM)
    }

    /// Represents the SIGCHLD signal.
    ///
    /// On Unix systems this signal is sent when the status of a child process
    /// has changed. By default, this signal is ignored.
    pub fn child() -> Self {
        Self(libc::SIGCHLD)
    }

    /// Represents the SIGHUP signal.
    ///
    /// On Unix systems this signal is sent when the terminal is disconnected.
    /// By default, the process is terminated by this signal.
    pub fn hangup() -> Self {
        Self(libc::SIGHUP)
    }

    /// Represents the SIGINFO signal.
    ///
    /// On Unix systems this signal is sent to request a status update from the
    /// process. By default, this signal is ignored.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    pub fn info() -> Self {
        Self(libc::SIGINFO)
    }

    /// Represents the SIGINT signal.
    ///
    /// On Unix systems this signal is sent to interrupt a program.
    /// By default, the process is terminated by this signal.
    pub fn interrupt() -> Self {
        Self(libc::SIGINT)
    }

    /// Represents the SIGIO signal.
    ///
    /// On Unix systems this signal is sent when I/O operations are possible
    /// on some file descriptor. By default, this signal is ignored.
    pub fn io() -> Self {
        Self(libc::SIGIO)
    }

    /// Represents the SIGPIPE signal.
    ///
    /// On Unix systems this signal is sent when the process attempts to write
    /// to a pipe which has no reader. By default, the process is terminated by
    /// this signal.
    pub fn pipe() -> Self {
        Self(libc::SIGPIPE)
    }

    /// Represents the SIGQUIT signal.
    ///
    /// On Unix systems this signal is sent to issue a shutdown of the
    /// process, after which the OS will dump the process core.
    /// By default, the process is terminated by this signal.
    pub fn quit() -> Self {
        Self(libc::SIGQUIT)
    }

    /// Represents the SIGTERM signal.
    ///
    /// On Unix systems this signal is sent to issue a shutdown of the
    /// process. By default, the process is terminated by this signal.
    pub fn terminate() -> Self {
        Self(libc::SIGTERM)
    }

    /// Represents the SIGUSR1 signal.
    ///
    /// On Unix systems this is a user defined signal.
    /// By default, the process is terminated by this signal.
    pub fn user_defined1() -> Self {
        Self(libc::SIGUSR1)
    }

    /// Represents the SIGUSR2 signal.
    ///
    /// On Unix systems this is a user defined signal.
    /// By default, the process is terminated by this signal.
    pub fn user_defined2() -> Self {
        Self(libc::SIGUSR2)
    }

    /// Represents the SIGWINCH signal.
    ///
    /// On Unix systems this signal is sent when the terminal window is resized.
    /// By default, this signal is ignored.
    pub fn window_change() -> Self {
        Self(libc::SIGWINCH)
    }
}

pub(crate) struct SignalInfo {
    event_info: EventInfo,
    init: Once,
    initialized: AtomicBool,
}

impl Default for SignalInfo {
    fn default() -> SignalInfo {
        SignalInfo {
            event_info: Default::default(),
            init: Once::new(),
            initialized: AtomicBool::new(false),
        }
    }
}

/// Our global signal handler for all signals registered by this module.
///
/// The purpose of this signal handler is to primarily:
///
/// 1. Flag that our specific signal was received (e.g. store an atomic flag)
/// 2. Wake up the driver by writing a byte to a pipe
///
/// Those two operations should both be async-signal safe.
fn action(globals: Pin<&'static Globals>, signal: c_int) {
    globals.record_event(signal as EventId);

    // Send a wakeup, ignore any errors (anything reasonably possible is
    // full pipe and then it will wake up anyway).
    let mut sender = &globals.sender;
    drop(sender.write(&[1]));
}

/// Enables this module to receive signal notifications for the `signal`
/// provided.
///
/// This will register the signal handler if it hasn't already been registered,
/// returning any error along the way if that fails.
fn signal_enable(signal: c_int, handle: Handle) -> io::Result<()> {
    if signal < 0 || signal_hook_registry::FORBIDDEN.contains(&signal) {
        return Err(Error::new(
            ErrorKind::Other,
            format!("Refusing to register signal {}", signal),
        ));
    }

    // Check that we have a signal driver running
    handle.check_inner()?;

    let globals = globals();
    let siginfo = match globals.storage().get(signal as EventId) {
        Some(slot) => slot,
        None => return Err(io::Error::new(io::ErrorKind::Other, "signal too large")),
    };
    let mut registered = Ok(());
    siginfo.init.call_once(|| {
        registered = unsafe {
            signal_hook_registry::register(signal, move || action(globals, signal)).map(|_| ())
        };
        if registered.is_ok() {
            siginfo.initialized.store(true, Ordering::Relaxed);
        }
    });
    registered?;
    // If the call_once failed, it won't be retried on the next attempt to register the signal. In
    // such case it is not run, registered is still `Ok(())`, initialized is still `false`.
    if siginfo.initialized.load(Ordering::Relaxed) {
        Ok(())
    } else {
        Err(Error::new(
            ErrorKind::Other,
            "Failed to register signal handler",
        ))
    }
}

/// A stream of events for receiving a particular type of OS signal.
///
/// In general signal handling on Unix is a pretty tricky topic, and this
/// structure is no exception! There are some important limitations to keep in
/// mind when using `Signal` streams:
///
/// * Signals handling in Unix already necessitates coalescing signals
///   together sometimes. This `Signal` stream is also no exception here in
///   that it will also coalesce signals. That is, even if the signal handler
///   for this process runs multiple times, the `Signal` stream may only return
///   one signal notification. Specifically, before `poll` is called, all
///   signal notifications are coalesced into one item returned from `poll`.
///   Once `poll` has been called, however, a further signal is guaranteed to
///   be yielded as an item.
///
///   Put another way, any element pulled off the returned stream corresponds to
///   *at least one* signal, but possibly more.
///
/// * Signal handling in general is relatively inefficient. Although some
///   improvements are possible in this crate, it's recommended to not plan on
///   having millions of signal channels open.
///
/// If you've got any questions about this feel free to open an issue on the
/// repo! New approaches to alleviate some of these limitations are always
/// appreciated!
///
/// # Caveats
///
/// The first time that a `Signal` instance is registered for a particular
/// signal kind, an OS signal-handler is installed which replaces the default
/// platform behavior when that signal is received, **for the duration of the
/// entire process**.
///
/// For example, Unix systems will terminate a process by default when it
/// receives SIGINT. But, when a `Signal` instance is created to listen for
/// this signal, the next SIGINT that arrives will be translated to a stream
/// event, and the process will continue to execute. **Even if this `Signal`
/// instance is dropped, subsequent SIGINT deliveries will end up captured by
/// Tokio, and the default platform behavior will NOT be reset**.
///
/// Thus, applications should take care to ensure the expected signal behavior
/// occurs as expected after listening for specific signals.
///
/// # Examples
///
/// Wait for SIGHUP
///
/// ```rust,no_run
/// use tokio::signal::unix::{signal, SignalKind};
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     // An infinite stream of hangup signals.
///     let mut stream = signal(SignalKind::hangup())?;
///
///     // Print whenever a HUP signal is received
///     loop {
///         stream.recv().await;
///         println!("got signal HUP");
///     }
/// }
/// ```
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Signal {
    rx: Receiver<()>,
}

/// Creates a new stream which will receive notifications when the current
/// process receives the specified signal `kind`.
///
/// This function will create a new stream which binds to the default reactor.
/// The `Signal` stream is an infinite stream which will receive
/// notifications whenever a signal is received. More documentation can be
/// found on `Signal` itself, but to reiterate:
///
/// * Signals may be coalesced beyond what the kernel already does.
/// * Once a signal handler is registered with the process the underlying
///   libc signal handler is never unregistered.
///
/// A `Signal` stream can be created for a particular signal number
/// multiple times. When a signal is received then all the associated
/// channels will receive the signal notification.
///
/// # Errors
///
/// * If the lower-level C functions fail for some reason.
/// * If the previous initialization of this specific signal failed.
/// * If the signal is one of
///   [`signal_hook::FORBIDDEN`](fn@signal_hook_registry::register#panics)
pub fn signal(kind: SignalKind) -> io::Result<Signal> {
    signal_with_handle(kind, Handle::current())
}

pub(crate) fn signal_with_handle(kind: SignalKind, handle: Handle) -> io::Result<Signal> {
    let signal = kind.0;

    // Turn the signal delivery on once we are ready for it
    signal_enable(signal, handle)?;

    // One wakeup in a queue is enough, no need for us to buffer up any
    // more.
    let (tx, rx) = channel(1);
    globals().register_listener(signal as EventId, tx);

    Ok(Signal { rx })
}

impl Signal {
    /// Receives the next signal notification event.
    ///
    /// `None` is returned if no more events can be received by this stream.
    ///
    /// # Examples
    ///
    /// Wait for SIGHUP
    ///
    /// ```rust,no_run
    /// use tokio::signal::unix::{signal, SignalKind};
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
    ///     // An infinite stream of hangup signals.
    ///     let mut stream = signal(SignalKind::hangup())?;
    ///
    ///     // Print whenever a HUP signal is received
    ///     loop {
    ///         stream.recv().await;
    ///         println!("got signal HUP");
    ///     }
    /// }
    /// ```
    pub async fn recv(&mut self) -> Option<()> {
        use crate::future::poll_fn;
        poll_fn(|cx| self.poll_recv(cx)).await
    }

    pub(crate) fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
        self.rx.poll_recv(cx)
    }

    /// Try to receive a signal notification without blocking or registering a waker.
    pub(crate) fn try_recv(&mut self) -> Result<(), TryRecvError> {
        self.rx.try_recv()
    }
}

// Work around for abstracting streams internally
pub(crate) trait InternalStream: Unpin {
    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>>;
    fn try_recv(&mut self) -> Result<(), TryRecvError>;
}

impl InternalStream for Signal {
    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
        self.poll_recv(cx)
    }

    fn try_recv(&mut self) -> Result<(), TryRecvError> {
        self.try_recv()
    }
}

pub(crate) fn ctrl_c() -> io::Result<Signal> {
    signal(SignalKind::interrupt())
}

#[cfg(all(test, not(loom)))]
mod tests {
    use super::*;

    #[test]
    fn signal_enable_error_on_invalid_input() {
        signal_enable(-1, Handle::default()).unwrap_err();
    }

    #[test]
    fn signal_enable_error_on_forbidden_input() {
        signal_enable(signal_hook_registry::FORBIDDEN[0], Handle::default()).unwrap_err();
    }
}
<commit_msg>signal: make Signal::poll_recv method public (#3383)<commit_after>//! Unix-specific types for signal handling.
//!
//! This module is only defined on Unix platforms and contains the primary
//! `Signal` type for receiving notifications of signals.

#![cfg(unix)]

use crate::signal::registry::{globals, EventId, EventInfo, Globals, Init, Storage};
use crate::sync::mpsc::error::TryRecvError;
use crate::sync::mpsc::{channel, Receiver};

use libc::c_int;
use mio::net::UnixStream;
use std::io::{self, Error, ErrorKind, Write};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Once;
use std::task::{Context, Poll};

pub(crate) mod driver;
use self::driver::Handle;

pub(crate) type OsStorage = Vec<SignalInfo>;

// Number of different unix signals
// (FreeBSD has 33)
const SIGNUM: usize = 33;

impl Init for OsStorage {
    fn init() -> Self {
        (0..SIGNUM).map(|_| SignalInfo::default()).collect()
    }
}

impl Storage for OsStorage {
    fn event_info(&self, id: EventId) -> Option<&EventInfo> {
        self.get(id).map(|si| &si.event_info)
    }

    fn for_each<'a, F>(&'a self, f: F)
    where
        F: FnMut(&'a EventInfo),
    {
        self.iter().map(|si| &si.event_info).for_each(f)
    }
}

#[derive(Debug)]
pub(crate) struct OsExtraData {
    sender: UnixStream,
    receiver: UnixStream,
}

impl Init for OsExtraData {
    fn init() -> Self {
        let (receiver, sender) = UnixStream::pair().expect("failed to create UnixStream");

        Self { sender, receiver }
    }
}

/// Represents the specific kind of signal to listen for.
#[derive(Debug, Clone, Copy)]
pub struct SignalKind(c_int);

impl SignalKind {
    /// Allows for listening to any valid OS signal.
    ///
    /// For example, this can be used for listening for platform-specific
    /// signals.
    /// ```rust,no_run
    /// # use tokio::signal::unix::SignalKind;
    /// # let signum = -1;
    /// // let signum = libc::OS_SPECIFIC_SIGNAL;
    /// let kind = SignalKind::from_raw(signum);
    /// ```
    pub fn from_raw(signum: c_int) -> Self {
        Self(signum)
    }

    /// Represents the SIGALRM signal.
    ///
    /// On Unix systems this signal is sent when a real-time timer has expired.
    /// By default, the process is terminated by this signal.
    pub fn alarm() -> Self {
        Self(libc::SIGALRM)
    }

    /// Represents the SIGCHLD signal.
    ///
    /// On Unix systems this signal is sent when the status of a child process
    /// has changed. By default, this signal is ignored.
    pub fn child() -> Self {
        Self(libc::SIGCHLD)
    }

    /// Represents the SIGHUP signal.
    ///
    /// On Unix systems this signal is sent when the terminal is disconnected.
    /// By default, the process is terminated by this signal.
    pub fn hangup() -> Self {
        Self(libc::SIGHUP)
    }

    /// Represents the SIGINFO signal.
    ///
    /// On Unix systems this signal is sent to request a status update from the
    /// process. By default, this signal is ignored.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    pub fn info() -> Self {
        Self(libc::SIGINFO)
    }

    /// Represents the SIGINT signal.
    ///
    /// On Unix systems this signal is sent to interrupt a program.
    /// By default, the process is terminated by this signal.
    pub fn interrupt() -> Self {
        Self(libc::SIGINT)
    }

    /// Represents the SIGIO signal.
    ///
    /// On Unix systems this signal is sent when I/O operations are possible
    /// on some file descriptor. By default, this signal is ignored.
    pub fn io() -> Self {
        Self(libc::SIGIO)
    }

    /// Represents the SIGPIPE signal.
    ///
    /// On Unix systems this signal is sent when the process attempts to write
    /// to a pipe which has no reader. By default, the process is terminated by
    /// this signal.
    pub fn pipe() -> Self {
        Self(libc::SIGPIPE)
    }

    /// Represents the SIGQUIT signal.
    ///
    /// On Unix systems this signal is sent to issue a shutdown of the
    /// process, after which the OS will dump the process core.
    /// By default, the process is terminated by this signal.
    pub fn quit() -> Self {
        Self(libc::SIGQUIT)
    }

    /// Represents the SIGTERM signal.
    ///
    /// On Unix systems this signal is sent to issue a shutdown of the
    /// process. By default, the process is terminated by this signal.
    pub fn terminate() -> Self {
        Self(libc::SIGTERM)
    }

    /// Represents the SIGUSR1 signal.
    ///
    /// On Unix systems this is a user defined signal.
    /// By default, the process is terminated by this signal.
    pub fn user_defined1() -> Self {
        Self(libc::SIGUSR1)
    }

    /// Represents the SIGUSR2 signal.
    ///
    /// On Unix systems this is a user defined signal.
    /// By default, the process is terminated by this signal.
    pub fn user_defined2() -> Self {
        Self(libc::SIGUSR2)
    }

    /// Represents the SIGWINCH signal.
    ///
    /// On Unix systems this signal is sent when the terminal window is resized.
    /// By default, this signal is ignored.
    pub fn window_change() -> Self {
        Self(libc::SIGWINCH)
    }
}

pub(crate) struct SignalInfo {
    event_info: EventInfo,
    init: Once,
    initialized: AtomicBool,
}

impl Default for SignalInfo {
    fn default() -> SignalInfo {
        SignalInfo {
            event_info: Default::default(),
            init: Once::new(),
            initialized: AtomicBool::new(false),
        }
    }
}

/// Our global signal handler for all signals registered by this module.
///
/// The purpose of this signal handler is to primarily:
///
/// 1. Flag that our specific signal was received (e.g. store an atomic flag)
/// 2. Wake up the driver by writing a byte to a pipe
///
/// Those two operations should both be async-signal safe.
fn action(globals: Pin<&'static Globals>, signal: c_int) {
    globals.record_event(signal as EventId);

    // Send a wakeup, ignore any errors (anything reasonably possible is
    // full pipe and then it will wake up anyway).
    let mut sender = &globals.sender;
    drop(sender.write(&[1]));
}

/// Enables this module to receive signal notifications for the `signal`
/// provided.
///
/// This will register the signal handler if it hasn't already been registered,
/// returning any error along the way if that fails.
fn signal_enable(signal: c_int, handle: Handle) -> io::Result<()> {
    if signal < 0 || signal_hook_registry::FORBIDDEN.contains(&signal) {
        return Err(Error::new(
            ErrorKind::Other,
            format!("Refusing to register signal {}", signal),
        ));
    }

    // Check that we have a signal driver running
    handle.check_inner()?;

    let globals = globals();
    let siginfo = match globals.storage().get(signal as EventId) {
        Some(slot) => slot,
        None => return Err(io::Error::new(io::ErrorKind::Other, "signal too large")),
    };
    let mut registered = Ok(());
    siginfo.init.call_once(|| {
        registered = unsafe {
            signal_hook_registry::register(signal, move || action(globals, signal)).map(|_| ())
        };
        if registered.is_ok() {
            siginfo.initialized.store(true, Ordering::Relaxed);
        }
    });
    registered?;
    // If the call_once failed, it won't be retried on the next attempt to register the signal. In
    // such case it is not run, registered is still `Ok(())`, initialized is still `false`.
    if siginfo.initialized.load(Ordering::Relaxed) {
        Ok(())
    } else {
        Err(Error::new(
            ErrorKind::Other,
            "Failed to register signal handler",
        ))
    }
}

/// A stream of events for receiving a particular type of OS signal.
///
/// In general signal handling on Unix is a pretty tricky topic, and this
/// structure is no exception! There are some important limitations to keep in
/// mind when using `Signal` streams:
///
/// * Signals handling in Unix already necessitates coalescing signals
///   together sometimes. This `Signal` stream is also no exception here in
///   that it will also coalesce signals. That is, even if the signal handler
///   for this process runs multiple times, the `Signal` stream may only return
///   one signal notification. Specifically, before `poll` is called, all
///   signal notifications are coalesced into one item returned from `poll`.
///   Once `poll` has been called, however, a further signal is guaranteed to
///   be yielded as an item.
///
///   Put another way, any element pulled off the returned stream corresponds to
///   *at least one* signal, but possibly more.
///
/// * Signal handling in general is relatively inefficient. Although some
///   improvements are possible in this crate, it's recommended to not plan on
///   having millions of signal channels open.
///
/// If you've got any questions about this feel free to open an issue on the
/// repo! New approaches to alleviate some of these limitations are always
/// appreciated!
///
/// # Caveats
///
/// The first time that a `Signal` instance is registered for a particular
/// signal kind, an OS signal-handler is installed which replaces the default
/// platform behavior when that signal is received, **for the duration of the
/// entire process**.
///
/// For example, Unix systems will terminate a process by default when it
/// receives SIGINT. But, when a `Signal` instance is created to listen for
/// this signal, the next SIGINT that arrives will be translated to a stream
/// event, and the process will continue to execute. **Even if this `Signal`
/// instance is dropped, subsequent SIGINT deliveries will end up captured by
/// Tokio, and the default platform behavior will NOT be reset**.
///
/// Thus, applications should take care to ensure the expected signal behavior
/// occurs as expected after listening for specific signals.
///
/// # Examples
///
/// Wait for SIGHUP
///
/// ```rust,no_run
/// use tokio::signal::unix::{signal, SignalKind};
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     // An infinite stream of hangup signals.
///     let mut stream = signal(SignalKind::hangup())?;
///
///     // Print whenever a HUP signal is received
///     loop {
///         stream.recv().await;
///         println!("got signal HUP");
///     }
/// }
/// ```
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Signal {
    rx: Receiver<()>,
}

/// Creates a new stream which will receive notifications when the current
/// process receives the specified signal `kind`.
///
/// This function will create a new stream which binds to the default reactor.
/// The `Signal` stream is an infinite stream which will receive
/// notifications whenever a signal is received. More documentation can be
/// found on `Signal` itself, but to reiterate:
///
/// * Signals may be coalesced beyond what the kernel already does.
/// * Once a signal handler is registered with the process the underlying
///   libc signal handler is never unregistered.
///
/// A `Signal` stream can be created for a particular signal number
/// multiple times. When a signal is received then all the associated
/// channels will receive the signal notification.
///
/// # Errors
///
/// * If the lower-level C functions fail for some reason.
/// * If the previous initialization of this specific signal failed.
/// * If the signal is one of
///   [`signal_hook::FORBIDDEN`](fn@signal_hook_registry::register#panics)
pub fn signal(kind: SignalKind) -> io::Result<Signal> {
    signal_with_handle(kind, Handle::current())
}

pub(crate) fn signal_with_handle(kind: SignalKind, handle: Handle) -> io::Result<Signal> {
    let signal = kind.0;

    // Turn the signal delivery on once we are ready for it
    signal_enable(signal, handle)?;

    // One wakeup in a queue is enough, no need for us to buffer up any
    // more.
    let (tx, rx) = channel(1);
    globals().register_listener(signal as EventId, tx);

    Ok(Signal { rx })
}

impl Signal {
    /// Receives the next signal notification event.
    ///
    /// `None` is returned if no more events can be received by this stream.
    ///
    /// # Examples
    ///
    /// Wait for SIGHUP
    ///
    /// ```rust,no_run
    /// use tokio::signal::unix::{signal, SignalKind};
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
    ///     // An infinite stream of hangup signals.
    ///     let mut stream = signal(SignalKind::hangup())?;
    ///
    ///     // Print whenever a HUP signal is received
    ///     loop {
    ///         stream.recv().await;
    ///         println!("got signal HUP");
    ///     }
    /// }
    /// ```
    pub async fn recv(&mut self) -> Option<()> {
        use crate::future::poll_fn;
        poll_fn(|cx| self.poll_recv(cx)).await
    }

    /// Polls to receive the next signal notification event, outside of an
    /// `async` context.
    ///
    /// This method returns:
    ///
    ///  * `Poll::Pending` if no signals are available but the channel is not
    ///    closed.
    ///  * `Poll::Ready(Some(()))` if a signal is available.
    ///  * `Poll::Ready(None)` if the channel has been closed and all signals
    ///    sent before it was closed have been received.
    ///
    /// # Examples
    ///
    /// Polling from a manually implemented future
    ///
    /// ```rust,no_run
    /// use std::pin::Pin;
    /// use std::future::Future;
    /// use std::task::{Context, Poll};
    /// use tokio::signal::unix::Signal;
    ///
    /// struct MyFuture {
    ///     signal: Signal,
    /// }
    ///
    /// impl Future for MyFuture {
    ///     type Output = Option<()>;
    ///
    ///     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
    ///         println!("polling MyFuture");
    ///         self.signal.poll_recv(cx)
    ///     }
    /// }
    /// ```
    pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
        self.rx.poll_recv(cx)
    }

    /// Try to receive a signal notification without blocking or registering a waker.
    pub(crate) fn try_recv(&mut self) -> Result<(), TryRecvError> {
        self.rx.try_recv()
    }
}

// Work around for abstracting streams internally
pub(crate) trait InternalStream: Unpin {
    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>>;
    fn try_recv(&mut self) -> Result<(), TryRecvError>;
}

impl InternalStream for Signal {
    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
        self.poll_recv(cx)
    }

    fn try_recv(&mut self) -> Result<(), TryRecvError> {
        self.try_recv()
    }
}

pub(crate) fn ctrl_c() -> io::Result<Signal> {
    signal(SignalKind::interrupt())
}

#[cfg(all(test, not(loom)))]
mod tests {
    use super::*;

    #[test]
    fn signal_enable_error_on_invalid_input() {
        signal_enable(-1, Handle::default()).unwrap_err();
    }

    #[test]
    fn signal_enable_error_on_forbidden_input() {
        signal_enable(signal_hook_registry::FORBIDDEN[0], Handle::default()).unwrap_err();
    }
}
<|endoftext|>