Skip to content

Commit

Permalink
work: cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@prestwich
prestwich committed Sep 30, 2024
1 parent e452013 commit 8c6d7b5
Showing 1 changed file with 163 additions and 115 deletions.
278 changes: 163 additions & 115 deletions crates/json-rpc-router/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,89 +57,6 @@ impl<S> fmt::Debug for Router<S> {
}
}

/// A unique internal identifier for a method.
#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct MethodId(usize);

impl From<usize> for MethodId {
fn from(id: usize) -> Self {
Self(id)
}
}

impl Add for MethodId {
type Output = Self;

fn add(self, rhs: Self) -> Self::Output {
Self(self.0 + rhs.0)
}
}

impl Add<usize> for MethodId {
type Output = Self;

fn add(self, rhs: usize) -> Self::Output {
Self(self.0 + rhs)
}
}

impl AddAssign for MethodId {
fn add_assign(&mut self, rhs: Self) {
self.0 += rhs.0;
}
}

impl AddAssign<usize> for MethodId {
fn add_assign(&mut self, rhs: usize) {
self.0 += rhs;
}
}

/// A boxed, erased type that can be converted into a Handler. Similar to
/// axum's `ErasedIntoRoute`
///
/// Currently this is a placeholder to enable future convenience functions
pub trait ErasedIntoRoute<S>: Send {
/// Take a reference to this type, clone it, box it, and type erase it.
///
/// This allows it to be stored in a collection of `dyn
/// ErasedIntoRoute<S>`.
fn clone_box(&self) -> Box<dyn ErasedIntoRoute<S>>;

/// Convert this type into a handler.
fn into_route(self: Box<Self>, state: S) -> Route;

/// Call this handler with the given state.
#[allow(dead_code)]
fn call_with_state(
self: Box<Self>,
params: Box<RawValue>,
state: S,
) -> <Route as Service<Box<RawValue>>>::Future;
}

/// A boxed, erased type that can be converted into a handler.
///
/// Similar to axum's `BoxedIntoRoute`
struct BoxedIntoHandler<S>(Box<dyn ErasedIntoRoute<S>>);

impl<S> Clone for BoxedIntoHandler<S> {
fn clone(&self) -> Self {
Self(self.0.clone_box())
}
}

/// A method, which may be ready to handle requests or may need to be
/// initialized with some state.
///
/// Analagous to axum's `MethodEndpoint`
enum Method<S> {
/// A method that needs to be initialized with some state.
Needs(BoxedIntoHandler<S>),
/// A method that is ready to handle requests.
Ready(Route),
}

/// The inner state of a [`Router`]. Maps methods to their handlers.
#[derive(Default)]
pub struct RouterInner<S> {
Expand Down Expand Up @@ -167,6 +84,23 @@ impl<S> RouterInner<S> {
}
}

/// Add state to the router, readying methods that require that state.
pub fn with_state<S2>(self, state: &S) -> RouterInner<S2>
where
S: Clone,
{
RouterInner {
routes: self
.routes
.into_iter()
.map(|(id, method)| (id, method.with_state(state)))
.collect(),
last_id: self.last_id,
name_to_id: self.name_to_id,
id_to_name: self.id_to_name,
}
}

