configurable middleware and metrics & http middleware (#9596)

paradigmxyz · Jul 18, 2024 · 8c8702b · 8c8702b
1 parent ea47939
commit 8c8702b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/crates/rpc/rpc-builder/src/lib.rs b/crates/rpc/rpc-builder/src/lib.rs
@@ -178,8 +178,10 @@ use tower::Layer;
 use tower_http::cors::CorsLayer;
 
 use crate::{
-    auth::AuthRpcModule, cors::CorsDomainError, error::WsHttpSamePortError,
-    metrics::RpcRequestMetrics,
+    auth::AuthRpcModule,
+    cors::CorsDomainError,
+    error::WsHttpSamePortError,
+    metrics::{RpcRequestMetrics, RpcRequestMetricsService},
 };
 
 // re-export for convenience
@@ -1141,7 +1143,6 @@ pub struct RpcServerConfig<RpcMiddleware = Identity> {
     /// JWT secret for authentication
     jwt_secret: Option<JwtSecret>,
     /// Configurable RPC middleware
-    #[allow(dead_code)]
     rpc_middleware: RpcServiceBuilder<RpcMiddleware>,
 }
 
@@ -1337,8 +1338,9 @@ impl<RpcMiddleware> RpcServerConfig<RpcMiddleware> {
     /// Returns the [`RpcServerHandle`] with the handle to the started servers.
     pub async fn start(self, modules: &TransportRpcModules) -> Result<RpcServerHandle, RpcError>
     where
-        RpcMiddleware: for<'a> Layer<RpcService, Service: RpcServiceT<'a>> + Clone + Send + 'static,
-        <RpcMiddleware as Layer<RpcService>>::Service: Send + std::marker::Sync,
+        RpcMiddleware: Layer<RpcRequestMetricsService<RpcService>> + Clone + Send + 'static,
+        for<'a> <RpcMiddleware as Layer<RpcRequestMetricsService<RpcService>>>::Service:
+            Send + Sync + 'static + RpcServiceT<'a>,
     {
         let mut http_handle = None;
         let mut ws_handle = None;
@@ -1396,7 +1398,7 @@ impl<RpcMiddleware> RpcServerConfig<RpcMiddleware> {
                             .option_layer(Self::maybe_jwt_layer(self.jwt_secret)),
                     )
                     .set_rpc_middleware(
-                        RpcServiceBuilder::new().layer(
+                        self.rpc_middleware.clone().layer(
                             modules
                                 .http
                                 .as_ref()
@@ -1444,7 +1446,8 @@ impl<RpcMiddleware> RpcServerConfig<RpcMiddleware> {
                         .option_layer(Self::maybe_jwt_layer(self.jwt_secret)),
                 )
                 .set_rpc_middleware(
-                    RpcServiceBuilder::new()
+                    self.rpc_middleware
+                        .clone()
                         .layer(modules.ws.as_ref().map(RpcRequestMetrics::ws).unwrap_or_default()),
                 )
                 .build(ws_socket_addr)
@@ -1468,7 +1471,7 @@ impl<RpcMiddleware> RpcServerConfig<RpcMiddleware> {
                         .option_layer(Self::maybe_jwt_layer(self.jwt_secret)),
                 )
                 .set_rpc_middleware(
-                    RpcServiceBuilder::new().layer(
+                    self.rpc_middleware.clone().layer(
                         modules.http.as_ref().map(RpcRequestMetrics::http).unwrap_or_default(),
                     ),
                 )

diff --git a/crates/rpc/rpc-builder/src/metrics.rs b/crates/rpc/rpc-builder/src/metrics.rs
@@ -81,8 +81,8 @@ struct RpcServerMetricsInner {
 /// A [`RpcServiceT`] middleware that captures RPC metrics for the server.
 ///
 /// This is created per connection and captures metrics for each request.
-#[derive(Clone)]
-pub(crate) struct RpcRequestMetricsService<S> {
+#[derive(Clone, Debug)]
+pub struct RpcRequestMetricsService<S> {
     metrics: RpcRequestMetrics,
     inner: S,
 }
@@ -125,7 +125,7 @@ impl<S> Drop for RpcRequestMetricsService<S> {
 
 /// Response future to update the metrics for a single request/response pair.
 #[pin_project::pin_project]
-pub(crate) struct MeteredRequestFuture<F> {
+pub struct MeteredRequestFuture<F> {
     #[pin]
     fut: F,
     /// time when the request started