Merge pull request #4237 from ProvableHQ/feat/rest_enhancements

[Perf] REST enhancements

Author: vicsn

node/rest/Cargo.toml

@@ -61,6 +61,9 @@ workspace = true
 features = [ "parking_lot" ]
 optional = true
 
+[dependencies.lru]
+workspace = true
+
 [dependencies.once_cell]
 workspace = true
 

node/rest/src/lib.rs

@@ -51,15 +51,17 @@ use axum::{
 use axum_extra::response::ErasedJson;
 #[cfg(feature = "locktick")]
 use locktick::parking_lot::Mutex;
+use lru::LruCache;
 #[cfg(not(feature = "locktick"))]
 use parking_lot::Mutex;
-use std::{net::SocketAddr, sync::Arc};
+use std::{net::SocketAddr, num::NonZeroUsize, sync::Arc, time::Duration};
 use tokio::{net::TcpListener, sync::Semaphore, task::JoinHandle};
 use tower_governor::{GovernorLayer, governor::GovernorConfigBuilder};
 use tower_http::{
     cors::{Any, CorsLayer},
     trace::TraceLayer,
 };
+use tracing::Span;
 
 /// The default port used for the REST API
 pub const DEFAULT_REST_PORT: u16 = 3030;
@@ -68,6 +70,9 @@ pub const DEFAULT_REST_PORT: u16 = 3030;
 pub const API_VERSION_V1: &str = "v1";
 pub const API_VERSION_V2: &str = "v2";
 
+/// The capacity of the LRU holding recently requested blocks.
+const BLOCK_CACHE_SIZE: usize = 128;
+
 /// A REST API server for the ledger.
 #[derive(Clone)]
 pub struct Rest<N: Network, C: ConsensusStorage<N>, R: Routing<N>> {
@@ -89,6 +94,8 @@ pub struct Rest<N: Network, C: ConsensusStorage<N>, R: Routing<N>> {
     num_verifying_executions: Arc<Semaphore>,
     /// The number of ongoing solution verifications via REST.
     num_verifying_solutions: Arc<Semaphore>,
+    /// A cache containing recently requested blocks.
+    block_cache: Arc<Mutex<LruCache<N::BlockHash, ErasedJson>>>,
 }
 
 impl<N: Network, C: 'static + ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
@@ -113,6 +120,7 @@ impl<N: Network, C: 'static + ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R>
             num_verifying_deploys: Arc::new(Semaphore::new(VM::<N, C>::MAX_PARALLEL_DEPLOY_VERIFICATIONS)),
             num_verifying_executions: Arc::new(Semaphore::new(VM::<N, C>::MAX_PARALLEL_EXECUTE_VERIFICATIONS)),
             num_verifying_solutions: Arc::new(Semaphore::new(N::MAX_SOLUTIONS)),
+            block_cache: Arc::new(Mutex::new(LruCache::new(NonZeroUsize::new(BLOCK_CACHE_SIZE).unwrap()))),
         };
         // Spawn the server.
         server.spawn_server(rest_ip, rest_rps).await?;
@@ -263,20 +271,41 @@ impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
         #[cfg(feature = "history-staking-rewards")]
         let routes = routes.route("/staking/rewards/{address}/{height}", get(Self::get_staking_reward));
 
+        let trace_layer = TraceLayer::new_for_http()
+            .make_span_with(|request: &Request<_>| {
+                let addr = request
+                    .extensions()
+                    .get::<ConnectInfo<SocketAddr>>()
+                    .map(|ConnectInfo(addr)| addr.to_string())
+                    .unwrap_or_else(|| "unknown".to_string());
+
+                // Create a span that includes method, path, and our extracted IP
+                tracing::info_span!(
+                    "REST",
+                    method = %request.method(),
+                    uri = %request.uri().path(),
+                    addr = %addr,
+                )
+            })
+            .on_request(|_request: &Request<_>, _span: &Span| {
+                info!("Received a request");
+            })
+            .on_response(|_response: &Response<_>, latency: Duration, _span: &Span| {
+                info!("Finished request in {:?}", latency);
+            });
+
         routes
             // Pass in `Rest` to make things convenient.
             .with_state(self.clone())
-            // Enable tower-http tracing.
-            .layer(TraceLayer::new_for_http())
-            // Custom logging.
-            .layer(middleware::map_request(log_middleware))
-            // Enable CORS.
-            .layer(cors)
             // Cap the request body size at 512KiB.
             .layer(DefaultBodyLimit::max(512 * 1024))
             .layer(GovernorLayer {
                 config: governor_config.into(),
             })
+            // Enable CORS.
+            .layer(cors)
+            // Enable tower-http tracing.
+            .layer(trace_layer)
     }
 
     async fn spawn_server(&mut self, rest_ip: SocketAddr, rest_rps: u32) -> Result<()> {
@@ -314,12 +343,6 @@ impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
     }
 }
 