/// Get the next available ID.
fn get_id(&mut self) -> MethodId {
self.last_id += 1;
Expand All @@ -178,6 +112,8 @@ impl<S> RouterInner<S> {
self.name_to_id.get(name).and_then(|id| self.routes.get(id))
}

/// Enroll a method name, returning an ID assignment. Panics if the method
/// name already exists in the router. Future: don't panic.
fn enroll_method_name(&mut self, method: Cow<'static, str>) -> MethodId {
if self.name_to_id.contains_key(&method) {
panic!("Method name already exists in the router.");
Expand All @@ -189,8 +125,20 @@ impl<S> RouterInner<S> {
id
}

/// Add a handler to the router. This method is complete and ready to call.
fn route(mut self, method: impl Into<Cow<'static, str>>, handler: Route) -> Self {
let method = method.into();
let id = self.get_id();

self.name_to_id.insert(method.clone(), id);
self.id_to_name.insert(id, method.clone());
self.routes.insert(id, Method::Ready(handler));

self
}

/// Add a method to the router. This method may be missing state `S`.
pub fn add_into_route<H>(mut self, method: impl Into<Cow<'static, str>>, handler: H) -> Self
pub fn route_erased<H>(mut self, method: impl Into<Cow<'static, str>>, handler: H) -> Self
where
H: ErasedIntoRoute<S>,
{
Expand All @@ -212,18 +160,6 @@ impl<S> RouterInner<S> {
self
}

/// Add a handler to the router. This method is complete and ready to call.
pub fn add_route(mut self, method: impl Into<Cow<'static, str>>, handler: Route) -> Self {
let method = method.into();
let id = self.get_id();

self.name_to_id.insert(method.clone(), id);
self.id_to_name.insert(id, method.clone());
self.routes.insert(id, Method::Ready(handler));

self
}

/// Add a service to the router.
pub fn route_service<T>(self, method: impl Into<Cow<'static, str>>, service: T) -> Self
where
Expand All @@ -236,7 +172,7 @@ impl<S> RouterInner<S> {
+ Send
+ 'static,
{
self.add_route(method, BoxCloneService::new(service))
self.route(method, BoxCloneService::new(service))
}

/// Call a method on the router, with the provided state.
Expand Down Expand Up @@ -267,8 +203,89 @@ impl<S> RouterInner<S> {
}
}

trait Captures<'a> {}
impl<'a, T: ?Sized> Captures<'a> for T {}
/// A unique internal identifier for a method.
#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct MethodId(usize);

impl From<usize> for MethodId {
fn from(id: usize) -> Self {
Self(id)
}
}

impl Add<usize> for MethodId {
type Output = Self;

fn add(self, rhs: usize) -> Self::Output {
Self(self.0 + rhs)
}
}

impl AddAssign<usize> for MethodId {
fn add_assign(&mut self, rhs: usize) {
self.0 += rhs;
}
}

/// A boxed, erased type that can be converted into a Handler. Similar to
/// axum's `ErasedIntoRoute`
///
/// Currently this is a placeholder to enable future convenience functions
pub trait ErasedIntoRoute<S>: Send {
/// Take a reference to this type, clone it, box it, and type erase it.
///
/// This allows it to be stored in a collection of `dyn
/// ErasedIntoRoute<S>`.
fn clone_box(&self) -> Box<dyn ErasedIntoRoute<S>>;

/// Convert this type into a handler.
fn into_route(self: Box<Self>, state: S) -> Route;

/// Call this handler with the given state.
#[allow(dead_code)]
fn call_with_state(
self: Box<Self>,
params: Box<RawValue>,
state: S,
) -> <Route as Service<Box<RawValue>>>::Future;
}

/// A boxed, erased type that can be converted into a handler.
///
/// Similar to axum's `BoxedIntoRoute`
///
/// Currently this is a placeholder to enable future convenience functions
struct BoxedIntoHandler<S>(Box<dyn ErasedIntoRoute<S>>);

impl<S> Clone for BoxedIntoHandler<S> {
fn clone(&self) -> Self {
Self(self.0.clone_box())
}
}

/// A method, which may be ready to handle requests or may need to be
/// initialized with some state.
///
/// Analagous to axum's `MethodEndpoint`
enum Method<S> {
/// A method that needs to be initialized with some state.
Needs(BoxedIntoHandler<S>),
/// A method that is ready to handle requests.
Ready(Route),
}

impl<S> Method<S>
where
S: Clone,
{
/// Add state to a method, converting
fn with_state<S2>(self, state: &S) -> Method<S2> {
match self {
Self::Ready(route) => Method::Ready(route),
Self::Needs(handler) => Method::Ready(handler.0.into_route(state.clone())),
}
}
}

/// A handler for a JSON-RPC method.
pub trait Handler<T, S>: Clone + Send + Sized + 'static {
Expand Down Expand Up @@ -333,7 +350,7 @@ where
impl<F, Fut, Params, Payload, ErrData, S> Handler<(Params,), S> for F
where
F: FnOnce(Params) -> Fut + Clone + Send + 'static,
Fut: Future<Output = ResponsePayload<Payload, ErrData>> + Send + 'static,
Fut: Future<Output = Result<Payload, ErrData>> + Send + 'static,
Params: RpcObject,
Payload: RpcObject,
ErrData: RpcObject,
Expand All @@ -348,6 +365,8 @@ where

self(params)
.await
.map(ResponsePayload::Success)
.unwrap_or_else(ResponsePayload::internal_error_with_obj)
.serialize_payload()
.ok()
.unwrap_or_else(ResponsePayload::internal_error)
Expand All @@ -358,7 +377,7 @@ where
impl<F, Fut, Params, Payload, ErrData, S> Handler<(Params, S), S> for F
where
F: FnOnce(Params, S) -> Fut + Clone + Send + 'static,
Fut: Future<Output = ResponsePayload<Payload, ErrData>> + Send + 'static,
Fut: Future<Output = Result<Payload, ErrData>> + Send + 'static,
Params: RpcObject,
Payload: RpcObject,
ErrData: RpcObject,
Expand All @@ -372,13 +391,22 @@ where
return ResponsePayload::invalid_params();
};

self(params, state).await.serialize_payload().ok().unwrap_or_else(|| {
ResponsePayload::internal_error_message("Failed to serialize response".into())
})
self(params, state)
.await
.map(ResponsePayload::Success)
.unwrap_or_else(ResponsePayload::internal_error_with_obj)
.serialize_payload()
.ok()
.unwrap_or_else(|| {
ResponsePayload::internal_error_message("Failed to serialize response".into())
})
})
}
}

trait Captures<'a> {}
impl<'a, T: ?Sized> Captures<'a> for T {}

#[cfg(test)]
mod test {
use alloy_json_rpc::ErrorPayload;
Expand All @@ -388,17 +416,25 @@ mod test {

#[tokio::test]
async fn example() {
let router: RouterInner<()> = RouterInner::new().route_service(
"hello_world",
tower::service_fn(|_: Box<RawValue>| async {
Ok(ResponsePayload::<(), u8>::internal_error_with_message_and_obj(
Cow::Borrowed("Hello, world!"),
30u8,
)
.serialize_payload()
.unwrap())
}),
);
let router: RouterInner<()> = RouterInner::new()
.route_service(
"hello_world",
tower::service_fn(|_: Box<RawValue>| async {
Ok(ResponsePayload::<(), u8>::internal_error_with_message_and_obj(
Cow::Borrowed("Hello, world!"),
30u8,
)
.serialize_payload()
.unwrap())
}),
)
.route_service(
"foo",
Handler::with_state(
|a: Box<RawValue>| async move { Ok::<_, ()>(a.get().to_owned()) },
(),
),
);

let res = router
.call_with_state("hello_world", Default::default(), ())
Expand All @@ -414,5 +450,17 @@ mod test {
data: Some(30u8)
})
),);

let res2 = dbg!(
router
.call_with_state("foo", RawValue::from_string("{}".to_owned()).unwrap(), ())
.await
)
.deserialize_success::<String>()
.unwrap()
.deserialize_error::<()>()
.unwrap();

assert_eq!(res2, ResponsePayload::Success("{}".to_owned()));
}
}

0 comments on commit 8c6d7b5

Please sign in to comment.