GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added read deadline to AvailableReader

Browse Source
This commit is contained in:
Jędrzej Stuczyński
2023-03-24 18:10:46 +00:00
parent 3f0d4846df
commit 790220039b
1 changed files with 76 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/socks5/proxy-helpers/src/available_reader.rs
+76 -30
View File
@@ -1,8 +1,9 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// Copyright 2021-2023 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use bytes::{BufMut, Bytes, BytesMut};
use futures::Stream;
use log::error;
use std::cell::RefCell;
use std::future::Future;
use std::io;
@@ -15,6 +16,7 @@ use tokio_util::io::poll_read_buf;
const MAX_READ_AMOUNT: usize = 500 * 1000; // 0.5MB
const GRACE_DURATION: Duration = Duration::from_millis(1);
const READ_TIMEOUT: Duration = Duration::from_millis(10);
pub struct AvailableReader<'a, R: AsyncRead + Unpin> {
// TODO: come up with a way to avoid using RefCell (not sure if possible though due to having to
@@ -22,6 +24,7 @@ pub struct AvailableReader<'a, R: AsyncRead + Unpin> {
buf: RefCell<BytesMut>,
inner: RefCell<&'a mut R>,
grace_period: Option<Pin<Box<Sleep>>>,
read_deadline: Option<Pin<Box<Sleep>>>,
}
impl<'a, R> AvailableReader<'a, R>
@@ -31,10 +34,58 @@ where
const BUF_INCREMENT: usize = 4096;
pub fn new(reader: &'a mut R) -> Self {
// pub fn new(reader: &'a mut R, buffer_size: usize) -> Self {
AvailableReader {
buf: RefCell::new(BytesMut::with_capacity(Self::BUF_INCREMENT)),
inner: RefCell::new(reader),
grace_period: Some(Box::pin(sleep(GRACE_DURATION))),
grace_period: None,
read_deadline: None,
}
}
fn return_buf(mut self: Pin<&mut Self>) -> Poll<Option<<Self as Stream>::Item>> {
// reset timeouts so the poll wouldn't be accidentally called
self.grace_period = None;
self.read_deadline = None;
let buf = self.buf.borrow_mut().split();
// there's no data in the buffer, it means the underlying source is done
if buf.is_empty() {
Poll::Ready(None)
} else {
Poll::Ready(Some(Ok(buf.freeze())))
}
}
fn poll_read_deadline(&mut self, cx: &mut Context<'_>) -> Poll<()> {
match self.read_deadline.as_mut() {
Some(read_deadline) => Pin::new(read_deadline).poll(cx),
None => {
// this is the first time we're polling this reader - set the deadline
self.read_deadline = Some(Box::pin(sleep(READ_TIMEOUT)));
Poll::Pending
}
}
}
fn poll_grace_period(&mut self, cx: &mut Context<'_>) -> Poll<()> {
match self.grace_period.as_mut() {
Some(grace_period) => Pin::new(grace_period).poll(cx),
None => {
// this is the first time we're polling this reader - set the grace period
self.grace_period = Some(Box::pin(sleep(GRACE_DURATION)));
Poll::Pending
}
}
}
fn reset_grace_period(&mut self) {
if let Some(grace_period) = self.grace_period.as_mut() {
let now = Instant::now();
grace_period.as_mut().reset(now + GRACE_DURATION);
} else {
// this branch should be impossible!
error!("reached an impossible branch - attempted to reset a non-existent grace period timeout")
}
}
}
@@ -55,50 +106,46 @@ impl<'a, R: AsyncRead + Unpin> Stream for AvailableReader<'a, R> {
self.buf.borrow_mut().deref_mut(),
);
let deadline_poll_res = self.poll_read_deadline(cx);
let grace_period_poll_res = self.poll_grace_period(cx);
match poll_res {
Poll::Pending => {
// there's nothing for us here, just return whatever we have (assuming we read anything!)
// there's nothing for us here, just return whatever we have
// (assuming we read anything and we SHOULD HAVE read something)
if self.buf.borrow().is_empty() {
// this case shouldn't be possible - if something woke the future we MUST HAVE HAD some data
// (it's because we never set the sleeps until the first call, meaning some data was already there)
error!("AvailableReader got woken with not data to return!");
Poll::Pending
} else if grace_period_poll_res.is_pending() && deadline_poll_res.is_pending() {
// if neither timeout was reached yet, wait a bit more
Poll::Pending
} else {
// if exists - check grace period
if let Some(grace_period) = self.grace_period.as_mut() {
if Pin::new(grace_period).poll(cx).is_pending() {
return Poll::Pending;
}
}
let buf = self.buf.replace(BytesMut::new());
Poll::Ready(Some(Ok(buf.freeze())))
// otherwise return what we managed to read (and reset timeouts)
self.return_buf()
}
}
Poll::Ready(Err(err)) => Poll::Ready(Some(Err(err))),
Poll::Ready(Ok(n)) => {
// if exists - reset grace period
if let Some(grace_period) = self.grace_period.as_mut() {
let now = Instant::now();
grace_period.as_mut().reset(now + GRACE_DURATION);
}
self.reset_grace_period();
// if we read a non-0 amount, we're not done yet!
if n == 0 {
let buf = self.buf.replace(BytesMut::new());
if !buf.is_empty() {
Poll::Ready(Some(Ok(buf.freeze())))
} else {
Poll::Ready(None)
}
// if we read 0 bytes, it means the underlying source is done, so return what we have
// (note, this might get called twice if we already had something in the buffer)
self.return_buf()
} else {
// tell the waker we should be polled again!
cx.waker().wake_by_ref();
// if we reached our maximum amount - return it
// if we reached our maximum amount or we've been trying to read the data for too long
// return what we have
let read_bytes_len = self.buf.borrow().len();
if read_bytes_len >= MAX_READ_AMOUNT {
let buf = self.buf.replace(BytesMut::new());
return Poll::Ready(Some(Ok(buf.freeze())));
if read_bytes_len >= MAX_READ_AMOUNT || deadline_poll_res.is_ready() {
self.return_buf()
} else {
Poll::Pending
}
Poll::Pending
}
}
}
@@ -154,7 +201,6 @@ mod tests {
assert_eq!(read_data, first_data_chunk);
assert_pending!(poll!(available_reader.next()));
// before dropping the mock, we need to empty it
let mut buf = vec![0u8; second_data_chunk.len()];
assert_eq!(reader_mock.read(&mut buf).await.unwrap(), 100);