Skip to content

Commit

Permalink
Move request hooks off of the Serve trait.
Browse files Browse the repository at this point in the history 
They now exist on an extension trait called RequestHook.
  • Loading branch information
@tikue
tikue committed Feb 5, 2024
1 parent 5fd6138 commit b413559
Show file tree
Hide file tree
Showing 2 changed files with 147 additions and 143 deletions.
146 changes: 3 additions & 143 deletions tarpc/src/server.rs
View file
Expand Up @@ -39,10 +39,6 @@ pub mod limits;
/// Provides helper methods for streams of Channels.
pub mod incoming;

use request_hook::{
AfterRequest, BeforeRequest, HookThenServe, HookThenServeThenHook, ServeThenHook,
};

/// Settings that control the behavior of [channels](Channel).
#[derive(Clone, Debug)]
pub struct Config {
Expand Down Expand Up @@ -81,143 +77,6 @@ pub trait Serve {

/// Responds to a single request.
async fn serve(self, ctx: context::Context, req: Self::Req) -> Result<Self::Resp, ServerError>;

/// Runs a hook before execution of the request.
///
/// If the hook returns an error, the request will not be executed and the error will be
/// returned instead.
///
/// The hook can also modify the request context. This could be used, for example, to enforce a
/// maximum deadline on all requests.
///
/// Any type that implements [`BeforeRequest`] can be used as the hook. Types that implement
/// `FnMut(&mut Context, &RequestType) -> impl Future<Output = Result<(), ServerError>>` can
/// also be used.
///
/// # Example
///
/// ```rust
/// use futures::{executor::block_on, future};
/// use tarpc::{context, ServerError, server::{Serve, serve}};
/// use std::io;
///
/// let serve = serve(|_ctx, i| async move { Ok(i + 1) })
/// .before(|_ctx: &mut context::Context, req: &i32| {
/// future::ready(
/// if *req == 1 {
/// Err(ServerError::new(
/// io::ErrorKind::Other,
/// format!("I don't like {req}")))
/// } else {
/// Ok(())
/// })
/// });
/// let response = serve.serve(context::current(), 1);
/// assert!(block_on(response).is_err());
/// ```
fn before<Hook>(self, hook: Hook) -> HookThenServe<Self, Hook>
where
Hook: BeforeRequest<Self::Req>,
Self: Sized,
{
HookThenServe::new(self, hook)
}

/// Runs a hook after completion of a request.
///
/// The hook can modify the request context and the response.
///
/// Any type that implements [`AfterRequest`] can be used as the hook. Types that implement
/// `FnMut(&mut Context, &mut Result<ResponseType, ServerError>) -> impl Future<Output = ()>`
/// can also be used.
///
/// # Example
///
/// ```rust
/// use futures::{executor::block_on, future};
/// use tarpc::{context, ServerError, server::{Serve, serve}};
/// use std::io;
///
/// let serve = serve(
/// |_ctx, i| async move {
/// if i == 1 {
/// Err(ServerError::new(
/// io::ErrorKind::Other,
/// format!("{i} is the loneliest number")))
/// } else {
/// Ok(i + 1)
/// }
/// })
/// .after(|_ctx: &mut context::Context, resp: &mut Result<i32, ServerError>| {
/// if let Err(e) = resp {
/// eprintln!("server error: {e:?}");
/// }
/// future::ready(())
/// });
///
/// let response = serve.serve(context::current(), 1);
/// assert!(block_on(response).is_err());
/// ```
fn after<Hook>(self, hook: Hook) -> ServeThenHook<Self, Hook>
where
Hook: AfterRequest<Self::Resp>,
Self: Sized,
{
ServeThenHook::new(self, hook)
}

/// Runs a hook before and after execution of the request.
///
/// If the hook returns an error, the request will not be executed and the error will be
/// returned instead.
///
/// The hook can also modify the request context and the response. This could be used, for
/// example, to enforce a maximum deadline on all requests.
///
/// # Example
///
/// ```rust
/// use futures::{executor::block_on, future};
/// use tarpc::{
/// context, ServerError, server::{Serve, serve, request_hook::{BeforeRequest, AfterRequest}}
/// };
/// use std::{io, time::Instant};
///
/// struct PrintLatency(Instant);
///
/// impl<Req> BeforeRequest<Req> for PrintLatency {
/// async fn before(&mut self, _: &mut context::Context, _: &Req) -> Result<(), ServerError> {
/// self.0 = Instant::now();
/// Ok(())
/// }
/// }
///
/// impl<Resp> AfterRequest<Resp> for PrintLatency {
/// async fn after(
/// &mut self,
/// _: &mut context::Context,
/// _: &mut Result<Resp, ServerError>,
/// ) {
/// tracing::info!("Elapsed: {:?}", self.0.elapsed());
/// }
/// }
///
/// let serve = serve(|_ctx, i| async move {
/// Ok(i + 1)
/// }).before_and_after(PrintLatency(Instant::now()));
/// let response = serve.serve(context::current(), 1);
/// assert!(block_on(response).is_ok());
/// ```
fn before_and_after<Hook>(
self,
hook: Hook,
) -> HookThenServeThenHook<Self::Req, Self::Resp, Self, Hook>
where
Hook: BeforeRequest<Self::Req> + AfterRequest<Self::Resp>,
Self: Sized,
{
HookThenServeThenHook::new(self, hook)
}
}

/// A Serve wrapper around a Fn.
Expand Down Expand Up @@ -1096,8 +955,9 @@ where
#[cfg(test)]
mod tests {
use super::{
in_flight_requests::AlreadyExistsError, serve, AfterRequest, BaseChannel, BeforeRequest,
Channel, Config, Requests, Serve,
in_flight_requests::AlreadyExistsError,
request_hook::{AfterRequest, BeforeRequest, RequestHook},
serve, BaseChannel, Channel, Config, Requests, Serve,
};
use crate::{
context, trace,
Expand Down
144 changes: 144 additions & 0 deletions tarpc/src/server/request_hook.rs
View file
Expand Up @@ -6,6 +6,8 @@

//! Hooks for horizontal functionality that can run either before or after a request is executed.

use crate::server::Serve;

/// A request hook that runs before a request is executed.
mod before;

Expand All @@ -23,3 +25,145 @@ pub use {
},
before_and_after::HookThenServeThenHook,
};

/// Hooks that run before and/or after serving a request.
pub trait RequestHook: Serve {
/// Runs a hook before execution of the request.
///
/// If the hook returns an error, the request will not be executed and the error will be
/// returned instead.
///
/// The hook can also modify the request context. This could be used, for example, to enforce a
/// maximum deadline on all requests.
///
/// Any type that implements [`BeforeRequest`] can be used as the hook. Types that implement
/// `FnMut(&mut Context, &RequestType) -> impl Future<Output = Result<(), ServerError>>` can
/// also be used.
///
/// # Example
///
/// ```rust
/// use futures::{executor::block_on, future};
/// use tarpc::{context, ServerError, server::{Serve, request_hook::RequestHook, serve}};
/// use std::io;
///
/// let serve = serve(|_ctx, i| async move { Ok(i + 1) })
/// .before(|_ctx: &mut context::Context, req: &i32| {
/// future::ready(
/// if *req == 1 {
/// Err(ServerError::new(
/// io::ErrorKind::Other,
/// format!("I don't like {req}")))
/// } else {
/// Ok(())
/// })
/// });
/// let response = serve.serve(context::current(), 1);
/// assert!(block_on(response).is_err());
/// ```
fn before<Hook>(self, hook: Hook) -> HookThenServe<Self, Hook>
where
Hook: BeforeRequest<Self::Req>,
Self: Sized,
{
HookThenServe::new(self, hook)
}

/// Runs a hook after completion of a request.
///
/// The hook can modify the request context and the response.
///
/// Any type that implements [`AfterRequest`] can be used as the hook. Types that implement
/// `FnMut(&mut Context, &mut Result<ResponseType, ServerError>) -> impl Future<Output = ()>`
/// can also be used.
///
/// # Example
///
/// ```rust
/// use futures::{executor::block_on, future};
/// use tarpc::{context, ServerError, server::{Serve, request_hook::RequestHook, serve}};
/// use std::io;
///
/// let serve = serve(
/// |_ctx, i| async move {
/// if i == 1 {
/// Err(ServerError::new(
/// io::ErrorKind::Other,
/// format!("{i} is the loneliest number")))
/// } else {
/// Ok(i + 1)
/// }
/// })
/// .after(|_ctx: &mut context::Context, resp: &mut Result<i32, ServerError>| {
/// if let Err(e) = resp {
/// eprintln!("server error: {e:?}");
/// }
/// future::ready(())
/// });
///
/// let response = serve.serve(context::current(), 1);
/// assert!(block_on(response).is_err());
/// ```
fn after<Hook>(self, hook: Hook) -> ServeThenHook<Self, Hook>
where
Hook: AfterRequest<Self::Resp>,
Self: Sized,
{
ServeThenHook::new(self, hook)
}

/// Runs a hook before and after execution of the request.
///
/// If the hook returns an error, the request will not be executed and the error will be
/// returned instead.
///
/// The hook can also modify the request context and the response. This could be used, for
/// example, to enforce a maximum deadline on all requests.
///
/// # Example
///
/// ```rust
/// use futures::{executor::block_on, future};
/// use tarpc::{
/// context, ServerError,
/// server::{Serve, serve, request_hook::{BeforeRequest, AfterRequest, RequestHook}}
/// };
/// use std::{io, time::Instant};
///
/// struct PrintLatency(Instant);
///
/// impl<Req> BeforeRequest<Req> for PrintLatency {
/// async fn before(&mut self, _: &mut context::Context, _: &Req) -> Result<(), ServerError> {
/// self.0 = Instant::now();
/// Ok(())
/// }
/// }
///
/// impl<Resp> AfterRequest<Resp> for PrintLatency {
/// async fn after(
/// &mut self,
/// _: &mut context::Context,
/// _: &mut Result<Resp, ServerError>,
/// ) {
/// tracing::info!("Elapsed: {:?}", self.0.elapsed());
/// }
/// }
///
/// let serve = serve(|_ctx, i| async move {
/// Ok(i + 1)
/// }).before_and_after(PrintLatency(Instant::now()));
/// let response = serve.serve(context::current(), 1);
/// assert!(block_on(response).is_ok());
/// ```
fn before_and_after<Hook>(
self,
hook: Hook,
) -> HookThenServeThenHook<Self::Req, Self::Resp, Self, Hook>
where
Hook: BeforeRequest<Self::Req> + AfterRequest<Self::Resp>,
Self: Sized,
{
HookThenServeThenHook::new(self, hook)
}
}
impl<S: Serve> RequestHook for S {}

0 comments on commit b413559

Please sign in to comment.