[json rpc server] add timeout (#22095)

## Description 

add timeouts to JSON RPC server

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] gRPC:
- [ ] JSON-RPC: 
- [ ] GraphQL: 
- [ ] CLI: 
- [ ] Rust SDK:

Author: phoenix-o

crates/sui-json-rpc/src/lib.rs

@@ -1,16 +1,16 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::net::SocketAddr;
-use std::sync::Arc;
-
 use axum::body::Body;
 use axum::http;
 use hyper::Request;
 use jsonrpsee::RpcModule;
 use metrics::Metrics;
 use metrics::MetricsLayer;
 use prometheus::Registry;
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::time::Duration;
 use sui_core::traffic_controller::metrics::TrafficControllerMetrics;
 use sui_core::traffic_controller::TrafficController;
 use sui_types::traffic_control::PolicyConfig;
@@ -163,7 +163,13 @@ impl JsonRpcServerBuilder {
         let (stop_handle, server_handle) = jsonrpsee::server::stop_channel();
         std::mem::forget(server_handle);
 
+        let timeout = std::env::var("JSON_RPC_TIMEOUT")
+            .ok()
+            .and_then(|value| value.parse::<u64>().ok())
+            .unwrap_or(60);
+
         let rpc_middleware = jsonrpsee::server::middleware::rpc::RpcServiceBuilder::new()
+            .layer_fn(move |s| TimeoutLayer::new(s, Duration::from_secs(timeout)))
             .layer_fn(move |s| MetricsLayer::new(s, metrics.clone()))
             .layer_fn(move |s| TrafficControllerService::new(s, traffic_controller.clone()));
         let service_builder = jsonrpsee::server::ServerBuilder::new()
@@ -284,6 +290,7 @@ where
     fn rpc_doc_module() -> Module;
 }
 
+use crate::metrics::TimeoutLayer;
 use jsonrpsee::core::BoxError;
 
 #[derive(Clone)]

crates/sui-json-rpc/src/metrics.rs

@@ -1,15 +1,16 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::collections::HashSet;
-use std::sync::Arc;
-
 use futures::FutureExt;
 use jsonrpsee::server::middleware::rpc::RpcServiceT;
+use jsonrpsee::types::{ErrorCode, ErrorObject, Id};
 use prometheus::{
     register_histogram_vec_with_registry, register_int_counter_vec_with_registry,
     register_int_gauge_vec_with_registry, HistogramVec, IntCounterVec, IntGaugeVec,
 };
+use std::collections::HashSet;
+use std::sync::Arc;
+use std::time::Duration;
 use sui_json_rpc_api::TRANSIENT_ERROR_CODE;
 use sui_json_rpc_api::{CLIENT_SDK_TYPE_HEADER, CLIENT_TARGET_API_VERSION_HEADER};
 use tokio::time::Instant;
@@ -241,3 +242,34 @@ where
         .boxed()
     }
 }
+
+#[derive(Clone)]
+pub struct TimeoutLayer<S>(S, Duration);
+
+impl<S> TimeoutLayer<S> {
+    pub fn new(service: S, duration: Duration) -> Self {
+        Self(service, duration)
+    }
+}
+
+impl<'a, S> RpcServiceT<'a> for TimeoutLayer<S>
+where
+    S: RpcServiceT<'a> + Send + Sync,
+    S::Future: 'a,
+{
+    type Future = futures::future::BoxFuture<'a, jsonrpsee::MethodResponse>;
+
+    fn call(&self, req: jsonrpsee::types::Request<'a>) -> Self::Future {
+        let future = tokio::time::timeout(self.1, self.0.call(req));
+        async move {
+            match future.await {
+                Ok(res) => res,
+                Err(_) => jsonrpsee::MethodResponse::error(
+                    Id::Null,
+                    ErrorObject::from(ErrorCode::InternalError),
+                ),
+            }
+        }
+        .boxed()
+    }
+}