From 99d405f7db30429d3668dd040b683f8632e7f01b Mon Sep 17 00:00:00 2001
From: Bryan Stitt <bryan@llamanodes.com>
Date: Tue, 22 Aug 2023 10:37:12 -0700
Subject: [PATCH] Revert "remove caching the /status things"

This reverts commit c5827b6f37ff44f50a63c07e6c03023429b97287.
---
 web3_proxy/src/frontend/mod.rs    | 26 +++++++++++++++++++++++++-
 web3_proxy/src/frontend/status.rs | 27 +++++++++++++++++++++++----
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/web3_proxy/src/frontend/mod.rs b/web3_proxy/src/frontend/mod.rs
index 9c1f1453..3032f271 100644
--- a/web3_proxy/src/frontend/mod.rs
+++ b/web3_proxy/src/frontend/mod.rs
@@ -17,12 +17,14 @@ use axum::{
     routing::{get, post},
     BoxError, Extension, Router,
 };
-use http::{header::AUTHORIZATION, Request};
+use http::{header::AUTHORIZATION, Request, StatusCode};
 use hyper::Body;
 use listenfd::ListenFd;
+use moka::future::{Cache, CacheBuilder};
 use std::sync::Arc;
 use std::{iter::once, time::Duration};
 use std::{net::SocketAddr, sync::atomic::Ordering};
+use strum::{EnumCount, EnumIter};
 use tokio::sync::broadcast;
 use tower::timeout::TimeoutLayer;
 use tower::ServiceBuilder;
@@ -31,12 +33,32 @@ use tower_http::{cors::CorsLayer, normalize_path::NormalizePathLayer, trace::Tra
 use tracing::{error_span, info};
 use ulid::Ulid;
 
+/// simple keys for caching responses
+#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, EnumCount, EnumIter)]
+pub enum ResponseCacheKey {
+    BackupsNeeded,
+    Health,
+    Status,
+}
+
+pub type ResponseCache = Cache<ResponseCacheKey, (StatusCode, &'static str, axum::body::Bytes)>;
+
 /// Start the frontend server.
 pub async fn serve(
     app: Arc<Web3ProxyApp>,
     mut shutdown_receiver: broadcast::Receiver<()>,
     shutdown_complete_sender: broadcast::Sender<()>,
 ) -> Web3ProxyResult<()> {
+    // setup caches for whatever the frontend needs
+    // no need for max items since it is limited by the enum key
+    // TODO: latest moka allows for different ttls for different
+    let response_cache_size = ResponseCacheKey::COUNT;
+
+    let response_cache: ResponseCache = CacheBuilder::new(response_cache_size as u64)
+        .name("frontend_response")
+        .time_to_live(Duration::from_secs(1))
+        .build();
+
     // TODO: read config for if fastest/versus should be available publicly. default off
 
     // build our axum Router
@@ -242,6 +264,8 @@ pub async fn serve(
         .layer(CorsLayer::very_permissive())
         // application state
         .layer(Extension(app.clone()))
+        // frontend caches
+        .layer(Extension(Arc::new(response_cache)))
         // request timeout
         .layer(
             ServiceBuilder::new()
diff --git a/web3_proxy/src/frontend/status.rs b/web3_proxy/src/frontend/status.rs
index 2340fac9..6ef5d8c5 100644
--- a/web3_proxy/src/frontend/status.rs
+++ b/web3_proxy/src/frontend/status.rs
@@ -3,6 +3,7 @@
 //! For ease of development, users can currently access these endponts.
 //! They will eventually move to another port.
 
+use super::{ResponseCache, ResponseCacheKey};
 use crate::{
     app::{Web3ProxyApp, APP_USER_AGENT},
     errors::Web3ProxyError,
@@ -21,7 +22,8 @@ use moka::future::Cache;
 use once_cell::sync::Lazy;
 use serde::{ser::SerializeStruct, Serialize};
 use serde_json::json;
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
+use tokio::time::timeout;
 use tracing::trace;
 
 static HEALTH_OK: Lazy<Bytes> = Lazy::new(|| Bytes::from("OK\n"));
@@ -75,8 +77,13 @@ pub async fn debug_request(
 #[debug_handler]
 pub async fn health(
     Extension(app): Extension<Arc<Web3ProxyApp>>,
+    Extension(cache): Extension<Arc<ResponseCache>>,
 ) -> Result<impl IntoResponse, Web3ProxyError> {
-    let (code, content_type, body) = _health(app).await;
+    let (code, content_type, body) = timeout(
+        Duration::from_secs(3),
+        cache.get_with(ResponseCacheKey::Health, async move { _health(app).await }),
+    )
+    .await?;
 
     let x = Response::builder()
         .status(code)
@@ -107,8 +114,15 @@ async fn _health(app: Arc<Web3ProxyApp>) -> (StatusCode, &'static str, Bytes) {
 #[debug_handler]
 pub async fn backups_needed(
     Extension(app): Extension<Arc<Web3ProxyApp>>,
+    Extension(cache): Extension<Arc<ResponseCache>>,
 ) -> Result<impl IntoResponse, Web3ProxyError> {
-    let (code, content_type, body) = _backups_needed(app).await;
+    let (code, content_type, body) = timeout(
+        Duration::from_secs(3),
+        cache.get_with(ResponseCacheKey::BackupsNeeded, async move {
+            _backups_needed(app).await
+        }),
+    )
+    .await?;
 
     let x = Response::builder()
         .status(code)
@@ -151,8 +165,13 @@ async fn _backups_needed(app: Arc<Web3ProxyApp>) -> (StatusCode, &'static str, B
 #[debug_handler]
 pub async fn status(
     Extension(app): Extension<Arc<Web3ProxyApp>>,
+    Extension(cache): Extension<Arc<ResponseCache>>,
 ) -> Result<impl IntoResponse, Web3ProxyError> {
-    let (code, content_type, body) = _status(app).await;
+    let (code, content_type, body) = timeout(
+        Duration::from_secs(3),
+        cache.get_with(ResponseCacheKey::Status, async move { _status(app).await }),
+    )
+    .await?;
 
     let x = Response::builder()
         .status(code)