-/// Creates a log message for every HTTP request.
-async fn log_middleware(ConnectInfo(addr): ConnectInfo<SocketAddr>, request: Request<Body>) -> Request<Body> {
-    info!("Received '{} {}' from '{addr}'", request.method(), request.uri());
-    request
-}
-
 /// Converts errors to the old style for the v1 API.
 /// The error code will always be 500 and the content a simple string.
 async fn v1_error_middleware(response: Response) -> Response {

node/rest/src/routes.rs

@@ -146,7 +146,10 @@ impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
 
     /// GET /<network>/block/latest
     pub(crate) async fn get_block_latest(State(rest): State<Self>) -> ErasedJson {
-        ErasedJson::pretty(rest.ledger.latest_block())
+        let block = rest.ledger.latest_block();
+        let hash = block.hash();
+        // When present, this is 3x faster than serializing the block from the ledger.
+        rest.block_cache.lock().get_or_insert(hash, || ErasedJson::pretty(block)).clone()
     }
 
     /// GET /<network>/block/{height}
@@ -157,23 +160,38 @@ impl<N: Network, C: ConsensusStorage<N>, R: Routing<N>> Rest<N, C, R> {
     ) -> Result<ErasedJson, RestError> {
         // Manually parse the height or the height of the hash, axum doesn't support different types
         // for the same path param.
-        let id_name;
-        let block = if let Ok(height) = height_or_hash.parse::<u32>() {
-            id_name = "hash";
-            rest.ledger.try_get_block(height).with_context(|| "Failed to get block by height")?
+        let hash = if let Ok(height) = height_or_hash.parse::<u32>() {
+            rest.ledger.get_hash(height).with_context(|| "Failed to get a block's hash")?
         } else if let Ok(hash) = height_or_hash.parse::<N::BlockHash>() {
-            id_name = "height";
-            rest.ledger.try_get_block_by_hash(&hash).with_context(|| "Failed to get block by hash")?
+            hash
         } else {
-            return Err(RestError::bad_request(anyhow!(
-                "invalid input, it is neither a block height nor a block hash"
-            )));
+            return Err(RestError::bad_request(anyhow!("invalid input: neither a block height nor a block hash")));
         };
 
-        match block {
-            Some(block) => Ok(ErasedJson::pretty(block)),
-            None => Err(RestError::not_found(anyhow!("No block with {id_name} {height_or_hash} found"))),
+        // Attempt to find a serialized block in the cache.
+        if let Some(json_block) = rest.block_cache.lock().get(&hash) {
+            return Ok(json_block.clone());
         }
+
+        // Retrieve the block from the database.
+        let json_block = match tokio::task::spawn_blocking(move || match rest.ledger.try_get_block_by_hash(&hash) {
+            Ok(Some(block)) => Some(ErasedJson::pretty(block)),
+            Ok(None) => None,
+            Err(e) => {
+                error!("Couldn't find a block: {e}");
+                None
+            }
+        })
+        .await
+        {
+            Ok(Some(block)) => Ok(block),
+            Ok(None) => Err(RestError::not_found(anyhow!("Couldn't find block {height_or_hash}"))),
+            Err(e) => Err(RestError::internal_server_error(anyhow!("tokio error: {e}"))),
+        }?;
+
+        rest.block_cache.lock().put(hash, json_block.clone());
+
+        Ok(json_block)
     }
 
     /// GET /<network>/blocks?start={start_height}&end={end_height}