From c1b0093d9e865a7928e58e203483862859b819b2 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Mon, 27 Sep 2021 17:42:43 +1000
Subject: [PATCH] Add LRU cache for execution blocks

---
 Cargo.lock                                    |  1 +
 beacon_node/execution_layer/Cargo.toml        |  1 +
 beacon_node/execution_layer/src/engine_api.rs |  1 -
 beacon_node/execution_layer/src/lib.rs        | 28 +++++++++++++++++--
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 92bf5da0ee..e189a564a2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2075,6 +2075,7 @@ dependencies = [
  "eth2_ssz_types",
  "futures",
  "hex",
+ "lru",
  "reqwest",
  "sensitive_url",
  "serde",
diff --git a/beacon_node/execution_layer/Cargo.toml b/beacon_node/execution_layer/Cargo.toml
index e9de26e833..a8f472e64c 100644
--- a/beacon_node/execution_layer/Cargo.toml
+++ b/beacon_node/execution_layer/Cargo.toml
@@ -23,3 +23,4 @@ bytes = "1.1.0"
 task_executor = { path = "../../common/task_executor" }
 hex = "0.4.2"
 eth2_ssz_types = { path = "../../consensus/ssz_types"}
+lru = "0.6.0"
diff --git a/beacon_node/execution_layer/src/engine_api.rs b/beacon_node/execution_layer/src/engine_api.rs
index e03a9e6dae..2b0e89e00e 100644
--- a/beacon_node/execution_layer/src/engine_api.rs
+++ b/beacon_node/execution_layer/src/engine_api.rs
@@ -102,7 +102,6 @@ pub enum BlockByNumberQuery<'a> {
 #[serde(rename_all = "camelCase")]
 pub struct ExecutionBlock {
     pub block_hash: Hash256,
-    pub block_number: u64,
     pub parent_hash: Hash256,
     pub total_difficulty: Uint256,
 }
diff --git a/beacon_node/execution_layer/src/lib.rs b/beacon_node/execution_layer/src/lib.rs
index 7b34a31b2e..ad3d103f25 100644
--- a/beacon_node/execution_layer/src/lib.rs
+++ b/beacon_node/execution_layer/src/lib.rs
@@ -1,10 +1,12 @@
 use engine_api::{Error as ApiError, *};
 use engines::{Engine, EngineError, Engines};
+use lru::LruCache;
 use sensitive_url::SensitiveUrl;
 use slog::{crit, Logger};
 use std::future::Future;
 use std::sync::Arc;
 use task_executor::TaskExecutor;
+use tokio::sync::{Mutex, MutexGuard};
 
 pub use engine_api::{http::HttpJsonRpc, ConsensusStatus, ExecutePayloadResponse};
 pub use execute_payload_handle::ExecutePayloadHandle;
@@ -14,6 +16,8 @@ mod engines;
 mod execute_payload_handle;
 pub mod test_utils;
 
+const EXECUTION_BLOCKS_LRU_CACHE_SIZE: usize = 128;
+
 #[derive(Debug)]
 pub enum Error {
     ApiError(ApiError),
@@ -33,6 +37,7 @@ struct Inner {
     engines: Engines<HttpJsonRpc>,
     terminal_total_difficulty: Uint256,
     fee_recipient: Option<Address>,
+    execution_blocks: Mutex<LruCache<Hash256, ExecutionBlock>>,
     executor: TaskExecutor,
     log: Logger,
 }
@@ -66,6 +71,7 @@ impl ExecutionLayer {
             },
             terminal_total_difficulty,
             fee_recipient,
+            execution_blocks: Mutex::new(LruCache::new(EXECUTION_BLOCKS_LRU_CACHE_SIZE)),
             executor,
             log,
         };
@@ -95,6 +101,10 @@ impl ExecutionLayer {
             .ok_or(Error::FeeRecipientUnspecified)
     }
 
+    async fn execution_blocks(&self) -> MutexGuard<'_, LruCache<Hash256, ExecutionBlock>> {
+        self.inner.execution_blocks.lock().await
+    }
+
     fn log(&self) -> &Logger {
         &self.inner.log
     }
@@ -266,13 +276,27 @@ impl ExecutionLayer {
                     .api
                     .get_block_by_number(BlockByNumberQuery::Tag(LATEST_TAG))
                     .await?;
+                self.execution_blocks().await.put(block.parent_hash, block);
 
                 loop {
                     if block.total_difficulty >= self.terminal_total_difficulty() {
                         ttd_exceeding_block = Some(block.block_hash);
 
-                        // TODO(paul): add a LRU cache for these lookups.
-                        block = engine.api.get_block_by_hash(block.parent_hash).await?;
+                        let cached = self
+                            .execution_blocks()
+                            .await
+                            .get(&block.parent_hash)
+                            .copied();
+                        if let Some(cached_block) = cached {
+                            // The block was in the cache, no need to request it from the execution
+                            // engine.
+                            block = cached_block;
+                        } else {
+                            // The block was *not* in the cache, request it from the execution
+                            // engine and cache it for future reference.
+                            block = engine.api.get_block_by_hash(block.parent_hash).await?;
+                            self.execution_blocks().await.put(block.parent_hash, block);
+                        }
                     } else {
                         return Ok::<_, ApiError>(ttd_exceeding_block);
                     }