mirror of
https://github.com/OMGeeky/tarpc.git
synced 2025-12-30 08:08:58 +01:00
Remove deprecated tokio-proto and replace with homegrown rpc framework (#199)
# New Crates - crate rpc contains the core client/server request-response framework, as well as a transport trait. - crate bincode-transport implements a transport that works almost exactly as tarpc works today (not to say it's wire-compatible). - crate trace has some foundational types for tracing. This isn't really fleshed out yet, but it's useful for in-process log tracing, at least. All crates are now at the top level. e.g. tarpc-plugins is now tarpc/plugins rather than tarpc/src/plugins. tarpc itself is now a *very* small code surface, as most functionality has been moved into the other more granular crates. # New Features - deadlines: all requests specify a deadline, and a server will stop processing a response when past its deadline. - client cancellation propagation: when a client drops a request, the client sends a message to the server informing it to cancel its response. This means cancellations can propagate across multiple server hops. - trace context stuff as mentioned above - more server configuration for total connection limits, per-connection request limits, etc. # Removals - no more shutdown handle. I left it out for now because of time and not being sure what the right solution is. - all async now, no blocking stub or server interface. This helps with maintainability, and async/await makes async code much more usable. The service trait is thusly renamed Service, and the client is renamed Client. - no built-in transport. Tarpc is now transport agnostic (see bincode-transport for transitioning existing uses). - going along with the previous bullet, no preferred transport means no TLS support at this time. We could make a tls transport or make bincode-transport compatible with TLS. - a lot of examples were removed because I couldn't keep up with maintaining all of them. Hopefully the ones I kept are still illustrative. - no more plugins! # Open Questions 1. Should client.send() return `Future<Response>` or `Future<Future<Response>>`? The former appears more ergonomic but it doesn’t allow concurrent requests with a single client handle. The latter is less ergonomic but yields back control of the client once it’s successfully sent out the request. Should we offer fns for both? 2. Should rpc service! Fns take &mut self or &self or self? The service needs to impl Clone anyway, technically we only need to clone it once per connection, and then leave it up to the user to decide if they want to clone it per RPC. In practice, everyone doing nontrivial stuff will need to clone it per RPC, I think. 3. Do the request/response structs look ok? 4. Is supporting server shutdown/lameduck important? Fixes #178 #155 #124 #104 #83 #38
This commit is contained in:
Tim
GitHub
251
rpc/src/server/filter.rs
Normal file
251
rpc/src/server/filter.rs
Normal file
@@ -0,0 +1,251 @@
|
||||
use crate::{
|
||||
server::{Channel, Config},
|
||||
util::Compact,
|
||||
ClientMessage, Response, Transport,
|
||||
};
|
||||
use fnv::FnvHashMap;
|
||||
use futures::{channel::mpsc, prelude::*, ready, stream::Fuse, task::{LocalWaker, Poll}};
|
||||
use log::{debug, error, info, trace, warn};
|
||||
use pin_utils::unsafe_pinned;
|
||||
use std::{
|
||||
collections::hash_map::Entry,
|
||||
io,
|
||||
marker::PhantomData,
|
||||
net::{IpAddr, SocketAddr},
|
||||
ops::Try,
|
||||
option::NoneError,
|
||||
pin::Pin,
|
||||
};
|
||||
|
||||
/// Drops connections under configurable conditions:
|
||||
///
|
||||
/// 1. If the max number of connections is reached.
|
||||
/// 2. If the max number of connections for a single IP is reached.
|
||||
#[derive(Debug)]
|
||||
pub struct ConnectionFilter<S, Req, Resp> {
|
||||
listener: Fuse<S>,
|
||||
closed_connections: mpsc::UnboundedSender<SocketAddr>,
|
||||
closed_connections_rx: mpsc::UnboundedReceiver<SocketAddr>,
|
||||
config: Config,
|
||||
connections_per_ip: FnvHashMap<IpAddr, usize>,
|
||||
open_connections: usize,
|
||||
ghost: PhantomData<(Req, Resp)>,
|
||||
}
|
||||
|
||||
enum NewConnection<Req, Resp, C> {
|
||||
Filtered,
|
||||
Accepted(Channel<Req, Resp, C>),
|
||||
}
|
||||
|
||||
impl<Req, Resp, C> Try for NewConnection<Req, Resp, C> {
|
||||
type Ok = Channel<Req, Resp, C>;
|
||||
type Error = NoneError;
|
||||
|
||||
fn into_result(self) -> Result<Channel<Req, Resp, C>, NoneError> {
|
||||
match self {
|
||||
NewConnection::Filtered => Err(NoneError),
|
||||
NewConnection::Accepted(channel) => Ok(channel),
|
||||
}
|
||||
}
|
||||
|
||||
fn from_error(_: NoneError) -> Self {
|
||||
NewConnection::Filtered
|
||||
}
|
||||
|
||||
fn from_ok(channel: Channel<Req, Resp, C>) -> Self {
|
||||
NewConnection::Accepted(channel)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, Req, Resp> ConnectionFilter<S, Req, Resp> {
|
||||
unsafe_pinned!(open_connections: usize);
|
||||
unsafe_pinned!(config: Config);
|
||||
unsafe_pinned!(connections_per_ip: FnvHashMap<IpAddr, usize>);
|
||||
unsafe_pinned!(closed_connections_rx: mpsc::UnboundedReceiver<SocketAddr>);
|
||||
unsafe_pinned!(listener: Fuse<S>);
|
||||
|
||||
/// Sheds new connections to stay under configured limits.
|
||||
pub fn filter<C>(listener: S, config: Config) -> Self
|
||||
where
|
||||
S: Stream<Item = Result<C, io::Error>>,
|
||||
C: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
{
|
||||
let (closed_connections, closed_connections_rx) = mpsc::unbounded();
|
||||
|
||||
ConnectionFilter {
|
||||
listener: listener.fuse(),
|
||||
closed_connections,
|
||||
closed_connections_rx,
|
||||
config,
|
||||
connections_per_ip: FnvHashMap::default(),
|
||||
open_connections: 0,
|
||||
ghost: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_new_connection<C>(self: &mut Pin<&mut Self>, stream: C) -> NewConnection<Req, Resp, C>
|
||||
where
|
||||
C: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
{
|
||||
let peer = match stream.peer_addr() {
|
||||
Ok(peer) => peer,
|
||||
Err(e) => {
|
||||
warn!("Could not get peer_addr of new connection: {}", e);
|
||||
return NewConnection::Filtered;
|
||||
}
|
||||
};
|
||||
|
||||
let open_connections = *self.open_connections();
|
||||
if open_connections >= self.config().max_connections {
|
||||
warn!(
|
||||
"[{}] Shedding connection because the maximum open connections \
|
||||
limit is reached ({}/{}).",
|
||||
peer,
|
||||
open_connections,
|
||||
self.config().max_connections
|
||||
);
|
||||
return NewConnection::Filtered;
|
||||
}
|
||||
|
||||
let config = self.config.clone();
|
||||
let open_connections_for_ip = self.increment_connections_for_ip(&peer)?;
|
||||
*self.open_connections() += 1;
|
||||
|
||||
debug!(
|
||||
"[{}] Opening channel ({}/{} connections for IP, {} total).",
|
||||
peer,
|
||||
open_connections_for_ip,
|
||||
config.max_connections_per_ip,
|
||||
self.open_connections(),
|
||||
);
|
||||
|
||||
NewConnection::Accepted(Channel {
|
||||
client_addr: peer,
|
||||
closed_connections: self.closed_connections.clone(),
|
||||
transport: stream.fuse(),
|
||||
config,
|
||||
ghost: PhantomData,
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_closed_connection(self: &mut Pin<&mut Self>, addr: &SocketAddr) {
|
||||
*self.open_connections() -= 1;
|
||||
debug!(
|
||||
"[{}] Closing channel. {} open connections remaining.",
|
||||
addr, self.open_connections
|
||||
);
|
||||
self.decrement_connections_for_ip(&addr);
|
||||
self.connections_per_ip().compact(0.1);
|
||||
}
|
||||
|
||||
fn increment_connections_for_ip(self: &mut Pin<&mut Self>, peer: &SocketAddr) -> Option<usize> {
|
||||
let max_connections_per_ip = self.config().max_connections_per_ip;
|
||||
let mut occupied;
|
||||
let mut connections_per_ip = self.connections_per_ip();
|
||||
let occupied = match connections_per_ip.entry(peer.ip()) {
|
||||
Entry::Vacant(vacant) => vacant.insert(0),
|
||||
Entry::Occupied(o) => {
|
||||
if *o.get() < max_connections_per_ip {
|
||||
// Store the reference outside the block to extend the lifetime.
|
||||
occupied = o;
|
||||
occupied.get_mut()
|
||||
} else {
|
||||
info!(
|
||||
"[{}] Opened max connections from IP ({}/{}).",
|
||||
peer,
|
||||
o.get(),
|
||||
max_connections_per_ip
|
||||
);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
};
|
||||
*occupied += 1;
|
||||
Some(*occupied)
|
||||
}
|
||||
|
||||
fn decrement_connections_for_ip(self: &mut Pin<&mut Self>, addr: &SocketAddr) {
|
||||
let should_compact = match self.connections_per_ip().entry(addr.ip()) {
|
||||
Entry::Vacant(_) => {
|
||||
error!("[{}] Got vacant entry when closing connection.", addr);
|
||||
return;
|
||||
}
|
||||
Entry::Occupied(mut occupied) => {
|
||||
*occupied.get_mut() -= 1;
|
||||
if *occupied.get() == 0 {
|
||||
occupied.remove();
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
};
|
||||
if should_compact {
|
||||
self.connections_per_ip().compact(0.1);
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_listener<C>(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<Option<io::Result<NewConnection<Req, Resp, C>>>>
|
||||
where
|
||||
S: Stream<Item = Result<C, io::Error>>,
|
||||
C: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
{
|
||||
match ready!(self.listener().poll_next_unpin(cx)?) {
|
||||
Some(codec) => Poll::Ready(Some(Ok(self.handle_new_connection(codec)))),
|
||||
None => Poll::Ready(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_closed_connections(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<io::Result<()>> {
|
||||
match ready!(self.closed_connections_rx().poll_next_unpin(cx)) {
|
||||
Some(addr) => {
|
||||
self.handle_closed_connection(&addr);
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
None => unreachable!("Holding a copy of closed_connections and didn't close it."),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, Req, Resp, T> Stream for ConnectionFilter<S, Req, Resp>
|
||||
where
|
||||
S: Stream<Item = Result<T, io::Error>>,
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
{
|
||||
type Item = io::Result<Channel<Req, Resp, T>>;
|
||||
|
||||
fn poll_next(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<Option<io::Result<Channel<Req, Resp, T>>>> {
|
||||
loop {
|
||||
match (self.poll_listener(cx)?, self.poll_closed_connections(cx)?) {
|
||||
(Poll::Ready(Some(NewConnection::Accepted(channel))), _) => {
|
||||
return Poll::Ready(Some(Ok(channel)))
|
||||
}
|
||||
(Poll::Ready(Some(NewConnection::Filtered)), _) | (_, Poll::Ready(())) => {
|
||||
trace!("Filtered a connection; {} open.", self.open_connections());
|
||||
continue;
|
||||
}
|
||||
(Poll::Pending, Poll::Pending) => return Poll::Pending,
|
||||
(Poll::Ready(None), Poll::Pending) => {
|
||||
if *self.open_connections() > 0 {
|
||||
trace!(
|
||||
"Listener closed; {} open connections.",
|
||||
self.open_connections()
|
||||
);
|
||||
return Poll::Pending;
|
||||
}
|
||||
trace!("Shutting down listener: all connections closed, and no more coming.");
|
||||
return Poll::Ready(None);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
599
rpc/src/server/mod.rs
Normal file
599
rpc/src/server/mod.rs
Normal file
@@ -0,0 +1,599 @@
|
||||
//! Provides a server that concurrently handles many connections sending multiplexed requests.
|
||||
|
||||
use crate::{
|
||||
context::Context, util::deadline_compat, util::AsDuration, util::Compact, ClientMessage,
|
||||
ClientMessageKind, Request, Response, ServerError, Transport,
|
||||
};
|
||||
use fnv::FnvHashMap;
|
||||
use futures::{
|
||||
channel::mpsc,
|
||||
future::{abortable, AbortHandle},
|
||||
prelude::*,
|
||||
ready,
|
||||
stream::Fuse,
|
||||
task::{LocalWaker, Poll},
|
||||
try_ready,
|
||||
};
|
||||
use humantime::format_rfc3339;
|
||||
use log::{debug, error, info, trace, warn};
|
||||
use pin_utils::{unsafe_pinned, unsafe_unpinned};
|
||||
use std::{
|
||||
error::Error as StdError,
|
||||
io,
|
||||
marker::PhantomData,
|
||||
net::SocketAddr,
|
||||
pin::Pin,
|
||||
time::{Instant, SystemTime},
|
||||
};
|
||||
use tokio_timer::timeout;
|
||||
use trace::{self, TraceId};
|
||||
|
||||
mod filter;
|
||||
|
||||
/// Manages clients, serving multiplexed requests over each connection.
|
||||
#[derive(Debug)]
|
||||
pub struct Server<Req, Resp> {
|
||||
config: Config,
|
||||
ghost: PhantomData<(Req, Resp)>,
|
||||
}
|
||||
|
||||
/// Settings that control the behavior of the server.
|
||||
#[non_exhaustive]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Config {
|
||||
/// The maximum number of clients that can be connected to the server at once. When at the
|
||||
/// limit, existing connections are honored and new connections are rejected.
|
||||
pub max_connections: usize,
|
||||
/// The maximum number of clients per IP address that can be connected to the server at once.
|
||||
/// When an IP is at the limit, existing connections are honored and new connections on that IP
|
||||
/// address are rejected.
|
||||
pub max_connections_per_ip: usize,
|
||||
/// The maximum number of requests that can be in flight for each client. When a client is at
|
||||
/// the in-flight request limit, existing requests are fulfilled and new requests are rejected.
|
||||
/// Rejected requests are sent a response error.
|
||||
pub max_in_flight_requests_per_connection: usize,
|
||||
/// The number of responses per client that can be buffered server-side before being sent.
|
||||
/// `pending_response_buffer` controls the buffer size of the channel that a server's
|
||||
/// response tasks use to send responses to the client handler task.
|
||||
pub pending_response_buffer: usize,
|
||||
}
|
||||
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Config {
|
||||
max_connections: 1_000_000,
|
||||
max_connections_per_ip: 1_000,
|
||||
max_in_flight_requests_per_connection: 1_000,
|
||||
pending_response_buffer: 100,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<Req, Resp> Server<Req, Resp> {
|
||||
/// Returns a new server with configuration specified `config`.
|
||||
pub fn new(config: Config) -> Self {
|
||||
Server {
|
||||
config,
|
||||
ghost: PhantomData,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the config for this server.
|
||||
pub fn config(&self) -> &Config {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Returns a stream of the incoming connections to the server.
|
||||
pub fn incoming<S, T>(
|
||||
self,
|
||||
listener: S,
|
||||
) -> impl Stream<Item = io::Result<Channel<Req, Resp, T>>>
|
||||
where
|
||||
Req: Send,
|
||||
Resp: Send,
|
||||
S: Stream<Item = io::Result<T>>,
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
{
|
||||
self::filter::ConnectionFilter::filter(listener, self.config.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// The future driving the server.
|
||||
#[derive(Debug)]
|
||||
pub struct Running<S, F> {
|
||||
incoming: S,
|
||||
request_handler: F,
|
||||
}
|
||||
|
||||
impl<S, F> Running<S, F> {
|
||||
unsafe_pinned!(incoming: S);
|
||||
unsafe_unpinned!(request_handler: F);
|
||||
}
|
||||
|
||||
impl<S, T, Req, Resp, F, Fut> Future for Running<S, F>
|
||||
where
|
||||
S: Sized + Stream<Item = io::Result<Channel<Req, Resp, T>>>,
|
||||
Req: Send + 'static,
|
||||
Resp: Send + 'static,
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send + 'static,
|
||||
F: FnMut(Context, Req) -> Fut + Send + 'static + Clone,
|
||||
Fut: Future<Output = io::Result<Resp>> + Send + 'static,
|
||||
{
|
||||
type Output = ();
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &LocalWaker) -> Poll<()> {
|
||||
while let Some(channel) = ready!(self.incoming().poll_next(cx)) {
|
||||
match channel {
|
||||
Ok(channel) => {
|
||||
let peer = channel.client_addr;
|
||||
if let Err(e) = crate::spawn(channel.respond_with(self.request_handler().clone()))
|
||||
{
|
||||
warn!("[{}] Failed to spawn connection handler: {:?}", peer, e);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("Incoming connection error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
info!("Server shutting down.");
|
||||
return Poll::Ready(());
|
||||
}
|
||||
}
|
||||
|
||||
/// A utility trait enabling a stream to fluently chain a request handler.
|
||||
pub trait Handler<T, Req, Resp>
|
||||
where
|
||||
Self: Sized + Stream<Item = io::Result<Channel<Req, Resp, T>>>,
|
||||
Req: Send,
|
||||
Resp: Send,
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
{
|
||||
/// Responds to all requests with `request_handler`.
|
||||
fn respond_with<F, Fut>(self, request_handler: F) -> Running<Self, F>
|
||||
where
|
||||
F: FnMut(Context, Req) -> Fut + Send + 'static + Clone,
|
||||
Fut: Future<Output = io::Result<Resp>> + Send + 'static,
|
||||
{
|
||||
Running {
|
||||
incoming: self,
|
||||
request_handler,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, Req, Resp, S> Handler<T, Req, Resp> for S
|
||||
where
|
||||
S: Sized + Stream<Item = io::Result<Channel<Req, Resp, T>>>,
|
||||
Req: Send,
|
||||
Resp: Send,
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
{}
|
||||
|
||||
/// Responds to all requests with `request_handler`.
|
||||
/// The server end of an open connection with a client.
|
||||
#[derive(Debug)]
|
||||
pub struct Channel<Req, Resp, T> {
|
||||
/// Writes responses to the wire and reads requests off the wire.
|
||||
transport: Fuse<T>,
|
||||
/// Signals the connection is closed when `Channel` is dropped.
|
||||
closed_connections: mpsc::UnboundedSender<SocketAddr>,
|
||||
/// Channel limits to prevent unlimited resource usage.
|
||||
config: Config,
|
||||
/// The address of the server connected to.
|
||||
client_addr: SocketAddr,
|
||||
/// Types the request and response.
|
||||
ghost: PhantomData<(Req, Resp)>,
|
||||
}
|
||||
|
||||
impl<Req, Resp, T> Drop for Channel<Req, Resp, T> {
|
||||
fn drop(&mut self) {
|
||||
trace!("[{}] Closing channel.", self.client_addr);
|
||||
|
||||
// Even in a bounded channel, each connection would have a guaranteed slot, so using
|
||||
// an unbounded sender is actually no different. And, the bound is on the maximum number
|
||||
// of open connections.
|
||||
if self
|
||||
.closed_connections
|
||||
.unbounded_send(self.client_addr)
|
||||
.is_err()
|
||||
{
|
||||
warn!(
|
||||
"[{}] Failed to send closed connection message.",
|
||||
self.client_addr
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<Req, Resp, T> Channel<Req, Resp, T> {
|
||||
unsafe_pinned!(transport: Fuse<T>);
|
||||
}
|
||||
|
||||
impl<Req, Resp, T> Channel<Req, Resp, T>
|
||||
where
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
Req: Send,
|
||||
Resp: Send,
|
||||
{
|
||||
pub(crate) fn start_send(self: &mut Pin<&mut Self>, response: Response<Resp>) -> io::Result<()> {
|
||||
self.transport().start_send(response)
|
||||
}
|
||||
|
||||
pub(crate) fn poll_ready(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<io::Result<()>> {
|
||||
self.transport().poll_ready(cx)
|
||||
}
|
||||
|
||||
pub(crate) fn poll_flush(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<io::Result<()>> {
|
||||
self.transport().poll_flush(cx)
|
||||
}
|
||||
|
||||
pub(crate) fn poll_next(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<Option<io::Result<ClientMessage<Req>>>> {
|
||||
self.transport().poll_next(cx)
|
||||
}
|
||||
|
||||
/// Returns the address of the client connected to the channel.
|
||||
pub fn client_addr(&self) -> &SocketAddr {
|
||||
&self.client_addr
|
||||
}
|
||||
|
||||
/// Respond to requests coming over the channel with `f`. Returns a future that drives the
|
||||
/// responses and resolves when the connection is closed.
|
||||
pub fn respond_with<F, Fut>(self, f: F) -> impl Future<Output = ()>
|
||||
where
|
||||
F: FnMut(Context, Req) -> Fut + Send + 'static,
|
||||
Fut: Future<Output = io::Result<Resp>> + Send + 'static,
|
||||
Req: 'static,
|
||||
Resp: 'static,
|
||||
{
|
||||
let (responses_tx, responses) = mpsc::channel(self.config.pending_response_buffer);
|
||||
let responses = responses.fuse();
|
||||
let peer = self.client_addr;
|
||||
|
||||
ClientHandler {
|
||||
channel: self,
|
||||
f,
|
||||
pending_responses: responses,
|
||||
responses_tx,
|
||||
in_flight_requests: FnvHashMap::default(),
|
||||
}.unwrap_or_else(move |e| {
|
||||
info!("[{}] ClientHandler errored out: {}", peer, e);
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ClientHandler<Req, Resp, T, F> {
|
||||
channel: Channel<Req, Resp, T>,
|
||||
/// Responses waiting to be written to the wire.
|
||||
pending_responses: Fuse<mpsc::Receiver<(Context, Response<Resp>)>>,
|
||||
/// Handed out to request handlers to fan in responses.
|
||||
responses_tx: mpsc::Sender<(Context, Response<Resp>)>,
|
||||
/// Number of requests currently being responded to.
|
||||
in_flight_requests: FnvHashMap<u64, AbortHandle>,
|
||||
/// Request handler.
|
||||
f: F,
|
||||
}
|
||||
|
||||
impl<Req, Resp, T, F> ClientHandler<Req, Resp, T, F> {
|
||||
unsafe_pinned!(channel: Channel<Req, Resp, T>);
|
||||
unsafe_pinned!(in_flight_requests: FnvHashMap<u64, AbortHandle>);
|
||||
unsafe_pinned!(pending_responses: Fuse<mpsc::Receiver<(Context, Response<Resp>)>>);
|
||||
unsafe_pinned!(responses_tx: mpsc::Sender<(Context, Response<Resp>)>);
|
||||
// For this to be safe, field f must be private, and code in this module must never
|
||||
// construct PinMut<F>.
|
||||
unsafe_unpinned!(f: F);
|
||||
}
|
||||
|
||||
impl<Req, Resp, T, F, Fut> ClientHandler<Req, Resp, T, F>
|
||||
where
|
||||
Req: Send + 'static,
|
||||
Resp: Send + 'static,
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
F: FnMut(Context, Req) -> Fut + Send + 'static,
|
||||
Fut: Future<Output = io::Result<Resp>> + Send + 'static,
|
||||
{
|
||||
/// If at max in-flight requests, check that there's room to immediately write a throttled
|
||||
/// response.
|
||||
fn poll_ready_if_throttling(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<io::Result<()>> {
|
||||
if self.in_flight_requests.len()
|
||||
>= self.channel.config.max_in_flight_requests_per_connection
|
||||
{
|
||||
let peer = self.channel().client_addr;
|
||||
|
||||
while let Poll::Pending = self.channel().poll_ready(cx)? {
|
||||
info!(
|
||||
"[{}] In-flight requests at max ({}), and transport is not ready.",
|
||||
peer,
|
||||
self.in_flight_requests().len(),
|
||||
);
|
||||
try_ready!(self.channel().poll_flush(cx));
|
||||
}
|
||||
}
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn pump_read(self: &mut Pin<&mut Self>, cx: &LocalWaker) -> Poll<Option<io::Result<()>>> {
|
||||
ready!(self.poll_ready_if_throttling(cx)?);
|
||||
|
||||
Poll::Ready(match ready!(self.channel().poll_next(cx)?) {
|
||||
Some(message) => {
|
||||
match message.message {
|
||||
ClientMessageKind::Request(request) => {
|
||||
self.handle_request(message.trace_context, request)?;
|
||||
}
|
||||
ClientMessageKind::Cancel { request_id } => {
|
||||
self.cancel_request(&message.trace_context, request_id);
|
||||
}
|
||||
}
|
||||
Some(Ok(()))
|
||||
}
|
||||
None => {
|
||||
trace!("[{}] Read half closed", self.channel.client_addr);
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn pump_write(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
read_half_closed: bool,
|
||||
) -> Poll<Option<io::Result<()>>> {
|
||||
match self.poll_next_response(cx)? {
|
||||
Poll::Ready(Some((_, response))) => {
|
||||
self.channel().start_send(response)?;
|
||||
Poll::Ready(Some(Ok(())))
|
||||
}
|
||||
Poll::Ready(None) => {
|
||||
// Shutdown can't be done before we finish pumping out remaining responses.
|
||||
ready!(self.channel().poll_flush(cx)?);
|
||||
Poll::Ready(None)
|
||||
}
|
||||
Poll::Pending => {
|
||||
// No more requests to process, so flush any requests buffered in the transport.
|
||||
ready!(self.channel().poll_flush(cx)?);
|
||||
|
||||
// Being here means there are no staged requests and all written responses are
|
||||
// fully flushed. So, if the read half is closed and there are no in-flight
|
||||
// requests, then we can close the write half.
|
||||
if read_half_closed && self.in_flight_requests().is_empty() {
|
||||
Poll::Ready(None)
|
||||
} else {
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_next_response(
|
||||
self: &mut Pin<&mut Self>,
|
||||
cx: &LocalWaker,
|
||||
) -> Poll<Option<io::Result<(Context, Response<Resp>)>>> {
|
||||
// Ensure there's room to write a response.
|
||||
while let Poll::Pending = self.channel().poll_ready(cx)? {
|
||||
ready!(self.channel().poll_flush(cx)?);
|
||||
}
|
||||
|
||||
let peer = self.channel().client_addr;
|
||||
|
||||
match ready!(self.pending_responses().poll_next(cx)) {
|
||||
Some((ctx, response)) => {
|
||||
if let Some(_) = self.in_flight_requests().remove(&response.request_id) {
|
||||
self.in_flight_requests().compact(0.1);
|
||||
}
|
||||
trace!(
|
||||
"[{}/{}] Staging response. In-flight requests = {}.",
|
||||
ctx.trace_id(),
|
||||
peer,
|
||||
self.in_flight_requests().len(),
|
||||
);
|
||||
return Poll::Ready(Some(Ok((ctx, response))));
|
||||
}
|
||||
None => {
|
||||
// This branch likely won't happen, since the ClientHandler is holding a Sender.
|
||||
trace!("[{}] No new responses.", peer);
|
||||
Poll::Ready(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_request(
|
||||
self: &mut Pin<&mut Self>,
|
||||
trace_context: trace::Context,
|
||||
request: Request<Req>,
|
||||
) -> io::Result<()> {
|
||||
let request_id = request.id;
|
||||
let peer = self.channel().client_addr;
|
||||
let ctx = Context {
|
||||
deadline: request.deadline,
|
||||
trace_context,
|
||||
};
|
||||
let request = request.message;
|
||||
|
||||
if self.in_flight_requests().len()
|
||||
>= self.channel().config.max_in_flight_requests_per_connection
|
||||
{
|
||||
debug!(
|
||||
"[{}/{}] Client has reached in-flight request limit ({}/{}).",
|
||||
ctx.trace_id(),
|
||||
peer,
|
||||
self.in_flight_requests().len(),
|
||||
self.channel().config.max_in_flight_requests_per_connection
|
||||
);
|
||||
|
||||
self.channel().start_send(Response {
|
||||
request_id,
|
||||
message: Err(ServerError {
|
||||
kind: io::ErrorKind::WouldBlock,
|
||||
detail: Some("Server throttled the request.".into()),
|
||||
}),
|
||||
})?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let deadline = ctx.deadline;
|
||||
let timeout = deadline.as_duration();
|
||||
trace!(
|
||||
"[{}/{}] Received request with deadline {} (timeout {:?}).",
|
||||
ctx.trace_id(),
|
||||
peer,
|
||||
format_rfc3339(deadline),
|
||||
timeout,
|
||||
);
|
||||
let mut response_tx = self.responses_tx().clone();
|
||||
|
||||
let trace_id = *ctx.trace_id();
|
||||
let response = self.f()(ctx.clone(), request);
|
||||
let response = deadline_compat::Deadline::new(response, Instant::now() + timeout).then(
|
||||
async move |result| {
|
||||
let response = Response {
|
||||
request_id,
|
||||
message: match result {
|
||||
Ok(message) => Ok(message),
|
||||
Err(e) => Err(make_server_error(e, trace_id, peer, deadline)),
|
||||
},
|
||||
};
|
||||
trace!("[{}/{}] Sending response.", trace_id, peer);
|
||||
await!(response_tx.send((ctx, response)).unwrap_or_else(|_| ()));
|
||||
},
|
||||
);
|
||||
let (abortable_response, abort_handle) = abortable(response);
|
||||
crate::spawn(abortable_response.map(|_| ()))
|
||||
.map_err(|e| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::Other,
|
||||
format!(
|
||||
"Could not spawn response task. Is shutdown: {}",
|
||||
e.is_shutdown()
|
||||
),
|
||||
)
|
||||
})?;
|
||||
self.in_flight_requests().insert(request_id, abort_handle);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn cancel_request(self: &mut Pin<&mut Self>, trace_context: &trace::Context, request_id: u64) {
|
||||
// It's possible the request was already completed, so it's fine
|
||||
// if this is None.
|
||||
if let Some(cancel_handle) = self.in_flight_requests().remove(&request_id) {
|
||||
self.in_flight_requests().compact(0.1);
|
||||
|
||||
cancel_handle.abort();
|
||||
let remaining = self.in_flight_requests().len();
|
||||
trace!(
|
||||
"[{}/{}] Request canceled. In-flight requests = {}",
|
||||
trace_context.trace_id,
|
||||
self.channel.client_addr,
|
||||
remaining,
|
||||
);
|
||||
} else {
|
||||
trace!(
|
||||
"[{}/{}] Received cancellation, but response handler \
|
||||
is already complete.",
|
||||
trace_context.trace_id,
|
||||
self.channel.client_addr
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<Req, Resp, T, F, Fut> Future for ClientHandler<Req, Resp, T, F>
|
||||
where
|
||||
Req: Send + 'static,
|
||||
Resp: Send + 'static,
|
||||
T: Transport<Item = ClientMessage<Req>, SinkItem = Response<Resp>> + Send,
|
||||
F: FnMut(Context, Req) -> Fut + Send + 'static,
|
||||
Fut: Future<Output = io::Result<Resp>> + Send + 'static,
|
||||
{
|
||||
type Output = io::Result<()>;
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &LocalWaker) -> Poll<io::Result<()>> {
|
||||
trace!("[{}] ClientHandler::poll", self.channel.client_addr);
|
||||
loop {
|
||||
let read = self.pump_read(cx)?;
|
||||
match (read, self.pump_write(cx, read == Poll::Ready(None))?) {
|
||||
(Poll::Ready(None), Poll::Ready(None)) => {
|
||||
info!("[{}] Client disconnected.", self.channel.client_addr);
|
||||
return Poll::Ready(Ok(()));
|
||||
}
|
||||
(read @ Poll::Ready(Some(())), write) | (read, write @ Poll::Ready(Some(()))) => {
|
||||
trace!(
|
||||
"[{}] read: {:?}, write: {:?}.",
|
||||
self.channel.client_addr,
|
||||
read,
|
||||
write
|
||||
)
|
||||
}
|
||||
(read, write) => {
|
||||
trace!(
|
||||
"[{}] read: {:?}, write: {:?} (not ready).",
|
||||
self.channel.client_addr,
|
||||
read,
|
||||
write,
|
||||
);
|
||||
return Poll::Pending;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn make_server_error(
|
||||
e: timeout::Error<io::Error>,
|
||||
trace_id: TraceId,
|
||||
peer: SocketAddr,
|
||||
deadline: SystemTime,
|
||||
) -> ServerError {
|
||||
if e.is_elapsed() {
|
||||
debug!(
|
||||
"[{}/{}] Response did not complete before deadline of {}s.",
|
||||
trace_id,
|
||||
peer,
|
||||
format_rfc3339(deadline)
|
||||
);
|
||||
// No point in responding, since the client will have dropped the request.
|
||||
ServerError {
|
||||
kind: io::ErrorKind::TimedOut,
|
||||
detail: Some(format!(
|
||||
"Response did not complete before deadline of {}s.",
|
||||
format_rfc3339(deadline)
|
||||
)),
|
||||
}
|
||||
} else if e.is_timer() {
|
||||
error!(
|
||||
"[{}/{}] Response failed because of an issue with a timer: {}",
|
||||
trace_id, peer, e
|
||||
);
|
||||
|
||||
ServerError {
|
||||
kind: io::ErrorKind::Other,
|
||||
detail: Some(format!("{}", e)),
|
||||
}
|
||||
} else if e.is_inner() {
|
||||
let e = e.into_inner().unwrap();
|
||||
ServerError {
|
||||
kind: e.kind(),
|
||||
detail: Some(e.description().into()),
|
||||
}
|
||||
} else {
|
||||
error!("[{}/{}] Unexpected response failure: {}", trace_id, peer, e);
|
||||
|
||||
ServerError {
|
||||
kind: io::ErrorKind::Other,
|
||||
detail: Some(format!("Server unexpectedly failed to respond: {}", e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user