use super::*;
use std::sync::Arc;
#[derive(Debug)]
pub struct TokioRuntime;
struct TokioJoinHandle(::tokio::task::JoinHandle<()>);
impl super::JoinHandleInner for TokioJoinHandle {
fn detach(&self) {
}
fn abort(&self) {
self.0.abort();
}
fn is_finished(&self) -> bool {
self.0.is_finished()
}
}
impl Runtime for TokioRuntime {
fn spawn(&self, future: Pin<Box<dyn Future<Output = ()> + Send>>) -> super::JoinHandle {
let handle = ::tokio::spawn(future);
super::JoinHandle {
inner: Box::new(TokioJoinHandle(handle)),
}
}
fn wrap_udp_socket(&self, sock: std::net::UdpSocket) -> io::Result<Arc<dyn AsyncUdpSocket>> {
sock.set_nonblocking(true)?;
Ok(Arc::new(UdpSocket {
io: Arc::new(::tokio::net::UdpSocket::from_std(sock)?),
}))
}
fn wrap_tcp_listener(
&self,
listener: std::net::TcpListener,
) -> io::Result<Arc<dyn AsyncTcpListener>> {
listener.set_nonblocking(true)?;
Ok(Arc::new(TcpListener {
io: ::tokio::net::TcpListener::from_std(listener)?,
}))
}
fn connect_tcp<'a>(
&'a self,
remote_addr: SocketAddr,
) -> Pin<Box<dyn Future<Output = io::Result<Arc<dyn AsyncTcpStream>>> + Send + 'a>> {
Box::pin(async move {
let stream = ::tokio::net::TcpStream::connect(remote_addr).await?;
let local_addr = stream.local_addr()?;
let peer_addr = stream.peer_addr()?;
let (read_half, write_half) = stream.into_split();
Ok(Arc::new(TcpStream {
read_half,
write_half,
local_addr,
peer_addr,
}) as Arc<dyn AsyncTcpStream>)
})
}
}
#[derive(Debug, Clone)]
struct UdpSocket {
io: Arc<::tokio::net::UdpSocket>,
}
impl AsyncUdpSocket for UdpSocket {
fn send_to<'a>(
&'a self,
buf: &'a [u8],
target: SocketAddr,
) -> Pin<Box<dyn Future<Output = io::Result<usize>> + Send + 'a>> {
Box::pin(async move { self.io.send_to(buf, target).await })
}
fn recv_from<'a>(
&'a self,
buf: &'a mut [u8],
) -> Pin<Box<dyn Future<Output = io::Result<(usize, SocketAddr)>> + Send + 'a>> {
Box::pin(async move { self.io.recv_from(buf).await })
}
fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.local_addr()
}
}
#[derive(Debug)]
struct TcpListener {
io: ::tokio::net::TcpListener,
}
impl AsyncTcpListener for TcpListener {
fn accept<'a>(
&'a self,
) -> Pin<Box<dyn Future<Output = io::Result<(Arc<dyn AsyncTcpStream>, SocketAddr)>> + Send + 'a>>
{
Box::pin(async move {
let (stream, addr) = self.io.accept().await?;
let local_addr = stream.local_addr()?;
let peer_addr = stream.peer_addr()?;
let (read_half, write_half) = stream.into_split();
Ok((
Arc::new(TcpStream {
read_half,
write_half,
local_addr,
peer_addr,
}) as Arc<dyn AsyncTcpStream>,
addr,
))
})
}
fn local_addr(&self) -> io::Result<SocketAddr> {
self.io.local_addr()
}
}
#[derive(Debug)]
struct TcpStream {
read_half: ::tokio::net::tcp::OwnedReadHalf,
write_half: ::tokio::net::tcp::OwnedWriteHalf,
local_addr: SocketAddr,
peer_addr: SocketAddr,
}
impl AsyncTcpStream for TcpStream {
fn read<'a, 'b>(
&'a self,
buf: &'b mut [u8],
) -> Pin<Box<dyn Future<Output = io::Result<usize>> + Send + 'b>>
where
'a: 'b,
{
Box::pin(async move {
loop {
self.read_half.readable().await?;
match self.read_half.try_read(buf) {
Ok(n) => return Ok(n),
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
Err(e) => return Err(e),
}
}
})
}
fn write_all<'a, 'b>(
&'a self,
buf: &'b [u8],
) -> Pin<Box<dyn Future<Output = io::Result<()>> + Send + 'b>>
where
'a: 'b,
{
Box::pin(async move {
let mut remaining = buf;
while !remaining.is_empty() {
self.write_half.writable().await?;
match self.write_half.try_write(remaining) {
Ok(0) => {
return Err(io::Error::new(
io::ErrorKind::WriteZero,
"failed to write any bytes",
));
}
Ok(n) => remaining = &remaining[n..],
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
Err(e) => return Err(e),
}
}
Ok(())
})
}
fn local_addr(&self) -> io::Result<SocketAddr> {
Ok(self.local_addr)
}
fn peer_addr(&self) -> io::Result<SocketAddr> {
Ok(self.peer_addr)
}
}
pub async fn sleep(duration: Duration) {
::tokio::time::sleep(duration).await
}
pub struct TokioInterval(::tokio::time::Interval);
impl TokioInterval {
pub async fn tick(&mut self) {
self.0.tick().await;
}
}
pub fn interval(period: Duration) -> TokioInterval {
TokioInterval(::tokio::time::interval(period))
}
pub async fn timeout<F, T>(duration: Duration, future: F) -> Result<T, ()>
where
F: std::future::Future<Output = T>,
{
::tokio::time::timeout(duration, future)
.await
.map_err(|_| ())
}
pub async fn resolve_host(host: &str) -> io::Result<Vec<SocketAddr>> {
::tokio::net::lookup_host(host)
.await
.map(|iter| iter.collect())
}
pub struct TokioMutex<T: ?Sized>(pub Arc<::tokio::sync::Mutex<T>>);
impl<T: ?Sized> Clone for TokioMutex<T> {
fn clone(&self) -> Self {
Self(self.0.clone())
}
}
impl<T> TokioMutex<T> {
pub fn new(value: T) -> Self {
Self(Arc::new(::tokio::sync::Mutex::new(value)))
}
pub async fn lock(&self) -> ::tokio::sync::MutexGuard<'_, T> {
self.0.lock().await
}
}
impl<T: ?Sized + Send> AsyncMutex<T> for TokioMutex<T> {
type Guard<'a>
= ::tokio::sync::MutexGuard<'a, T>
where
T: 'a;
fn lock(&self) -> Pin<Box<dyn Future<Output = Self::Guard<'_>> + Send + '_>> {
Box::pin(self.0.lock())
}
}
pub struct TokioNotify(pub Arc<::tokio::sync::Notify>);
impl Clone for TokioNotify {
fn clone(&self) -> Self {
Self(self.0.clone())
}
}
impl Default for TokioNotify {
fn default() -> Self {
Self::new()
}
}
impl TokioNotify {
pub fn new() -> Self {
Self(Arc::new(::tokio::sync::Notify::new()))
}
pub fn notify_one(&self) {
self.0.notify_one();
}
pub fn notify_waiters(&self) {
self.0.notify_waiters();
}
pub async fn notified(&self) {
self.0.notified().await
}
}
impl AsyncNotify for TokioNotify {
fn notify_one(&self) {
self.0.notify_one();
}
fn notify_waiters(&self) {
self.0.notify_waiters();
}
fn notified(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
Box::pin(self.0.notified())
}
}
pub struct TokioSender<T>(pub ::tokio::sync::mpsc::Sender<T>);
impl<T> Clone for TokioSender<T> {
fn clone(&self) -> Self {
Self(self.0.clone())
}
}
impl<T: Send> TokioSender<T> {
pub async fn send(&self, value: T) -> Result<(), SendError<T>> {
self.0.send(value).await.map_err(|e| SendError(e.0))
}
pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
self.0.try_send(value).map_err(|e| match e {
::tokio::sync::mpsc::error::TrySendError::Full(v) => TrySendError::Full(v),
::tokio::sync::mpsc::error::TrySendError::Closed(v) => TrySendError::Disconnected(v),
})
}
}
impl<T: Send> AsyncSender<T> for TokioSender<T> {
fn send(
&self,
value: T,
) -> Pin<Box<dyn Future<Output = Result<(), SendError<T>>> + Send + '_>> {
Box::pin(async move { self.0.send(value).await.map_err(|e| SendError(e.0)) })
}
fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
self.0.try_send(value).map_err(|e| match e {
::tokio::sync::mpsc::error::TrySendError::Full(v) => TrySendError::Full(v),
::tokio::sync::mpsc::error::TrySendError::Closed(v) => TrySendError::Disconnected(v),
})
}
}
pub struct TokioReceiver<T>(pub ::tokio::sync::mpsc::Receiver<T>);
impl<T: Send> TokioReceiver<T> {
pub async fn recv(&mut self) -> Option<T> {
self.0.recv().await
}
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
self.0.try_recv().map_err(|e| match e {
::tokio::sync::mpsc::error::TryRecvError::Empty => TryRecvError::Empty,
::tokio::sync::mpsc::error::TryRecvError::Disconnected => TryRecvError::Disconnected,
})
}
}
impl<T: Send> AsyncReceiver<T> for TokioReceiver<T> {
fn recv(&mut self) -> Pin<Box<dyn Future<Output = Option<T>> + Send + '_>> {
Box::pin(self.0.recv())
}
fn try_recv(&mut self) -> Result<T, TryRecvError> {
self.0.try_recv().map_err(|e| match e {
::tokio::sync::mpsc::error::TryRecvError::Empty => TryRecvError::Empty,
::tokio::sync::mpsc::error::TryRecvError::Disconnected => TryRecvError::Disconnected,
})
}
}
pub fn channel<T: Send>(capacity: usize) -> (TokioSender<T>, TokioReceiver<T>) {
let (tx, rx) = ::tokio::sync::mpsc::channel(capacity);
(TokioSender(tx), TokioReceiver(rx))
}
#[derive(Clone)]
pub struct TokioBroadcastSender<T>(pub ::tokio::sync::broadcast::Sender<T>);
impl<T: Send + Clone + 'static> TokioBroadcastSender<T> {
pub fn send(&self, value: T) -> Result<usize, super::BroadcastSendError<T>> {
self.0
.send(value)
.map_err(|e| super::BroadcastSendError(e.0))
}
pub fn subscribe(&self) -> TokioBroadcastReceiver<T> {
TokioBroadcastReceiver(self.0.subscribe())
}
pub fn receiver_count(&self) -> usize {
self.0.receiver_count()
}
}
pub struct TokioBroadcastReceiver<T>(pub ::tokio::sync::broadcast::Receiver<T>);
impl<T: Send + Clone + 'static> TokioBroadcastReceiver<T> {
pub async fn recv(&mut self) -> Result<T, super::BroadcastRecvError> {
self.0.recv().await.map_err(|e| match e {
::tokio::sync::broadcast::error::RecvError::Closed => super::BroadcastRecvError::Closed,
::tokio::sync::broadcast::error::RecvError::Lagged(n) => {
super::BroadcastRecvError::Lagged(n)
}
})
}
}
pub fn broadcast_channel<T: Send + Clone + 'static>(capacity: usize) -> TokioBroadcastSender<T> {
let (tx, _) = ::tokio::sync::broadcast::channel(capacity);
TokioBroadcastSender(tx)
}
pub fn block_on<F: std::future::Future>(future: F) -> F::Output {
::tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.expect("failed to build tokio runtime")
.block_on(future)
}