Something went wrong. Try again.
[READ-ONLY] Mirror of https://github.com/chrisvander/crate. The future of cloud storage.
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344use std::{ pin::Pin, task::{Context, Poll},};use tokio::{ io::{AsyncRead, ReadBuf}, sync::OwnedSemaphorePermit,};
// Keep the request slot until its response has drained, including slow downloads.pub struct LimitedBody { inner: Pin<Box<dyn AsyncRead + Send>>, permit: Option<OwnedSemaphorePermit>, deadline: Pin<Box<tokio::time::Sleep>>,}
impl LimitedBody { pub fn wrap(body: poem::Body, permit: OwnedSemaphorePermit) -> poem::Body { poem::Body::from_async_read(Self { inner: Box::pin(body.into_async_read()), permit: Some(permit), deadline: Box::pin(tokio::time::sleep(std::time::Duration::from_secs(300))), }) }}
impl AsyncRead for LimitedBody { fn poll_read( mut self: Pin<&mut Self>, context: &mut Context<'_>, buffer: &mut ReadBuf<'_>, ) -> Poll<std::io::Result<()>> { use std::future::Future; if self.deadline.as_mut().poll(context).is_ready() { self.permit.take(); return Poll::Ready(Err(std::io::Error::new( std::io::ErrorKind::TimedOut, "The response transfer timed out.", ))); } self.inner.as_mut().poll_read(context, buffer) }}