mirrors/ntex
1
0
Fork 0
mirror of https://github.com/ntex-rs/ntex.git synced 2025-04-04 21:37:58 +03:00
Code Issues Projects Releases Packages Wiki Activity

wip

Browse source
This commit is contained in:
Nikolay Kim 2024-09-10 19:06:11 +05:00
parent eecfcc6d36
commit f27996273c
9 changed files with 563 additions and 1761 deletions
Download patch file Download diff file Expand all files Collapse all files

460
ntex-async-std/src/io.rs
View file

@ -1,20 +1,22 @@
use std::future::{poll_fn, Future}; use std::{any, cell::RefCell, future::poll_fn, io, pin::Pin, task::Context, task::Poll};
use std::{any, cell::RefCell, io, pin::Pin, task::Context, task::Poll};
use async_std::io::{Read, Write}; use async_std::io::{Read as ARead, Write as AWrite};
use ntex_bytes::{Buf, BufMut, BytesVec}; use ntex_bytes::{Buf, BufMut, BytesVec};
use ntex_io::{types, Handle, IoStream, ReadContext, WriteContext, WriteStatus}; use ntex_io::{types, Handle, IoStream, ReadContext, WriteContext};
use ntex_util::{ready, time::sleep, time::Sleep}; use ntex_util::{future::lazy, ready};
use crate::TcpStream; use crate::TcpStream;
impl IoStream for TcpStream { impl IoStream for TcpStream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let mut rio = ReadTask(RefCell::new(self.clone())); let mut rio = Read(RefCell::new(self.clone()));
async_std::task::spawn_local(async move { async_std::task::spawn_local(async move {
read.handle(&mut rio).await; read.handle(&mut rio).await;
}); });
async_std::task::spawn_local(WriteTask::new(self.clone(), write)); let mut wio = Write(RefCell::new(self.clone()));
async_std::task::spawn_local(async move {
write.handle(&mut wio).await;
});
Some(Box::new(self)) Some(Box::new(self))
} }
} }
@ -31,9 +33,9 @@ impl Handle for TcpStream {
} }
/// Read io task /// Read io task
struct ReadTask(RefCell<TcpStream>); struct Read(RefCell<TcpStream>);
impl ntex_io::AsyncRead for ReadTask { impl ntex_io::AsyncRead for Read {
async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) {
// read data from socket // read data from socket
let result = poll_fn(|cx| { let result = poll_fn(|cx| {
@ -45,235 +47,91 @@ impl ntex_io::AsyncRead for ReadTask {
} }
} }
#[derive(Debug)] struct Write(RefCell<TcpStream>);
enum IoWriteState {
Processing(Option<Sleep>),
Shutdown(Sleep, Shutdown),
}
#[derive(Debug)] impl ntex_io::AsyncWrite for Write {
enum Shutdown { #[inline]
None, async fn write(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<()>) {
Stopping(u16), match lazy(|cx| flush_io(&mut self.0.borrow_mut().0, &mut buf, cx)).await {
} Poll::Ready(res) => (buf, res),
Poll::Pending => (buf, Ok(())),
/// Write io task
struct WriteTask {
st: IoWriteState,
io: TcpStream,
state: WriteContext,
}
impl WriteTask {
/// Create new write io task
fn new(io: TcpStream, state: WriteContext) -> Self {
Self {
io,
state,
st: IoWriteState::Processing(None),
} }
} }
}
impl Future for WriteTask { #[inline]
type Output = (); async fn flush(&mut self) -> io::Result<()> {
Ok(())
}
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { #[inline]
let this = self.as_mut().get_mut(); async fn shutdown(&mut self) -> io::Result<()> {
self.0.borrow().0.shutdown(std::net::Shutdown::Both)
match this.st {
IoWriteState::Processing(ref mut delay) => {
match this.state.poll_ready(cx) {
Poll::Ready(WriteStatus::Ready) => {
if let Some(delay) = delay {
if delay.poll_elapsed(cx).is_ready() {
this.state.close(Some(io::Error::new(
io::ErrorKind::TimedOut,
"Operation timedout",
)));
return Poll::Ready(());
}
}
// flush io stream
let io = &mut this.io.0;
match ready!(this.state.with_buf(|buf| flush_io(io, buf, cx))) {
Ok(()) => Poll::Pending,
Err(e) => {
this.state.close(Some(e));
Poll::Ready(())
}
}
}
Poll::Ready(WriteStatus::Timeout(time)) => {
log::trace!("initiate timeout delay for {:?}", time);
if delay.is_none() {
*delay = Some(sleep(time));
}
self.poll(cx)
}
Poll::Ready(WriteStatus::Shutdown(time)) => {
log::trace!("write task is instructed to shutdown");
let timeout = if let Some(delay) = delay.take() {
delay
} else {
sleep(time)
};
this.st = IoWriteState::Shutdown(timeout, Shutdown::None);
self.poll(cx)
}
Poll::Ready(WriteStatus::Terminate) => {
log::trace!("write task is instructed to terminate");
let _ = Pin::new(&mut this.io.0).poll_close(cx);
this.state.close(None);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
IoWriteState::Shutdown(ref mut delay, ref mut st) => {
// close WRITE side and wait for disconnect on read side.
// use disconnect timeout, otherwise it could hang forever.
loop {
match st {
Shutdown::None => {
// flush write buffer
let io = &mut this.io.0;
match this.state.with_buf(|buf| flush_io(io, buf, cx)) {
Poll::Ready(Ok(())) => {
if let Err(e) =
this.io.0.shutdown(std::net::Shutdown::Write)
{
this.state.close(Some(e));
return Poll::Ready(());
}
*st = Shutdown::Stopping(0);
continue;
}
Poll::Ready(Err(err)) => {
log::trace!(
"write task is closed with err during flush, {:?}",
err
);
this.state.close(Some(err));
return Poll::Ready(());
}
Poll::Pending => (),
}
}
Shutdown::Stopping(ref mut count) => {
// read until 0 or err
let mut buf = [0u8; 512];
let io = &mut this.io;
loop {
match Pin::new(&mut io.0).poll_read(cx, &mut buf) {
Poll::Ready(Err(e)) => {
log::trace!("write task is stopped");
this.state.close(Some(e));
return Poll::Ready(());
}
Poll::Ready(Ok(0)) => {
log::trace!("async-std socket is disconnected");
this.state.close(None);
return Poll::Ready(());
}
Poll::Ready(Ok(n)) => {
*count += n as u16;
if *count > 4096 {
log::trace!(
"write task is stopped, too much input"
);
this.state.close(None);
return Poll::Ready(());
}
}
Poll::Pending => break,
}
}
}
}
// disconnect timeout
if delay.poll_elapsed(cx).is_pending() {
return Poll::Pending;
}
log::trace!("write task is stopped after delay");
this.state.close(None);
let _ = Pin::new(&mut this.io.0).poll_close(cx);
return Poll::Ready(());
}
}
}
} }
} }
/// Flush write buffer to underlying I/O stream. /// Flush write buffer to underlying I/O stream.
pub(super) fn flush_io<T: Read + Write + Unpin>( pub(super) fn flush_io<T: ARead + AWrite + Unpin>(
io: &mut T, io: &mut T,
buf: &mut Option<BytesVec>, buf: &mut BytesVec,
cx: &mut Context<'_>, cx: &mut Context<'_>,
) -> Poll<io::Result<()>> { ) -> Poll<io::Result<()>> {
if let Some(buf) = buf { let len = buf.len();
let len = buf.len();
if len != 0 { if len != 0 {
// log::trace!("flushing framed transport: {:?}", buf.len()); // log::trace!("flushing framed transport: {:?}", buf.len());
let mut written = 0; let mut written = 0;
let result = loop { let result = loop {
break match Pin::new(&mut *io).poll_write(cx, &buf[written..]) { break match Pin::new(&mut *io).poll_write(cx, &buf[written..]) {
Poll::Ready(Ok(n)) => { Poll::Ready(Ok(n)) => {
if n == 0 { if n == 0 {
log::trace!("Disconnected during flush, written {}", written); log::trace!("Disconnected during flush, written {}", written);
Poll::Ready(Err(io::Error::new( Poll::Ready(Err(io::Error::new(
io::ErrorKind::WriteZero, io::ErrorKind::WriteZero,
"failed to write frame to transport", "failed to write frame to transport",
))) )))
} else {
written += n;
if written == len {
buf.clear();
Poll::Ready(Ok(()))
} else { } else {
written += n; continue;
if written == len {
buf.clear();
Poll::Ready(Ok(()))
} else {
continue;
}
} }
} }
Poll::Pending => {
// remove written data
buf.advance(written);
Poll::Pending
}
Poll::Ready(Err(e)) => {
log::trace!("Error during flush: {}", e);
Poll::Ready(Err(e))
}
};
};
// log::trace!("flushed {} bytes", written);
// flush
return if written > 0 {
match Pin::new(&mut *io).poll_flush(cx) {
Poll::Ready(Ok(_)) => result,
Poll::Pending => Poll::Pending,
Poll::Ready(Err(e)) => {
log::trace!("error during flush: {}", e);
Poll::Ready(Err(e))
}
} }
} else { Poll::Pending => {
result // remove written data
buf.advance(written);
Poll::Pending
}
Poll::Ready(Err(e)) => {
log::trace!("Error during flush: {}", e);
Poll::Ready(Err(e))
}
}; };
};
// log::trace!("flushed {} bytes", written);
// flush
if written > 0 {
match Pin::new(&mut *io).poll_flush(cx) {
Poll::Ready(Ok(_)) => result,
Poll::Pending => Poll::Pending,
Poll::Ready(Err(e)) => {
log::trace!("error during flush: {}", e);
Poll::Ready(Err(e))
}
}
} else {
result
} }
} else {
Poll::Ready(Ok(()))
} }
Poll::Ready(Ok(()))
} }
pub fn poll_read_buf<T: Read>( pub fn poll_read_buf<T: ARead>(
io: Pin<&mut T>, io: Pin<&mut T>,
cx: &mut Context<'_>, cx: &mut Context<'_>,
buf: &mut BytesVec, buf: &mut BytesVec,
@ -297,19 +155,22 @@ mod unixstream {
impl IoStream for UnixStream { impl IoStream for UnixStream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let mut rio = ReadTask(RefCell::new(self.clone())); let mut rio = Read(RefCell::new(self.clone()));
async_std::task::spawn_local(async move { async_std::task::spawn_local(async move {
read.handle(&mut rio).await; read.handle(&mut rio).await;
}); });
async_std::task::spawn_local(WriteTask::new(self, write)); let mut wio = Write(RefCell::new(self));
async_std::task::spawn_local(async move {
write.handle(&mut wio).await;
});
None None
} }
} }
/// Read io task /// Read io task
struct ReadTask(RefCell<UnixStream>); struct Read(RefCell<UnixStream>);
impl ntex_io::AsyncRead for ReadTask { impl ntex_io::AsyncRead for Read {
async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) {
// read data from socket // read data from socket
let result = poll_fn(|cx| { let result = poll_fn(|cx| {
@ -321,158 +182,25 @@ mod unixstream {
} }
} }
/// Write io task struct Write(RefCell<UnixStream>);
struct WriteTask {
st: IoWriteState,
io: UnixStream,
state: WriteContext,
}
impl WriteTask { impl ntex_io::AsyncWrite for Write {
/// Create new write io task #[inline]
fn new(io: UnixStream, state: WriteContext) -> Self { async fn write(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<()>) {
Self { match lazy(|cx| flush_io(&mut self.0.borrow_mut().0, &mut buf, cx)).await {
io, Poll::Ready(res) => (buf, res),
state, Poll::Pending => (buf, Ok(())),
st: IoWriteState::Processing(None),
} }
} }
}
impl Future for WriteTask { #[inline]
type Output = (); async fn flush(&mut self) -> io::Result<()> {
Ok(())
}
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { #[inline]
let this = self.as_mut().get_mut(); async fn shutdown(&mut self) -> io::Result<()> {
self.0.borrow().0.shutdown(std::net::Shutdown::Both)
match this.st {
IoWriteState::Processing(ref mut delay) => {
match this.state.poll_ready(cx) {
Poll::Ready(WriteStatus::Ready) => {
if let Some(delay) = delay {
if delay.poll_elapsed(cx).is_ready() {
this.state.close(Some(io::Error::new(
io::ErrorKind::TimedOut,
"Operation timedout",
)));
return Poll::Ready(());
}
}
// flush io stream
let io = &mut this.io.0;
match ready!(this.state.with_buf(|buf| flush_io(io, buf, cx))) {
Ok(()) => Poll::Pending,
Err(e) => {
this.state.close(Some(e));
Poll::Ready(())
}
}
}
Poll::Ready(WriteStatus::Timeout(time)) => {
log::trace!("initiate timeout delay for {:?}", time);
if delay.is_none() {
*delay = Some(sleep(time));
}
self.poll(cx)
}
Poll::Ready(WriteStatus::Shutdown(time)) => {
log::trace!("write task is instructed to shutdown");
let timeout = if let Some(delay) = delay.take() {
delay
} else {
sleep(time)
};
this.st = IoWriteState::Shutdown(timeout, Shutdown::None);
self.poll(cx)
}
Poll::Ready(WriteStatus::Terminate) => {
log::trace!("write task is instructed to terminate");
let _ = Pin::new(&mut this.io.0).poll_close(cx);
this.state.close(None);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
IoWriteState::Shutdown(ref mut delay, ref mut st) => {
// close WRITE side and wait for disconnect on read side.
// use disconnect timeout, otherwise it could hang forever.
loop {
match st {
Shutdown::None => {
// flush write buffer
let io = &mut this.io.0;
match this.state.with_buf(|buf| flush_io(io, buf, cx)) {
Poll::Ready(Ok(())) => {
if let Err(e) =
this.io.0.shutdown(std::net::Shutdown::Write)
{
this.state.close(Some(e));
return Poll::Ready(());
}
*st = Shutdown::Stopping(0);
continue;
}
Poll::Ready(Err(err)) => {
log::trace!(
"write task is closed with err during flush, {:?}",
err
);
this.state.close(Some(err));
return Poll::Ready(());
}
Poll::Pending => (),
}
}
Shutdown::Stopping(ref mut count) => {
// read until 0 or err
let mut buf = [0u8; 512];
let io = &mut this.io;
loop {
match Pin::new(&mut io.0).poll_read(cx, &mut buf) {
Poll::Ready(Err(e)) => {
log::trace!("write task is stopped");
this.state.close(Some(e));
return Poll::Ready(());
}
Poll::Ready(Ok(0)) => {
log::trace!(
"async-std unix socket is disconnected"
);
this.state.close(None);
return Poll::Ready(());
}
Poll::Ready(Ok(n)) => {
*count += n as u16;
if *count > 4096 {
log::trace!(
"write task is stopped, too much input"
);
this.state.close(None);
return Poll::Ready(());
}
}
Poll::Pending => break,
}
}
}
}
// disconnect timeout
if delay.poll_elapsed(cx).is_pending() {
return Poll::Pending;
}
log::trace!("write task is stopped after delay");
this.state.close(None);
let _ = Pin::new(&mut this.io.0).poll_close(cx);
return Poll::Ready(());
}
}
}
} }
} }
} }

135
ntex-compio/src/io.rs
View file

@ -4,16 +4,15 @@ use compio::buf::{BufResult, IoBuf, IoBufMut, SetBufInit};
use compio::io::{AsyncRead, AsyncWrite}; use compio::io::{AsyncRead, AsyncWrite};
use compio::net::TcpStream; use compio::net::TcpStream;
use ntex_bytes::{Buf, BufMut, BytesVec}; use ntex_bytes::{Buf, BufMut, BytesVec};
use ntex_io::{types, Handle, IoStream, ReadContext, WriteContext, WriteStatus}; use ntex_io::{types, Handle, IoStream, ReadContext, WriteContext};
use ntex_util::{future::select, future::Either, time::sleep};
impl IoStream for crate::TcpStream { impl IoStream for crate::TcpStream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let io = self.0.clone(); let io = self.0.clone();
compio::runtime::spawn(async move { compio::runtime::spawn(async move {
let mut wr_io = io.clone(); let mut wr_io = WriteIo(io.clone());
let wr_task = compio::runtime::spawn(async move { let wr_task = compio::runtime::spawn(async move {
write_task(&mut wr_io, &write).await; write.handle(&mut wr_io).await;
log::debug!("{} Write task is stopped", write.tag()); log::debug!("{} Write task is stopped", write.tag());
}); });
let mut io = ReadIo(io); let mut io = ReadIo(io);
@ -41,9 +40,9 @@ impl IoStream for crate::UnixStream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let io = self.0; let io = self.0;
compio::runtime::spawn(async move { compio::runtime::spawn(async move {
let mut wr_io = io.clone(); let mut wr_io = WriteIo(io.clone());
let wr_task = compio::runtime::spawn(async move { let wr_task = compio::runtime::spawn(async move {
write_task(&mut wr_io, &write).await; write.handle(&mut wr_io).await;
log::debug!("{} Write task is stopped", write.tag()); log::debug!("{} Write task is stopped", write.tag());
}); });
@ -111,105 +110,61 @@ impl SetBufInit for CompioBuf {
} }
} }
struct ReadIo<T: AsyncRead>(T); struct ReadIo<T>(T);
impl<T> ntex_io::AsyncRead for ReadIo<T> impl<T> ntex_io::AsyncRead for ReadIo<T>
where where
T: AsyncRead, T: AsyncRead,
{ {
#[inline]
async fn read(&mut self, buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, buf: BytesVec) -> (BytesVec, io::Result<usize>) {
let BufResult(result, buf) = self.0.read(CompioBuf(buf)).await; let BufResult(result, buf) = self.0.read(CompioBuf(buf)).await;
(buf.0, result) (buf.0, result)
} }
} }
/// Write io task struct WriteIo<T>(T);
async fn write_task<T: AsyncWrite>(mut io: T, state: &WriteContext) {
let mut delay = None;
loop { impl<T> ntex_io::AsyncWrite for WriteIo<T>
let result = if let Some(ref mut sleep) = delay { where
let result = match select(sleep, state.ready()).await { T: AsyncWrite,
Either::Left(_) => { {
state.close(Some(io::Error::new( #[inline]
io::ErrorKind::TimedOut, async fn write(&mut self, buf: BytesVec) -> (BytesVec, io::Result<()>) {
"Operation timedout", let mut buf = CompioBuf(buf);
))); loop {
return; let BufResult(result, buf1) = self.0.write(buf).await;
} buf = buf1;
Either::Right(res) => res,
};
delay = None;
result
} else {
state.ready().await
};
match result { return match result {
WriteStatus::Ready => { Ok(0) => (
// write io stream buf.0,
match write(&mut io, state).await { Err(io::Error::new(
Ok(()) => continue,
Err(e) => {
state.close(Some(e));
}
}
}
WriteStatus::Timeout(time) => {
log::trace!("{}: Initiate timeout delay for {:?}", state.tag(), time);
delay = Some(sleep(time));
continue;
}
WriteStatus::Shutdown(time) => {
log::trace!("{}: Write task is instructed to shutdown", state.tag());
let fut = async {
write(&mut io, state).await?;
io.flush().await?;
io.shutdown().await?;
Ok(())
};
match select(sleep(time), fut).await {
Either::Left(_) => state.close(None),
Either::Right(res) => state.close(res.err()),
}
}
WriteStatus::Terminate => {
log::trace!("{}: Write task is instructed to terminate", state.tag());
state.close(io.shutdown().await.err());
}
}
break;
}
}
// write to io stream
async fn write<T: AsyncWrite>(io: &mut T, state: &WriteContext) -> io::Result<()> {
state
.with_buf_async(|buf| async {
let mut buf = CompioBuf(buf);
loop {
let BufResult(result, buf1) = io.write(buf).await;
buf = buf1;
return match result {
Ok(0) => Err(io::Error::new(
io::ErrorKind::WriteZero, io::ErrorKind::WriteZero,
"failed to write frame to transport", "failed to write frame to transport",
)), )),
Ok(size) => { ),
if buf.0.len() == size { Ok(size) => {
// return io.flush().await; buf.0.advance(size);
state.memory_pool().release_write_buf(buf.0);
Ok(()) if buf.0.is_empty() {
} else { (buf.0, Ok(()))
buf.0.advance(size); } else {
continue; continue;
}
} }
Err(e) => Err(e), }
}; Err(e) => (buf.0, Err(e)),
} };
}) }
.await }
#[inline]
async fn flush(&mut self) -> io::Result<()> {
self.0.flush().await
}
#[inline]
async fn shutdown(&mut self) -> io::Result<()> {
self.0.shutdown().await
}
} }

476
ntex-glommio/src/io.rs
View file

@ -11,24 +11,23 @@ use crate::net_impl::{TcpStream, UnixStream};
impl IoStream for TcpStream { impl IoStream for TcpStream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let mut rio = ReadTask(self.clone()); let mut rio = Read(self.clone());
glommio::spawn_local(async move { glommio::spawn_local(async move { read.handle(&mut rio).await }).detach();
read.handle(&mut rio).await; let mut wio = Write(self.clone());
}) glommio::spawn_local(async move { write.handle(&mut wio).await }).detach();
.detach();
glommio::spawn_local(WriteTask::new(self.clone(), write)).detach();
Some(Box::new(self)) Some(Box::new(self))
} }
} }
impl IoStream for UnixStream { impl IoStream for UnixStream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let mut rio = UnixReadTask(self.clone()); let mut rio = UnixRead(self.clone());
glommio::spawn_local(async move { glommio::spawn_local(async move {
read.handle(&mut rio).await; read.handle(&mut rio).await;
}) })
.detach(); .detach();
glommio::spawn_local(UnixWriteTask::new(self, write)).detach(); let mut wio = UnixWrite(self);
glommio::spawn_local(async move { write.handle(&mut wio).await }).detach();
None None
} }
} }
@ -45,9 +44,9 @@ impl Handle for TcpStream {
} }
/// Read io task /// Read io task
struct ReadTask(TcpStream); struct Read(TcpStream);
impl ntex_io::AsyncRead for ReadTask { impl ntex_io::AsyncRead for Read {
async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) {
// read data from socket // read data from socket
let result = poll_fn(|cx| { let result = poll_fn(|cx| {
@ -59,245 +58,88 @@ impl ntex_io::AsyncRead for ReadTask {
} }
} }
enum IoWriteState { struct Write(TcpStream);
Processing(Option<Sleep>),
Shutdown(Sleep, Shutdown),
}
enum Shutdown { impl ntex_io::AsyncWrite for Write {
Flush, #[inline]
Close(Pin<Box<dyn Future<Output = glommio::Result<(), ()>>>>), async fn write(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<()>) {
Stopping(u16), match lazy(|cx| flush_io(&mut *self.0.borrow_mut(), &mut buf, cx)).await {
} Poll::Ready(res) => (buf, res),
Poll::Pending => (buf, Ok(())),
/// Write io task
struct WriteTask {
st: IoWriteState,
io: TcpStream,
state: WriteContext,
}
impl WriteTask {
/// Create new write io task
fn new(io: TcpStream, state: WriteContext) -> Self {
Self {
io,
state,
st: IoWriteState::Processing(None),
} }
} }
}
impl Future for WriteTask { #[inline]
type Output = (); async fn flush(&mut self) -> io::Result<()> {
Ok(())
}
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { #[inline]
let this = self.as_mut().get_mut(); async fn shutdown(&mut self) -> io::Result<()> {
poll_fn(|cx| Pin::new(&mut *self.0.borrow_mut()).poll_close(cx)).await
match this.st {
IoWriteState::Processing(ref mut delay) => {
match this.state.poll_ready(cx) {
Poll::Ready(WriteStatus::Ready) => {
if let Some(delay) = delay {
if delay.poll_elapsed(cx).is_ready() {
this.state.close(Some(io::Error::new(
io::ErrorKind::TimedOut,
"Operation timedout",
)));
return Poll::Ready(());
}
}
// flush io stream
match ready!(this.state.with_buf(|buf| flush_io(
&mut *this.io.0.borrow_mut(),
buf,
cx
))) {
Ok(()) => Poll::Pending,
Err(e) => {
this.state.close(Some(e));
Poll::Ready(())
}
}
}
Poll::Ready(WriteStatus::Timeout(time)) => {
log::trace!("initiate timeout delay for {:?}", time);
if delay.is_none() {
*delay = Some(sleep(time));
}
self.poll(cx)
}
Poll::Ready(WriteStatus::Shutdown(time)) => {
log::trace!("write task is instructed to shutdown");
let timeout = if let Some(delay) = delay.take() {
delay
} else {
sleep(time)
};
this.st = IoWriteState::Shutdown(timeout, Shutdown::Flush);
self.poll(cx)
}
Poll::Ready(WriteStatus::Terminate) => {
log::trace!("write task is instructed to terminate");
let _ = Pin::new(&mut *this.io.0.borrow_mut()).poll_close(cx);
this.state.close(None);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
IoWriteState::Shutdown(ref mut delay, ref mut st) => {
// close WRITE side and wait for disconnect on read side.
// use disconnect timeout, otherwise it could hang forever.
loop {
match st {
Shutdown::Flush => {
// flush write buffer
let mut io = this.io.0.borrow_mut();
match this.state.with_buf(|buf| flush_io(&mut *io, buf, cx)) {
Poll::Ready(Ok(())) => {
let io = this.io.clone();
#[allow(clippy::await_holding_refcell_ref)]
let fut = Box::pin(async move {
io.0.borrow()
.shutdown(std::net::Shutdown::Write)
.await
});
*st = Shutdown::Close(fut);
continue;
}
Poll::Ready(Err(err)) => {
log::trace!(
"write task is closed with err during flush, {:?}",
err
);
this.state.close(Some(err));
return Poll::Ready(());
}
Poll::Pending => (),
}
}
Shutdown::Close(ref mut fut) => {
if ready!(fut.poll(cx)).is_err() {
this.state.close(None);
return Poll::Ready(());
}
*st = Shutdown::Stopping(0);
continue;
}
Shutdown::Stopping(ref mut count) => {
// read until 0 or err
let mut buf = [0u8; 512];
let io = &mut this.io;
loop {
match Pin::new(&mut *io.0.borrow_mut())
.poll_read(cx, &mut buf)
{
Poll::Ready(Err(e)) => {
log::trace!("write task is stopped");
this.state.close(Some(e));
return Poll::Ready(());
}
Poll::Ready(Ok(0)) => {
log::trace!("glommio socket is disconnected");
this.state.close(None);
return Poll::Ready(());
}
Poll::Ready(Ok(n)) => {
*count += n as u16;
if *count > 4096 {
log::trace!(
"write task is stopped, too much input"
);
this.state.close(None);
return Poll::Ready(());
}
}
Poll::Pending => break,
}
}
}
}
// disconnect timeout
if delay.poll_elapsed(cx).is_pending() {
return Poll::Pending;
}
log::trace!("write task is stopped after delay");
this.state.close(None);
let _ = Pin::new(&mut *this.io.0.borrow_mut()).poll_close(cx);
return Poll::Ready(());
}
}
}
} }
} }
/// Flush write buffer to underlying I/O stream. /// Flush write buffer to underlying I/O stream.
pub(super) fn flush_io<T: AsyncRead + AsyncWrite + Unpin>( pub(super) fn flush_io<T: AsyncRead + AsyncWrite + Unpin>(
io: &mut T, io: &mut T,
buf: &mut Option<BytesVec>, buf: &mut BytesVec,
cx: &mut Context<'_>, cx: &mut Context<'_>,
) -> Poll<io::Result<()>> { ) -> Poll<io::Result<()>> {
if let Some(buf) = buf { let len = buf.len();
let len = buf.len();
if len != 0 { if len != 0 {
// log::trace!("flushing framed transport: {:?}", buf.len()); // log::trace!("flushing framed transport: {:?}", buf.len());
let mut written = 0; let mut written = 0;
let result = loop { let result = loop {
break match Pin::new(&mut *io).poll_write(cx, &buf[written..]) { break match Pin::new(&mut *io).poll_write(cx, &buf[written..]) {
Poll::Ready(Ok(n)) => { Poll::Ready(Ok(n)) => {
if n == 0 { if n == 0 {
log::trace!("Disconnected during flush, written {}", written); log::trace!("Disconnected during flush, written {}", written);
Poll::Ready(Err(io::Error::new( Poll::Ready(Err(io::Error::new(
io::ErrorKind::WriteZero, io::ErrorKind::WriteZero,
"failed to write frame to transport", "failed to write frame to transport",
))) )))
} else {
written += n;
if written == len {
buf.clear();
Poll::Ready(Ok(()))
} else { } else {
written += n; continue;
if written == len {
buf.clear();
Poll::Ready(Ok(()))
} else {
continue;
}
} }
} }
Poll::Pending => {
// remove written data
buf.advance(written);
Poll::Pending
}
Poll::Ready(Err(e)) => {
log::trace!("Error during flush: {}", e);
Poll::Ready(Err(e))
}
};
};
log::trace!("flushed {} bytes", written);
// flush
return if written > 0 {
match Pin::new(&mut *io).poll_flush(cx) {
Poll::Ready(Ok(_)) => result,
Poll::Pending => Poll::Pending,
Poll::Ready(Err(e)) => {
log::trace!("error during flush: {}", e);
Poll::Ready(Err(e))
}
} }
} else { Poll::Pending => {
result // remove written data
buf.advance(written);
Poll::Pending
}
Poll::Ready(Err(e)) => {
log::trace!("Error during flush: {}", e);
Poll::Ready(Err(e))
}
}; };
} };
// log::trace!("flushed {} bytes", written);
// flush
return if written > 0 {
match Pin::new(&mut *io).poll_flush(cx) {
Poll::Ready(Ok(_)) => result,
Poll::Pending => Poll::Pending,
Poll::Ready(Err(e)) => {
log::trace!("error during flush: {}", e);
Poll::Ready(Err(e))
}
}
} else {
result
};
} else {
Poll::Ready(Ok(()))
} }
Poll::Ready(Ok(()))
} }
pub fn poll_read_buf<T: AsyncRead>( pub fn poll_read_buf<T: AsyncRead>(
@ -317,10 +159,9 @@ pub fn poll_read_buf<T: AsyncRead>(
Poll::Ready(Ok(n)) Poll::Ready(Ok(n))
} }
/// Read io task struct UnixRead(UnixStream);
struct UnixReadTask(UnixStream);
impl ntex_io::AsyncRead for UnixReadTask { impl ntex_io::AsyncRead for UnixRead {
async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) {
// read data from socket // read data from socket
let result = poll_fn(|cx| { let result = poll_fn(|cx| {
@ -332,169 +173,24 @@ impl ntex_io::AsyncRead for UnixReadTask {
} }
} }
/// Write io task struct UnixWrite(UnixStream);
struct UnixWriteTask {
st: IoWriteState,
io: UnixStream,
state: WriteContext,
}
impl UnixWriteTask { impl ntex_io::AsyncWrite for UnixWrite {
/// Create new write io task #[inline]
fn new(io: UnixStream, state: WriteContext) -> Self { async fn write(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<()>) {
Self { match lazy(|cx| flush_io(&mut *self.0.borrow_mut(), &mut buf, cx)).await {
io, Poll::Ready(res) => (buf, res),
state, Poll::Pending => (buf, Ok(())),
st: IoWriteState::Processing(None),
} }
} }
}
impl Future for UnixWriteTask { #[inline]
type Output = (); async fn flush(&mut self) -> io::Result<()> {
Ok(())
}
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { #[inline]
let this = self.as_mut().get_mut(); async fn shutdown(&mut self) -> io::Result<()> {
poll_fn(|cx| Pin::new(&mut *self.0.borrow_mut()).poll_close(cx)).await
match this.st {
IoWriteState::Processing(ref mut delay) => {
match this.state.poll_ready(cx) {
Poll::Ready(WriteStatus::Ready) => {
if let Some(delay) = delay {
if delay.poll_elapsed(cx).is_ready() {
this.state.close(Some(io::Error::new(
io::ErrorKind::TimedOut,
"Operation timedout",
)));
return Poll::Ready(());
}
}
// flush io stream
match ready!(this.state.with_buf(|buf| flush_io(
&mut *this.io.0.borrow_mut(),
buf,
cx
))) {
Ok(()) => Poll::Pending,
Err(e) => {
this.state.close(Some(e));
Poll::Ready(())
}
}
}
Poll::Ready(WriteStatus::Timeout(time)) => {
log::trace!("initiate timeout delay for {:?}", time);
if delay.is_none() {
*delay = Some(sleep(time));
}
self.poll(cx)
}
Poll::Ready(WriteStatus::Shutdown(time)) => {
log::trace!("write task is instructed to shutdown");
let timeout = if let Some(delay) = delay.take() {
delay
} else {
sleep(time)
};
this.st = IoWriteState::Shutdown(timeout, Shutdown::Flush);
self.poll(cx)
}
Poll::Ready(WriteStatus::Terminate) => {
log::trace!("write task is instructed to terminate");
let _ = Pin::new(&mut *this.io.0.borrow_mut()).poll_close(cx);
this.state.close(None);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
IoWriteState::Shutdown(ref mut delay, ref mut st) => {
// close WRITE side and wait for disconnect on read side.
// use disconnect timeout, otherwise it could hang forever.
loop {
match st {
Shutdown::Flush => {
// flush write buffer
let mut io = this.io.0.borrow_mut();
match this.state.with_buf(|buf| flush_io(&mut *io, buf, cx)) {
Poll::Ready(Ok(())) => {
let io = this.io.clone();
#[allow(clippy::await_holding_refcell_ref)]
let fut = Box::pin(async move {
io.0.borrow()
.shutdown(std::net::Shutdown::Write)
.await
});
*st = Shutdown::Close(fut);
continue;
}
Poll::Ready(Err(err)) => {
log::trace!(
"write task is closed with err during flush, {:?}",
err
);
this.state.close(Some(err));
return Poll::Ready(());
}
Poll::Pending => (),
}
}
Shutdown::Close(ref mut fut) => {
if ready!(fut.poll(cx)).is_err() {
this.state.close(None);
return Poll::Ready(());
}
*st = Shutdown::Stopping(0);
continue;
}
Shutdown::Stopping(ref mut count) => {
// read until 0 or err
let mut buf = [0u8; 512];
let io = &mut this.io;
loop {
match Pin::new(&mut *io.0.borrow_mut())
.poll_read(cx, &mut buf)
{
Poll::Ready(Err(e)) => {
log::trace!("write task is stopped");
this.state.close(Some(e));
return Poll::Ready(());
}
Poll::Ready(Ok(0)) => {
log::trace!("glommio unix socket is disconnected");
this.state.close(None);
return Poll::Ready(());
}
Poll::Ready(Ok(n)) => {
*count += n as u16;
if *count > 4096 {
log::trace!(
"write task is stopped, too much input"
);
this.state.close(None);
return Poll::Ready(());
}
}
Poll::Pending => break,
}
}
}
}
// disconnect timeout
if delay.poll_elapsed(cx).is_pending() {
return Poll::Pending;
}
log::trace!("write task is stopped after delay");
this.state.close(None);
let _ = Pin::new(&mut *this.io.0.borrow_mut()).poll_close(cx);
return Poll::Ready(());
}
}
}
} }
} }

32
ntex-io/src/buf.rs
View file

@ -152,27 +152,6 @@ impl Stack {
} }
} }
pub(crate) fn with_read_source<F, R>(&self, io: &IoRef, f: F) -> R
where
F: FnOnce(&mut BytesVec) -> R,
{
let item = self.get_last_level();
let mut rb = item.0.take();
if rb.is_none() {
rb = Some(io.memory_pool().get_read_buf());
}
let result = f(rb.as_mut().unwrap());
if let Some(b) = rb {
if b.is_empty() {
io.memory_pool().release_read_buf(b);
} else {
item.0.set(Some(b));
}
}
result
}
pub(crate) fn with_read_destination<F, R>(&self, io: &IoRef, f: F) -> R pub(crate) fn with_read_destination<F, R>(&self, io: &IoRef, f: F) -> R
where where
F: FnOnce(&mut BytesVec) -> R, F: FnOnce(&mut BytesVec) -> R,
@ -226,6 +205,17 @@ impl Stack {
self.get_last_level().1.take() self.get_last_level().1.take()
} }
pub(crate) fn set_write_destination(&self, buf: BytesVec) -> Option<BytesVec> {
let b = self.get_last_level().1.take();
if b.is_some() {
self.get_last_level().1.set(b);
Some(buf)
} else {
self.get_last_level().1.set(Some(buf));
None
}
}
pub(crate) fn with_write_destination<F, R>(&self, io: &IoRef, f: F) -> R pub(crate) fn with_write_destination<F, R>(&self, io: &IoRef, f: F) -> R
where where
F: FnOnce(&mut Option<BytesVec>) -> R, F: FnOnce(&mut Option<BytesVec>) -> R,

4
ntex-io/src/ioref.rs
View file

@ -41,10 +41,6 @@ impl IoRef {
.intersects(Flags::IO_STOPPING | Flags::IO_STOPPED) .intersects(Flags::IO_STOPPING | Flags::IO_STOPPED)
} }
pub(crate) fn is_io_closed(&self) -> bool {
self.0.flags.get().intersects(Flags::IO_STOPPED)
}
#[inline] #[inline]
/// Check if write back-pressure is enabled /// Check if write back-pressure is enabled
pub fn is_wr_backpressure(&self) -> bool { pub fn is_wr_backpressure(&self) -> bool {

18
ntex-io/src/lib.rs
View file

@ -39,6 +39,19 @@ pub use self::utils::{seal, Decoded};
pub use self::flags::Flags; pub use self::flags::Flags;
#[doc(hidden)] #[doc(hidden)]
pub trait AsyncRead {
async fn read(&mut self, buf: BytesVec) -> (BytesVec, sio::Result<usize>);
}
#[doc(hidden)]
pub trait AsyncWrite {
async fn write(&mut self, buf: BytesVec) -> (BytesVec, sio::Result<()>);
async fn flush(&mut self) -> sio::Result<()>;
async fn shutdown(&mut self) -> sio::Result<()>;
}
/// Status for read task /// Status for read task
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum ReadStatus { pub enum ReadStatus {
@ -46,11 +59,6 @@ pub enum ReadStatus {
Terminate, Terminate,
} }
#[doc(hidden)]
pub trait AsyncRead {
async fn read(&mut self, buf: BytesVec) -> (BytesVec, sio::Result<usize>);
}
/// Status for write task /// Status for write task
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum WriteStatus { pub enum WriteStatus {

400
ntex-io/src/tasks.rs
View file

@ -1,9 +1,9 @@
use std::{future::poll_fn, future::Future, io, task::Context, task::Poll}; use std::{future::poll_fn, io, task::Poll};
use ntex_bytes::{BufMut, BytesVec, PoolRef}; use ntex_bytes::BufMut;
use ntex_util::future::{select, Either}; use ntex_util::{future::select, future::Either, time::sleep};
use crate::{AsyncRead, Flags, IoRef, ReadStatus, WriteStatus}; use crate::{AsyncRead, AsyncWrite, Flags, IoRef, ReadStatus, WriteStatus};
#[derive(Debug)] #[derive(Debug)]
/// Context for io read task /// Context for io read task
@ -20,22 +20,8 @@ impl ReadContext {
self.0.tag() self.0.tag()
} }
#[deprecated]
#[inline]
/// Check readiness for read operations
pub async fn ready(&self) -> ReadStatus {
poll_fn(|cx| self.0.filter().poll_read_ready(cx)).await
}
#[deprecated]
#[inline]
/// Wait when io get closed or preparing for close /// Wait when io get closed or preparing for close
pub async fn wait_for_close(&self) { async fn wait_for_close(&self) {
self.wait_for_close2().await
}
/// Wait when io get closed or preparing for close
async fn wait_for_close2(&self) {
poll_fn(|cx| { poll_fn(|cx| {
let flags = self.0.flags(); let flags = self.0.flags();
@ -52,113 +38,7 @@ impl ReadContext {
.await .await
} }
#[deprecated] /// Handle read io operations
#[inline]
/// Check readiness for read operations
pub fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<ReadStatus> {
self.0.filter().poll_read_ready(cx)
}
#[deprecated]
/// Get read buffer
pub fn with_buf<F>(&self, f: F) -> Poll<()>
where
F: FnOnce(&mut BytesVec, usize, usize) -> Poll<io::Result<()>>,
{
let inner = &self.0 .0;
let (hw, lw) = self.0.memory_pool().read_params().unpack();
let (result, nbytes, total) = inner.buffer.with_read_source(&self.0, |buf| {
let total = buf.len();
// call provided callback
let result = f(buf, hw, lw);
let total2 = buf.len();
let nbytes = if total2 > total { total2 - total } else { 0 };
(result, nbytes, total2)
});
// handle buffer changes
if nbytes > 0 {
let filter = self.0.filter();
let _ = filter
.process_read_buf(&self.0, &inner.buffer, 0, nbytes)
.and_then(|status| {
if status.nbytes > 0 {
// dest buffer has new data, wake up dispatcher
if inner.buffer.read_destination_size() >= hw {
log::trace!(
"{}: Io read buffer is too large {}, enable read back-pressure",
self.0.tag(),
total
);
inner.insert_flags(Flags::BUF_R_READY | Flags::BUF_R_FULL);
} else {
inner.insert_flags(Flags::BUF_R_READY);
if nbytes >= hw {
// read task is paused because of read back-pressure
// but there is no new data in top most read buffer
// so we need to wake up read task to read more data
// otherwise read task would sleep forever
inner.read_task.wake();
}
}
log::trace!(
"{}: New {} bytes available, wakeup dispatcher",
self.0.tag(),
nbytes
);
inner.dispatch_task.wake();
} else {
if nbytes >= hw {
// read task is paused because of read back-pressure
// but there is no new data in top most read buffer
// so we need to wake up read task to read more data
// otherwise read task would sleep forever
inner.read_task.wake();
}
if inner.flags.get().contains(Flags::RD_NOTIFY) {
// in case of "notify" we must wake up dispatch task
// if we read any data from source
inner.dispatch_task.wake();
}
}
// while reading, filter wrote some data
// in that case filters need to process write buffers
// and potentialy wake write task
if status.need_write {
filter.process_write_buf(&self.0, &inner.buffer, 0)
} else {
Ok(())
}
})
.map_err(|err| {
inner.dispatch_task.wake();
inner.io_stopped(Some(err));
inner.insert_flags(Flags::BUF_R_READY);
});
}
match result {
Poll::Ready(Ok(())) => {
inner.io_stopped(None);
Poll::Ready(())
}
Poll::Ready(Err(e)) => {
inner.io_stopped(Some(e));
Poll::Ready(())
}
Poll::Pending => {
if inner.flags.get().contains(Flags::IO_STOPPING_FILTERS) {
shutdown_filters(&self.0);
}
Poll::Pending
}
}
}
/// Get read buffer (async)
pub async fn handle<T>(&self, io: &mut T) pub async fn handle<T>(&self, io: &mut T)
where where
T: AsyncRead, T: AsyncRead,
@ -192,7 +72,7 @@ impl ReadContext {
let total = buf.len(); let total = buf.len();
// call provided callback // call provided callback
let (buf, result) = match select(io.read(buf), self.wait_for_close2()).await { let (buf, result) = match select(io.read(buf), self.wait_for_close()).await {
Either::Left(res) => res, Either::Left(res) => res,
Either::Right(_) => { Either::Right(_) => {
log::trace!("{}: Read io is closed, stop read task", self.tag()); log::trace!("{}: Read io is closed, stop read task", self.tag());
@ -297,128 +177,164 @@ impl WriteContext {
self.0.tag() self.0.tag()
} }
#[inline]
/// Return memory pool for this context
pub fn memory_pool(&self) -> PoolRef {
self.0.memory_pool()
}
#[inline]
/// Check readiness for write operations /// Check readiness for write operations
pub async fn ready(&self) -> WriteStatus { async fn ready(&self) -> WriteStatus {
poll_fn(|cx| self.0.filter().poll_write_ready(cx)).await poll_fn(|cx| self.0.filter().poll_write_ready(cx)).await
} }
#[inline]
/// Check readiness for write operations
pub fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<WriteStatus> {
self.0.filter().poll_write_ready(cx)
}
#[inline]
/// Check if io is closed
pub fn poll_close(&self, cx: &mut Context<'_>) -> Poll<()> {
if self.0.is_io_closed() {
Poll::Ready(())
} else {
self.0 .0.write_task.register(cx.waker());
Poll::Pending
}
}
/// Get write buffer
pub fn with_buf<F>(&self, f: F) -> Poll<io::Result<()>>
where
F: FnOnce(&mut Option<BytesVec>) -> Poll<io::Result<()>>,
{
let inner = &self.0 .0;
// call provided callback
let (result, len) = inner.buffer.with_write_destination(&self.0, |buf| {
let result = f(buf);
(result, buf.as_ref().map(|b| b.len()).unwrap_or(0))
});
// if write buffer is smaller than high watermark value, turn off back-pressure
let mut flags = inner.flags.get();
if len == 0 {
if flags.is_waiting_for_write() {
flags.waiting_for_write_is_done();
inner.dispatch_task.wake();
}
} else if flags.contains(Flags::BUF_W_BACKPRESSURE)
&& len < inner.pool.get().write_params_high() << 1
{
flags.remove(Flags::BUF_W_BACKPRESSURE);
inner.dispatch_task.wake();
}
match result {
Poll::Pending => flags.remove(Flags::WR_PAUSED),
Poll::Ready(Ok(())) => flags.insert(Flags::WR_PAUSED),
Poll::Ready(Err(_)) => {}
}
inner.flags.set(flags);
result
}
/// Get write buffer (async)
pub async fn with_buf_async<F, R>(&self, f: F) -> io::Result<()>
where
F: FnOnce(BytesVec) -> R,
R: Future<Output = io::Result<()>>,
{
let inner = &self.0 .0;
// running
let mut flags = inner.flags.get();
if flags.contains(Flags::WR_PAUSED) {
flags.remove(Flags::WR_PAUSED);
inner.flags.set(flags);
}
// buffer
let buf = inner.buffer.get_write_destination();
// call provided callback
let result = if let Some(buf) = buf {
if !buf.is_empty() {
f(buf).await
} else {
Ok(())
}
} else {
Ok(())
};
// if write buffer is smaller than high watermark value, turn off back-pressure
let mut flags = inner.flags.get();
let len = inner.buffer.write_destination_size();
if len == 0 {
if flags.is_waiting_for_write() {
flags.waiting_for_write_is_done();
inner.dispatch_task.wake();
}
flags.insert(Flags::WR_PAUSED);
inner.flags.set(flags);
} else if flags.contains(Flags::BUF_W_BACKPRESSURE)
&& len < inner.pool.get().write_params_high() << 1
{
flags.remove(Flags::BUF_W_BACKPRESSURE);
inner.flags.set(flags);
inner.dispatch_task.wake();
}
result
}
#[inline]
/// Indicate that write io task is stopped /// Indicate that write io task is stopped
pub fn close(&self, err: Option<io::Error>) { fn close(&self, err: Option<io::Error>) {
self.0 .0.io_stopped(err); self.0 .0.io_stopped(err);
} }
/// Handle write io operations
pub async fn handle<T>(&self, io: &mut T)
where
T: AsyncWrite,
{
let inner = &self.0 .0;
let mut delay = None;
let mut buf = None;
loop {
// check readiness
let result = if let Some(ref mut sleep) = delay {
let result = match select(sleep, self.ready()).await {
Either::Left(_) => {
self.close(Some(io::Error::new(
io::ErrorKind::TimedOut,
"Operation timedout",
)));
return;
}
Either::Right(res) => res,
};
delay = None;
result
} else {
self.ready().await
};
// running
let mut flags = inner.flags.get();
if flags.contains(Flags::WR_PAUSED) {
flags.remove(Flags::WR_PAUSED);
inner.flags.set(flags);
}
// call provided callback
match result {
WriteStatus::Ready => {
// write io stream
let (buf_result, result) = if let Some(b) = buf.take() {
io.write(b).await
} else if let Some(b) = inner.buffer.get_write_destination() {
io.write(b).await
} else {
// nothing to write, wait for wakeup
if flags.is_waiting_for_write() {
flags.waiting_for_write_is_done();
inner.dispatch_task.wake();
}
flags.insert(Flags::WR_PAUSED);
if flags.contains(Flags::BUF_W_BACKPRESSURE) {
flags.remove(
Flags::BUF_W_BACKPRESSURE | Flags::BUF_W_MUST_FLUSH,
);
inner.dispatch_task.wake();
}
inner.flags.set(flags);
continue;
};
match result {
Ok(_) => {
let len = if buf_result.is_empty() {
// return io.flush().await;
self.0.memory_pool().release_write_buf(buf_result);
0
} else if let Some(b) =
inner.buffer.set_write_destination(buf_result)
{
// write buffer is already set, we have to write
// current buffer
let l = b.len();
buf = Some(b);
l
} else {
0
};
// if write buffer is smaller than high watermark value, turn off back-pressure
let mut flags = inner.flags.get();
let len = len + inner.buffer.write_destination_size();
if len == 0 {
if flags.is_waiting_for_write() {
flags.waiting_for_write_is_done();
inner.dispatch_task.wake();
}
flags.insert(Flags::WR_PAUSED);
inner.flags.set(flags);
} else if flags.contains(Flags::BUF_W_BACKPRESSURE)
&& len < inner.pool.get().write_params_high() << 1
{
flags.remove(Flags::BUF_W_BACKPRESSURE);
inner.flags.set(flags);
inner.dispatch_task.wake();
}
continue;
}
Err(e) => self.close(Some(e)),
}
}
WriteStatus::Timeout(time) => {
log::trace!("{}: Initiate timeout delay for {:?}", self.tag(), time);
delay = Some(sleep(time));
continue;
}
WriteStatus::Shutdown(time) => {
log::trace!("{}: Write task is instructed to shutdown", self.tag());
let fut = async {
// write io stream
loop {
buf = if let Some(b) = buf {
let (b, result) = io.write(b).await;
result?;
if !b.is_empty() {
Some(b)
} else {
None
}
} else {
inner.buffer.get_write_destination()
};
if buf.is_none() {
break;
}
}
io.flush().await?;
io.shutdown().await?;
Ok(())
};
match select(sleep(time), fut).await {
Either::Left(_) => self.close(None),
Either::Right(res) => self.close(res.err()),
}
}
WriteStatus::Terminate => {
log::trace!("{}: Write task is instructed to terminate", self.tag());
self.close(io.shutdown().await.err());
}
}
return;
}
}
} }
fn shutdown_filters(io: &IoRef) { fn shutdown_filters(io: &IoRef) {

266
ntex-io/src/testing.rs
View file

@ -1,14 +1,13 @@
//! utilities and helpers for testing //! utilities and helpers for testing
#![allow(clippy::let_underscore_future)] #![allow(clippy::let_underscore_future)]
use std::future::{poll_fn, Future};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::task::{ready, Context, Poll, Waker}; use std::task::{Context, Poll, Waker};
use std::{any, cell::RefCell, cmp, fmt, io, mem, net, pin::Pin, rc::Rc}; use std::{any, cell::RefCell, cmp, fmt, future::poll_fn, io, mem, net, rc::Rc};
use ntex_bytes::{Buf, BufMut, Bytes, BytesVec}; use ntex_bytes::{Buf, BufMut, Bytes, BytesVec};
use ntex_util::time::{sleep, Millis, Sleep}; use ntex_util::time::{sleep, Millis};
use crate::{types, Handle, IoStream, ReadContext, WriteContext, WriteStatus}; use crate::{types, Handle, IoStream, ReadContext, WriteContext};
#[derive(Default)] #[derive(Default)]
struct AtomicWaker(Arc<Mutex<RefCell<Option<Waker>>>>); struct AtomicWaker(Arc<Mutex<RefCell<Option<Waker>>>>);
@ -356,14 +355,14 @@ impl IoStream for IoTest {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let io = Rc::new(self); let io = Rc::new(self);
let mut rio = ReadTask(io.clone()); let mut rio = Read(io.clone());
let _ = ntex_util::spawn(async move { let _ = ntex_util::spawn(async move {
read.handle(&mut rio).await; read.handle(&mut rio).await;
}); });
let _ = ntex_util::spawn(WriteTask {
io: io.clone(), let mut wio = Write(io.clone());
state: write, let _ = ntex_util::spawn(async move {
st: IoWriteState::Processing(None), write.handle(&mut wio).await;
}); });
Some(Box::new(io)) Some(Box::new(io))
@ -382,9 +381,9 @@ impl Handle for Rc<IoTest> {
} }
/// Read io task /// Read io task
struct ReadTask(Rc<IoTest>); struct Read(Rc<IoTest>);
impl crate::AsyncRead for ReadTask { impl crate::AsyncRead for Read {
async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) {
// read data from socket // read data from socket
let result = poll_fn(|cx| self.0.poll_read_buf(cx, &mut buf)).await; let result = poll_fn(|cx| self.0.poll_read_buf(cx, &mut buf)).await;
@ -392,205 +391,80 @@ impl crate::AsyncRead for ReadTask {
} }
} }
#[derive(Debug)] /// Write
enum IoWriteState { struct Write(Rc<IoTest>);
Processing(Option<Sleep>),
Shutdown(Option<Sleep>, Shutdown),
}
#[derive(Debug)] impl crate::AsyncWrite for Write {
enum Shutdown { async fn write(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<()>) {
None, let result = poll_fn(|cx| write_io(&self.0, &mut buf, cx)).await;
Flushed,
Stopping,
}
/// Write io task (buf, result)
struct WriteTask { }
st: IoWriteState,
io: Rc<IoTest>,
state: WriteContext,
}
impl Future for WriteTask { async fn flush(&mut self) -> io::Result<()> {
type Output = (); Ok(())
}
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { async fn shutdown(&mut self) -> io::Result<()> {
let this = self.as_mut().get_mut(); // shutdown WRITE side
self.0
match this.st { .local
IoWriteState::Processing(ref mut delay) => { .lock()
match this.state.poll_ready(cx) { .unwrap()
Poll::Ready(WriteStatus::Ready) => { .borrow_mut()
// flush framed instance .flags
match ready!(flush_io(&this.io, &this.state, cx)) { .insert(IoTestFlags::CLOSED);
Ok(()) => Poll::Pending, Ok(())
Err(e) => {
this.state.close(Some(e));
Poll::Ready(())
}
}
}
Poll::Ready(WriteStatus::Timeout(time)) => {
if delay.is_none() {
*delay = Some(sleep(time));
}
self.poll(cx)
}
Poll::Ready(WriteStatus::Shutdown(time)) => {
log::trace!("write task is instructed to shutdown");
let timeout = if let Some(delay) = delay.take() {
delay
} else {
sleep(time)
};
this.st = IoWriteState::Shutdown(Some(timeout), Shutdown::None);
self.poll(cx)
}
Poll::Ready(WriteStatus::Terminate) => {
log::trace!("write task is instructed to terminate");
// shutdown WRITE side
this.io
.local
.lock()
.unwrap()
.borrow_mut()
.flags
.insert(IoTestFlags::CLOSED);
this.state.close(None);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
IoWriteState::Shutdown(ref mut delay, ref mut st) => {
// close WRITE side and wait for disconnect on read side.
// use disconnect timeout, otherwise it could hang forever.
loop {
match st {
Shutdown::None => {
// flush write buffer
match flush_io(&this.io, &this.state, cx) {
Poll::Ready(Ok(())) => {
*st = Shutdown::Flushed;
continue;
}
Poll::Ready(Err(err)) => {
log::trace!(
"write task is closed with err during flush {:?}",
err
);
this.state.close(Some(err));
return Poll::Ready(());
}
Poll::Pending => (),
}
}
Shutdown::Flushed => {
// shutdown WRITE side
this.io
.local
.lock()
.unwrap()
.borrow_mut()
.flags
.insert(IoTestFlags::CLOSED);
*st = Shutdown::Stopping;
continue;
}
Shutdown::Stopping => {
// read until 0 or err
let io = &this.io;
loop {
let mut buf = BytesVec::new();
match io.poll_read_buf(cx, &mut buf) {
Poll::Ready(Err(e)) => {
this.state.close(Some(e));
log::trace!("write task is stopped");
return Poll::Ready(());
}
Poll::Ready(Ok(0)) => {
this.state.close(None);
log::trace!("write task is stopped");
return Poll::Ready(());
}
Poll::Pending => break,
_ => (),
}
}
}
}
// disconnect timeout
if let Some(ref delay) = delay {
if delay.poll_elapsed(cx).is_pending() {
return Poll::Pending;
}
}
log::trace!("write task is stopped after delay");
this.state.close(None);
return Poll::Ready(());
}
}
}
} }
} }
/// Flush write buffer to underlying I/O stream. /// Flush write buffer to underlying I/O stream.
pub(super) fn flush_io( pub(super) fn write_io(
io: &IoTest, io: &IoTest,
state: &WriteContext, buf: &mut BytesVec,
cx: &mut Context<'_>, cx: &mut Context<'_>,
) -> Poll<io::Result<()>> { ) -> Poll<io::Result<()>> {
state.with_buf(|buf| { let len = buf.len();
if let Some(buf) = buf {
let len = buf.len();
if len != 0 { if len != 0 {
log::trace!("flushing framed transport: {}", len); log::trace!("flushing framed transport: {}", len);
let mut written = 0; let mut written = 0;
let result = loop { let result = loop {
break match io.poll_write_buf(cx, &buf[written..]) { break match io.poll_write_buf(cx, &buf[written..]) {
Poll::Ready(Ok(n)) => { Poll::Ready(Ok(n)) => {
if n == 0 { if n == 0 {
log::trace!( log::trace!("disconnected during flush, written {}", written);
"disconnected during flush, written {}", Poll::Ready(Err(io::Error::new(
written io::ErrorKind::WriteZero,
); "failed to write frame to transport",
Poll::Ready(Err(io::Error::new( )))
io::ErrorKind::WriteZero, } else {
"failed to write frame to transport", written += n;
))) if written == len {
} else { buf.clear();
written += n; Poll::Ready(Ok(()))
if written == len { } else {
buf.clear(); continue;
Poll::Ready(Ok(()))
} else {
continue;
}
}
} }
Poll::Pending => { }
// remove written data }
buf.advance(written); Poll::Pending => {
Poll::Pending // remove written data
} buf.advance(written);
Poll::Ready(Err(e)) => { Poll::Pending
log::trace!("error during flush: {}", e); }
Poll::Ready(Err(e)) Poll::Ready(Err(e)) => {
} log::trace!("error during flush: {}", e);
}; Poll::Ready(Err(e))
}; }
log::trace!("flushed {} bytes", written); };
return result; };
} log::trace!("flushed {} bytes", written);
} result
} else {
Poll::Ready(Ok(())) Poll::Ready(Ok(()))
}) }
} }
#[cfg(test)] #[cfg(test)]

533
ntex-tokio/src/io.rs
View file

@ -1,12 +1,9 @@
use std::future::{poll_fn, Future};
use std::task::{Context, Poll}; use std::task::{Context, Poll};
use std::{any, cell::RefCell, cmp, io, mem, pin::Pin, rc::Rc, rc::Weak}; use std::{any, cell::RefCell, cmp, future::poll_fn, io, mem, pin::Pin, rc::Rc, rc::Weak};
use ntex_bytes::{Buf, BufMut, BytesVec}; use ntex_bytes::{Buf, BufMut, BytesVec};
use ntex_io::{ use ntex_io::{types, Filter, Handle, Io, IoBoxed, IoStream, ReadContext, WriteContext};
types, Filter, Handle, Io, IoBoxed, IoStream, ReadContext, WriteContext, WriteStatus, use ntex_util::{future::lazy, ready, time::Millis};
};
use ntex_util::{ready, time::sleep, time::Millis, time::Sleep};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::TcpStream; use tokio::net::TcpStream;
@ -14,11 +11,14 @@ impl IoStream for crate::TcpStream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let io = Rc::new(RefCell::new(self.0)); let io = Rc::new(RefCell::new(self.0));
let mut rio = ReadTask(io.clone()); let mut rio = Read(io.clone());
tokio::task::spawn_local(async move { tokio::task::spawn_local(async move {
read.handle(&mut rio).await; read.handle(&mut rio).await;
}); });
tokio::task::spawn_local(WriteTask::new(io.clone(), write)); let mut wio = Write(io.clone());
tokio::task::spawn_local(async move {
write.handle(&mut wio).await;
});
Some(Box::new(HandleWrapper(io))) Some(Box::new(HandleWrapper(io)))
} }
} }
@ -39,9 +39,10 @@ impl Handle for HandleWrapper {
} }
/// Read io task /// Read io task
struct ReadTask(Rc<RefCell<TcpStream>>); struct Read(Rc<RefCell<TcpStream>>);
impl ntex_io::AsyncRead for ReadTask { impl ntex_io::AsyncRead for Read {
#[inline]
async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) {
// read data from socket // read data from socket
let result = poll_fn(|cx| { let result = poll_fn(|cx| {
@ -53,281 +54,81 @@ impl ntex_io::AsyncRead for ReadTask {
} }
} }
#[derive(Debug)] struct Write(Rc<RefCell<TcpStream>>);
enum IoWriteState {
Processing(Option<Sleep>),
Shutdown(Sleep, Shutdown),
}
#[derive(Debug)] impl ntex_io::AsyncWrite for Write {
enum Shutdown { #[inline]
None, async fn write(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<()>) {
Flushed, match lazy(|cx| flush_io(&mut *self.0.borrow_mut(), &mut buf, cx)).await {
Stopping(u16), Poll::Ready(res) => (buf, res),
} Poll::Pending => (buf, Ok(())),
/// Write io task
struct WriteTask {
st: IoWriteState,
io: Rc<RefCell<TcpStream>>,
state: WriteContext,
}
impl WriteTask {
/// Create new write io task
fn new(io: Rc<RefCell<TcpStream>>, state: WriteContext) -> Self {
Self {
io,
state,
st: IoWriteState::Processing(None),
} }
} }
}
impl Future for WriteTask { #[inline]
type Output = (); async fn flush(&mut self) -> io::Result<()> {
Ok(())
}
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { #[inline]
let this = self.as_mut().get_mut(); async fn shutdown(&mut self) -> io::Result<()> {
poll_fn(|cx| Pin::new(&mut *self.0.borrow_mut()).poll_shutdown(cx)).await
if this.state.poll_close(cx).is_ready() {
return Poll::Ready(());
}
match this.st {
IoWriteState::Processing(ref mut delay) => {
match ready!(this.state.poll_ready(cx)) {
WriteStatus::Ready => {
if let Some(delay) = delay {
if delay.poll_elapsed(cx).is_ready() {
this.state.close(Some(io::Error::new(
io::ErrorKind::TimedOut,
"Operation timedout",
)));
return Poll::Ready(());
}
}
// flush io stream
match ready!(this.state.with_buf(|buf| flush_io(
&mut *this.io.borrow_mut(),
buf,
cx,
&this.state
))) {
Ok(()) => Poll::Pending,
Err(e) => {
this.state.close(Some(e));
Poll::Ready(())
}
}
}
WriteStatus::Timeout(time) => {
log::trace!(
"{}: Initiate timeout delay for {:?}",
this.state.tag(),
time
);
if delay.is_none() {
*delay = Some(sleep(time));
}
self.poll(cx)
}
WriteStatus::Shutdown(time) => {
log::trace!(
"{}: Write task is instructed to shutdown",
this.state.tag()
);
let timeout = if let Some(delay) = delay.take() {
delay
} else {
sleep(time)
};
this.st = IoWriteState::Shutdown(timeout, Shutdown::None);
self.poll(cx)
}
WriteStatus::Terminate => {
log::trace!(
"{}: Write task is instructed to terminate",
this.state.tag()
);
if !matches!(
this.io.borrow().linger(),
Ok(Some(std::time::Duration::ZERO))
) {
// call shutdown to prevent flushing data on terminated Io. when
// linger is set to zero, closing will reset the connection, so
// shutdown is not neccessary.
let _ = Pin::new(&mut *this.io.borrow_mut()).poll_shutdown(cx);
}
this.state.close(None);
Poll::Ready(())
}
}
}
IoWriteState::Shutdown(ref mut delay, ref mut st) => {
// close WRITE side and wait for disconnect on read side.
// use disconnect timeout, otherwise it could hang forever.
loop {
if this.state.poll_close(cx).is_ready() {
return Poll::Ready(());
}
match st {
Shutdown::None => {
// flush write buffer
let mut io = this.io.borrow_mut();
match this
.state
.with_buf(|buf| flush_io(&mut *io, buf, cx, &this.state))
{
Poll::Ready(Ok(())) => {
*st = Shutdown::Flushed;
continue;
}
Poll::Ready(Err(err)) => {
log::trace!(
"{}: Write task is closed with err during flush, {:?}", this.state.tag(),
err
);
this.state.close(Some(err));
return Poll::Ready(());
}
Poll::Pending => (),
}
}
Shutdown::Flushed => {
// shutdown WRITE side
match Pin::new(&mut *this.io.borrow_mut()).poll_shutdown(cx) {
Poll::Ready(Ok(_)) => {
*st = Shutdown::Stopping(0);
continue;
}
Poll::Ready(Err(e)) => {
log::trace!(
"{}: Write task is closed with err during shutdown",
this.state.tag()
);
this.state.close(Some(e));
return Poll::Ready(());
}
_ => (),
}
}
Shutdown::Stopping(ref mut count) => {
// read until 0 or err
let mut buf = [0u8; 512];
loop {
let mut read_buf = ReadBuf::new(&mut buf);
match Pin::new(&mut *this.io.borrow_mut())
.poll_read(cx, &mut read_buf)
{
Poll::Ready(Err(_)) | Poll::Ready(Ok(_))
if read_buf.filled().is_empty() =>
{
this.state.close(None);
log::trace!(
"{}: Tokio write task is stopped",
this.state.tag()
);
return Poll::Ready(());
}
Poll::Pending => {
*count += read_buf.filled().len() as u16;
if *count > 4096 {
log::trace!("{}: Tokio write task is stopped, too much input", this.state.tag());
this.state.close(None);
return Poll::Ready(());
}
break;
}
_ => (),
}
}
}
}
// disconnect timeout
if delay.poll_elapsed(cx).is_pending() {
return Poll::Pending;
}
log::trace!("{}: Write task is stopped after delay", this.state.tag());
this.state.close(None);
return Poll::Ready(());
}
}
}
} }
} }
/// Flush write buffer to underlying I/O stream. /// Flush write buffer to underlying I/O stream.
pub(super) fn flush_io<T: AsyncRead + AsyncWrite + Unpin>( pub(super) fn flush_io<T: AsyncRead + AsyncWrite + Unpin>(
io: &mut T, io: &mut T,
buf: &mut Option<BytesVec>, buf: &mut BytesVec,
cx: &mut Context<'_>, cx: &mut Context<'_>,
st: &WriteContext,
) -> Poll<io::Result<()>> { ) -> Poll<io::Result<()>> {
if let Some(buf) = buf { let len = buf.len();
let len = buf.len();
if len != 0 { if len != 0 {
// log::trace!("{}: Flushing framed transport: {:?}", st.tag(), buf.len()); // log::trace!("{}: Flushing framed transport: {:?}", st.tag(), buf.len());
let mut written = 0; let mut written = 0;
let result = loop { let result = loop {
break match Pin::new(&mut *io).poll_write(cx, &buf[written..]) { break match Pin::new(&mut *io).poll_write(cx, &buf[written..]) {
Poll::Ready(Ok(n)) => { Poll::Ready(Ok(n)) => {
if n == 0 { if n == 0 {
log::trace!( Poll::Ready(Err(io::Error::new(
"{}: Disconnected during flush, written {}", io::ErrorKind::WriteZero,
st.tag(), "failed to write frame to transport",
written )))
); } else {
Poll::Ready(Err(io::Error::new( written += n;
io::ErrorKind::WriteZero, if written == len {
"failed to write frame to transport", buf.clear();
))) Poll::Ready(Ok(()))
} else { } else {
written += n; continue;
if written == len {
buf.clear();
Poll::Ready(Ok(()))
} else {
continue;
}
} }
} }
Poll::Pending => {
// remove written data
buf.advance(written);
Poll::Pending
}
Poll::Ready(Err(e)) => {
log::trace!("{}: Error during flush: {}", st.tag(), e);
Poll::Ready(Err(e))
}
};
};
// log::trace!("{}: flushed {} bytes", st.tag(), written);
// flush
return if written > 0 {
match Pin::new(&mut *io).poll_flush(cx) {
Poll::Ready(Ok(_)) => result,
Poll::Pending => Poll::Pending,
Poll::Ready(Err(e)) => {
log::trace!("{}: Error during flush: {}", st.tag(), e);
Poll::Ready(Err(e))
}
} }
} else { Poll::Pending => {
result // remove written data
buf.advance(written);
Poll::Pending
}
Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
}; };
};
// log::trace!("{}: flushed {} bytes", st.tag(), written);
// flush
if written > 0 {
match Pin::new(&mut *io).poll_flush(cx) {
Poll::Ready(Ok(_)) => result,
Poll::Pending => Poll::Pending,
Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
}
} else {
result
} }
} else {
Poll::Ready(Ok(()))
} }
Poll::Ready(Ok(()))
} }
pub struct TokioIoBoxed(IoBoxed); pub struct TokioIoBoxed(IoBoxed);
@ -425,18 +226,22 @@ mod unixstream {
fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> { fn start(self, read: ReadContext, write: WriteContext) -> Option<Box<dyn Handle>> {
let io = Rc::new(RefCell::new(self.0)); let io = Rc::new(RefCell::new(self.0));
let mut rio = ReadTask(io.clone()); let mut rio = Read(io.clone());
tokio::task::spawn_local(async move { tokio::task::spawn_local(async move {
read.handle(&mut rio).await; read.handle(&mut rio).await;
}); });
tokio::task::spawn_local(WriteTask::new(io, write)); let mut wio = Write(io.clone());
tokio::task::spawn_local(async move {
write.handle(&mut wio).await;
});
None None
} }
} }
struct ReadTask(Rc<RefCell<UnixStream>>); struct Read(Rc<RefCell<UnixStream>>);
impl ntex_io::AsyncRead for ReadTask { impl ntex_io::AsyncRead for Read {
#[inline]
async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) { async fn read(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<usize>) {
// read data from socket // read data from socket
let result = poll_fn(|cx| { let result = poll_fn(|cx| {
@ -448,191 +253,25 @@ mod unixstream {
} }
} }
/// Write io task struct Write(Rc<RefCell<UnixStream>>);
struct WriteTask {
st: IoWriteState,
io: Rc<RefCell<UnixStream>>,
state: WriteContext,
}
impl WriteTask { impl ntex_io::AsyncWrite for Write {
/// Create new write io task #[inline]
fn new(io: Rc<RefCell<UnixStream>>, state: WriteContext) -> Self { async fn write(&mut self, mut buf: BytesVec) -> (BytesVec, io::Result<()>) {
Self { match lazy(|cx| flush_io(&mut *self.0.borrow_mut(), &mut buf, cx)).await {
io, Poll::Ready(res) => (buf, res),
state, Poll::Pending => (buf, Ok(())),
st: IoWriteState::Processing(None),
} }
} }
}
impl Future for WriteTask { #[inline]
type Output = (); async fn flush(&mut self) -> io::Result<()> {
Ok(())
}
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { #[inline]
let this = self.as_mut().get_mut(); async fn shutdown(&mut self) -> io::Result<()> {
poll_fn(|cx| Pin::new(&mut *self.0.borrow_mut()).poll_shutdown(cx)).await
if this.state.poll_close(cx).is_ready() {
return Poll::Ready(());
}
match this.st {
IoWriteState::Processing(ref mut delay) => {
match this.state.poll_ready(cx) {
Poll::Ready(WriteStatus::Ready) => {
if let Some(delay) = delay {
if delay.poll_elapsed(cx).is_ready() {
this.state.close(Some(io::Error::new(
io::ErrorKind::TimedOut,
"Operation timedout",
)));
return Poll::Ready(());
}
}
// flush io stream
match ready!(this.state.with_buf(|buf| flush_io(
&mut *this.io.borrow_mut(),
buf,
cx,
&this.state
))) {
Ok(()) => Poll::Pending,
Err(e) => {
this.state.close(Some(e));
Poll::Ready(())
}
}
}
Poll::Ready(WriteStatus::Timeout(time)) => {
if delay.is_none() {
*delay = Some(sleep(time));
}
self.poll(cx)
}
Poll::Ready(WriteStatus::Shutdown(time)) => {
log::trace!(
"{}: Write task is instructed to shutdown",
this.state.tag()
);
let timeout = if let Some(delay) = delay.take() {
delay
} else {
sleep(time)
};
this.st = IoWriteState::Shutdown(timeout, Shutdown::None);
self.poll(cx)
}
Poll::Ready(WriteStatus::Terminate) => {
log::trace!(
"{}: Write task is instructed to terminate",
this.state.tag()
);
let _ = Pin::new(&mut *this.io.borrow_mut()).poll_shutdown(cx);
this.state.close(None);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
IoWriteState::Shutdown(ref mut delay, ref mut st) => {
// close WRITE side and wait for disconnect on read side.
// use disconnect timeout, otherwise it could hang forever.
loop {
if this.state.poll_close(cx).is_ready() {
return Poll::Ready(());
}
match st {
Shutdown::None => {
// flush write buffer
let mut io = this.io.borrow_mut();
match this.state.with_buf(|buf| {
flush_io(&mut *io, buf, cx, &this.state)
}) {
Poll::Ready(Ok(())) => {
*st = Shutdown::Flushed;
continue;
}
Poll::Ready(Err(err)) => {
log::trace!(
"{}: Write task is closed with err during flush, {:?}", this.state.tag(),
err
);
this.state.close(Some(err));
return Poll::Ready(());
}
Poll::Pending => (),
}
}
Shutdown::Flushed => {
// shutdown WRITE side
match Pin::new(&mut *this.io.borrow_mut()).poll_shutdown(cx)
{
Poll::Ready(Ok(_)) => {
*st = Shutdown::Stopping(0);
continue;
}
Poll::Ready(Err(e)) => {
log::trace!(
"{}: Write task is closed with err during shutdown", this.state.tag()
);
this.state.close(Some(e));
return Poll::Ready(());
}
_ => (),
}
}
Shutdown::Stopping(ref mut count) => {
// read until 0 or err
let mut buf = [0u8; 512];
loop {
let mut read_buf = ReadBuf::new(&mut buf);
match Pin::new(&mut *this.io.borrow_mut())
.poll_read(cx, &mut read_buf)
{
Poll::Ready(Err(_)) | Poll::Ready(Ok(_))
if read_buf.filled().is_empty() =>
{
this.state.close(None);
log::trace!(
"{}: Write task is stopped",
this.state.tag()
);
return Poll::Ready(());
}
Poll::Pending => {
*count += read_buf.filled().len() as u16;
if *count > 4096 {
log::trace!(
"{}: Write task is stopped, too much input", this.state.tag()
);
this.state.close(None);
return Poll::Ready(());
}
break;
}
_ => (),
}
}
}
}
// disconnect timeout
if delay.poll_elapsed(cx).is_pending() {
return Poll::Pending;
}
log::trace!(
"{}: Write task is stopped after delay",
this.state.tag()
);
this.state.close(None);
return Poll::Ready(());
}
}
}
} }
} }
} }