2023-09-30 21:26:43 +00:00
|
|
|
use std::{
|
|
|
|
future::{ready, Ready},
|
2023-09-30 22:24:48 +00:00
|
|
|
rc::Rc,
|
2023-12-11 01:26:16 +00:00
|
|
|
time::Duration,
|
2023-09-30 21:26:43 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
use actix_web::{
|
|
|
|
dev::{Service, ServiceRequest, Transform},
|
|
|
|
HttpMessage,
|
|
|
|
};
|
|
|
|
use streem::IntoStreamer;
|
|
|
|
use tokio::task::JoinSet;
|
|
|
|
|
|
|
|
use crate::{future::NowOrNever, stream::LocalBoxStream};
|
|
|
|
|
|
|
|
const LIMIT: usize = 256;
|
|
|
|
|
|
|
|
async fn drain(rx: flume::Receiver<actix_web::dev::Payload>) {
|
|
|
|
let mut set = JoinSet::new();
|
|
|
|
|
|
|
|
while let Ok(payload) = rx.recv_async().await {
|
2023-12-28 17:58:38 +00:00
|
|
|
tracing::trace!("drain: looping");
|
|
|
|
|
2023-12-11 01:26:16 +00:00
|
|
|
// draining a payload is a best-effort task - if we can't collect in 2 minutes we bail
|
|
|
|
set.spawn_local(tokio::time::timeout(Duration::from_secs(120), async move {
|
2023-09-30 21:26:43 +00:00
|
|
|
let mut streamer = payload.into_streamer();
|
2023-12-28 17:58:38 +00:00
|
|
|
while streamer.next().await.is_some() {
|
|
|
|
tracing::trace!("drain drop bytes: looping");
|
|
|
|
}
|
2023-12-11 01:26:16 +00:00
|
|
|
}));
|
2023-09-30 21:26:43 +00:00
|
|
|
|
|
|
|
let mut count = 0;
|
|
|
|
|
|
|
|
// drain completed tasks
|
|
|
|
while set.join_next().now_or_never().is_some() {
|
2023-12-28 17:58:38 +00:00
|
|
|
tracing::trace!("drain join now: looping");
|
|
|
|
|
2023-09-30 21:26:43 +00:00
|
|
|
count += 1;
|
|
|
|
}
|
|
|
|
|
|
|
|
// if we're past the limit, wait for completions
|
|
|
|
while set.len() > LIMIT {
|
2023-12-28 17:58:38 +00:00
|
|
|
tracing::trace!("drain join await: looping");
|
|
|
|
|
2023-09-30 21:26:43 +00:00
|
|
|
if set.join_next().await.is_some() {
|
|
|
|
count += 1;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
if count > 0 {
|
|
|
|
tracing::info!("Drained {count} dropped payloads");
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// drain set
|
2023-12-28 17:58:38 +00:00
|
|
|
while set.join_next().await.is_some() {
|
|
|
|
tracing::trace!("drain join await cleanup: looping");
|
|
|
|
}
|
2023-09-30 21:26:43 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Clone)]
|
2023-10-21 00:08:11 +00:00
|
|
|
struct DrainHandle(Option<Rc<tokio::task::JoinHandle<()>>>);
|
2023-09-30 21:26:43 +00:00
|
|
|
|
|
|
|
pub(crate) struct Payload {
|
|
|
|
sender: flume::Sender<actix_web::dev::Payload>,
|
|
|
|
handle: DrainHandle,
|
|
|
|
}
|
|
|
|
pub(crate) struct PayloadMiddleware<S> {
|
|
|
|
inner: S,
|
|
|
|
sender: flume::Sender<actix_web::dev::Payload>,
|
|
|
|
_handle: DrainHandle,
|
|
|
|
}
|
|
|
|
|
|
|
|
pub(crate) struct PayloadStream {
|
|
|
|
inner: Option<actix_web::dev::Payload>,
|
|
|
|
sender: flume::Sender<actix_web::dev::Payload>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl DrainHandle {
|
2023-10-21 00:08:11 +00:00
|
|
|
fn new(handle: tokio::task::JoinHandle<()>) -> Self {
|
2023-09-30 22:24:48 +00:00
|
|
|
Self(Some(Rc::new(handle)))
|
2023-09-30 21:26:43 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl Payload {
|
|
|
|
pub(crate) fn new() -> Self {
|
|
|
|
let (tx, rx) = crate::sync::channel(LIMIT);
|
|
|
|
|
2023-10-21 00:08:11 +00:00
|
|
|
let handle = DrainHandle::new(crate::sync::spawn("drain-payloads", drain(rx)));
|
2023-09-30 21:26:43 +00:00
|
|
|
|
|
|
|
Payload { sender: tx, handle }
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl Drop for DrainHandle {
|
|
|
|
fn drop(&mut self) {
|
2023-09-30 22:24:48 +00:00
|
|
|
if let Some(handle) = self.0.take().and_then(Rc::into_inner) {
|
2023-09-30 21:26:43 +00:00
|
|
|
handle.abort();
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl Drop for PayloadStream {
|
|
|
|
fn drop(&mut self) {
|
|
|
|
if let Some(payload) = self.inner.take() {
|
|
|
|
tracing::warn!("Dropped unclosed payload, draining");
|
|
|
|
if self.sender.try_send(payload).is_err() {
|
|
|
|
tracing::error!("Failed to send unclosed payload for draining");
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl futures_core::Stream for PayloadStream {
|
|
|
|
type Item = Result<actix_web::web::Bytes, actix_web::error::PayloadError>;
|
|
|
|
|
|
|
|
fn poll_next(
|
|
|
|
mut self: std::pin::Pin<&mut Self>,
|
|
|
|
cx: &mut std::task::Context<'_>,
|
|
|
|
) -> std::task::Poll<Option<Self::Item>> {
|
|
|
|
if let Some(inner) = self.inner.as_mut() {
|
|
|
|
let opt = std::task::ready!(std::pin::Pin::new(inner).poll_next(cx));
|
|
|
|
|
|
|
|
if opt.is_none() {
|
|
|
|
self.inner.take();
|
|
|
|
}
|
|
|
|
|
|
|
|
std::task::Poll::Ready(opt)
|
|
|
|
} else {
|
|
|
|
std::task::Poll::Ready(None)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn size_hint(&self) -> (usize, Option<usize>) {
|
|
|
|
if let Some(inner) = self.inner.as_ref() {
|
|
|
|
inner.size_hint()
|
|
|
|
} else {
|
|
|
|
(0, Some(0))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<S> Transform<S, ServiceRequest> for Payload
|
|
|
|
where
|
|
|
|
S: Service<ServiceRequest>,
|
|
|
|
S::Future: 'static,
|
|
|
|
{
|
|
|
|
type Response = S::Response;
|
|
|
|
type Error = S::Error;
|
|
|
|
type InitError = ();
|
|
|
|
type Transform = PayloadMiddleware<S>;
|
|
|
|
type Future = Ready<Result<Self::Transform, Self::InitError>>;
|
|
|
|
|
|
|
|
fn new_transform(&self, service: S) -> Self::Future {
|
|
|
|
ready(Ok(PayloadMiddleware {
|
|
|
|
inner: service,
|
|
|
|
sender: self.sender.clone(),
|
|
|
|
_handle: self.handle.clone(),
|
|
|
|
}))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<S> Service<ServiceRequest> for PayloadMiddleware<S>
|
|
|
|
where
|
|
|
|
S: Service<ServiceRequest>,
|
|
|
|
S::Future: 'static,
|
|
|
|
{
|
|
|
|
type Response = S::Response;
|
|
|
|
type Error = S::Error;
|
|
|
|
type Future = S::Future;
|
|
|
|
|
|
|
|
fn poll_ready(
|
|
|
|
&self,
|
|
|
|
ctx: &mut core::task::Context<'_>,
|
|
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
|
|
self.inner.poll_ready(ctx)
|
|
|
|
}
|
|
|
|
|
|
|
|
fn call(&self, mut req: ServiceRequest) -> Self::Future {
|
2023-09-30 22:24:48 +00:00
|
|
|
let payload = req.take_payload();
|
|
|
|
|
|
|
|
if !matches!(payload, actix_web::dev::Payload::None) {
|
2023-09-30 21:26:43 +00:00
|
|
|
let payload: LocalBoxStream<'static, _> = Box::pin(PayloadStream {
|
|
|
|
inner: Some(payload),
|
|
|
|
sender: self.sender.clone(),
|
|
|
|
});
|
|
|
|
req.set_payload(payload.into());
|
|
|
|
}
|
|
|
|
|
|
|
|
self.inner.call(req)
|
|
|
|
}
|
|
|
|
}
|