2018-12-09 10:15:49 -08:00
|
|
|
use actix_service::{IntoNewService, IntoService, NewService, Service};
|
2018-12-06 14:04:42 -08:00
|
|
|
use futures::{try_ready, Async, Future, Poll};
|
2018-09-14 13:30:29 -07:00
|
|
|
|
|
|
|
use super::counter::{Counter, CounterGuard};
|
|
|
|
|
|
|
|
/// InFlight - new service for service that can limit number of in-flight
|
|
|
|
/// async requests.
|
|
|
|
///
|
|
|
|
/// Default number of in-flight requests is 15
|
|
|
|
pub struct InFlight<T> {
|
|
|
|
factory: T,
|
|
|
|
max_inflight: usize,
|
|
|
|
}
|
|
|
|
|
2018-11-29 16:56:15 -10:00
|
|
|
impl<T> InFlight<T> {
|
2019-02-01 19:53:13 -08:00
|
|
|
pub fn new<F>(factory: F) -> Self
|
2018-11-29 16:56:15 -10:00
|
|
|
where
|
2019-02-01 19:53:13 -08:00
|
|
|
T: NewService,
|
|
|
|
F: IntoNewService<T>,
|
2018-11-29 16:56:15 -10:00
|
|
|
{
|
2018-09-14 13:30:29 -07:00
|
|
|
Self {
|
|
|
|
factory: factory.into_new_service(),
|
|
|
|
max_inflight: 15,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Set max number of in-flight requests.
|
|
|
|
///
|
|
|
|
/// By default max in-flight requests is 15.
|
|
|
|
pub fn max_inflight(mut self, max: usize) -> Self {
|
|
|
|
self.max_inflight = max;
|
|
|
|
self
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-02-01 19:53:13 -08:00
|
|
|
impl<T> NewService for InFlight<T>
|
2018-09-14 13:30:29 -07:00
|
|
|
where
|
2019-02-01 19:53:13 -08:00
|
|
|
T: NewService,
|
2018-09-14 13:30:29 -07:00
|
|
|
{
|
2019-02-01 19:53:13 -08:00
|
|
|
type Request = T::Request;
|
2018-09-14 13:30:29 -07:00
|
|
|
type Response = T::Response;
|
|
|
|
type Error = T::Error;
|
|
|
|
type InitError = T::InitError;
|
|
|
|
type Service = InFlightService<T::Service>;
|
2019-02-01 19:53:13 -08:00
|
|
|
type Future = InFlightResponseFuture<T>;
|
2018-09-14 13:30:29 -07:00
|
|
|
|
|
|
|
fn new_service(&self) -> Self::Future {
|
|
|
|
InFlightResponseFuture {
|
|
|
|
fut: self.factory.new_service(),
|
|
|
|
max_inflight: self.max_inflight,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-02-01 19:53:13 -08:00
|
|
|
pub struct InFlightResponseFuture<T: NewService> {
|
2018-09-14 13:30:29 -07:00
|
|
|
fut: T::Future,
|
|
|
|
max_inflight: usize,
|
|
|
|
}
|
|
|
|
|
2019-02-01 19:53:13 -08:00
|
|
|
impl<T: NewService> Future for InFlightResponseFuture<T> {
|
2018-09-14 13:30:29 -07:00
|
|
|
type Item = InFlightService<T::Service>;
|
|
|
|
type Error = T::InitError;
|
|
|
|
|
|
|
|
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
|
|
|
|
Ok(Async::Ready(InFlightService::with_max_inflight(
|
|
|
|
self.max_inflight,
|
|
|
|
try_ready!(self.fut.poll()),
|
|
|
|
)))
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
pub struct InFlightService<T> {
|
|
|
|
service: T,
|
|
|
|
count: Counter,
|
|
|
|
}
|
|
|
|
|
2018-11-29 16:56:15 -10:00
|
|
|
impl<T> InFlightService<T> {
|
2019-02-01 19:53:13 -08:00
|
|
|
pub fn new<F>(service: F) -> Self
|
2018-11-29 16:56:15 -10:00
|
|
|
where
|
2019-02-01 19:53:13 -08:00
|
|
|
T: Service,
|
|
|
|
F: IntoService<T>,
|
2018-11-29 16:56:15 -10:00
|
|
|
{
|
2018-09-14 13:30:29 -07:00
|
|
|
Self {
|
|
|
|
service: service.into_service(),
|
|
|
|
count: Counter::new(15),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-02-01 19:53:13 -08:00
|
|
|
pub fn with_max_inflight<F>(max: usize, service: F) -> Self
|
2018-11-29 16:56:15 -10:00
|
|
|
where
|
2019-02-01 19:53:13 -08:00
|
|
|
T: Service,
|
|
|
|
F: IntoService<T>,
|
2018-11-29 16:56:15 -10:00
|
|
|
{
|
2018-09-14 13:30:29 -07:00
|
|
|
Self {
|
|
|
|
service: service.into_service(),
|
|
|
|
count: Counter::new(max),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2019-02-01 19:53:13 -08:00
|
|
|
impl<T> Service for InFlightService<T>
|
2018-11-29 16:56:15 -10:00
|
|
|
where
|
2019-02-01 19:53:13 -08:00
|
|
|
T: Service,
|
2018-11-29 16:56:15 -10:00
|
|
|
{
|
2019-02-01 19:53:13 -08:00
|
|
|
type Request = T::Request;
|
2018-09-14 13:30:29 -07:00
|
|
|
type Response = T::Response;
|
|
|
|
type Error = T::Error;
|
2019-02-01 19:53:13 -08:00
|
|
|
type Future = InFlightServiceResponse<T>;
|
2018-09-14 13:30:29 -07:00
|
|
|
|
|
|
|
fn poll_ready(&mut self) -> Poll<(), Self::Error> {
|
2019-02-01 15:15:53 -08:00
|
|
|
let res = self.service.poll_ready()?;
|
|
|
|
if res.is_ready() && !self.count.available() {
|
|
|
|
log::trace!("InFlight limit exceeded");
|
2018-09-14 13:30:29 -07:00
|
|
|
return Ok(Async::NotReady);
|
|
|
|
}
|
2019-02-01 15:15:53 -08:00
|
|
|
Ok(res)
|
2018-09-14 13:30:29 -07:00
|
|
|
}
|
|
|
|
|
2019-02-01 19:53:13 -08:00
|
|
|
fn call(&mut self, req: T::Request) -> Self::Future {
|
2018-09-14 13:30:29 -07:00
|
|
|
InFlightServiceResponse {
|
|
|
|
fut: self.service.call(req),
|
2018-11-29 16:56:15 -10:00
|
|
|
_guard: self.count.get(),
|
2018-09-14 13:30:29 -07:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[doc(hidden)]
|
2019-02-01 19:53:13 -08:00
|
|
|
pub struct InFlightServiceResponse<T: Service> {
|
2018-09-14 13:30:29 -07:00
|
|
|
fut: T::Future,
|
2018-11-29 16:56:15 -10:00
|
|
|
_guard: CounterGuard,
|
2018-09-14 13:30:29 -07:00
|
|
|
}
|
|
|
|
|
2019-02-01 19:53:13 -08:00
|
|
|
impl<T: Service> Future for InFlightServiceResponse<T> {
|
2018-09-14 13:30:29 -07:00
|
|
|
type Item = T::Response;
|
|
|
|
type Error = T::Error;
|
|
|
|
|
|
|
|
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
|
|
|
|
self.fut.poll()
|
|
|
|
}
|
|
|
|
}
|