From d80df3640380735406b6c6f0fab5f7326f0dc2ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Gr=C3=BCner?=
 <47506558+MegaRedHand@users.noreply.github.com>
Date: Wed, 17 Jan 2024 09:59:59 -0300
Subject: [PATCH] feat: add blocks cache (#605)

---
 bench/block_processing.exs                    |  5 +-
 lib/lambda_ethereum_consensus/application.ex  |  9 ++-
 .../beacon/beacon_node.ex                     | 10 +--
 .../fork_choice/helpers.ex                    | 14 ++--
 .../incoming_requests/incoming_requests.ex    |  2 +-
 lib/lambda_ethereum_consensus/store/blocks.ex | 74 +++++++++++++++++++
 lib/types/store.ex                            | 11 +--
 7 files changed, 101 insertions(+), 24 deletions(-)
 create mode 100644 lib/lambda_ethereum_consensus/store/blocks.ex

diff --git a/bench/block_processing.exs b/bench/block_processing.exs
index 907090b6d..ca07d61c5 100644
--- a/bench/block_processing.exs
+++ b/bench/block_processing.exs
@@ -2,12 +2,13 @@ alias LambdaEthereumConsensus.ForkChoice
 alias LambdaEthereumConsensus.ForkChoice.Handlers
 alias LambdaEthereumConsensus.ForkChoice.Helpers
 alias LambdaEthereumConsensus.StateTransition.Cache
+alias LambdaEthereumConsensus.Store
 alias LambdaEthereumConsensus.Store.BlockStore
-alias LambdaEthereumConsensus.Store.Db
 alias LambdaEthereumConsensus.Store.StateStore
 alias Types.{BeaconState, SignedBeaconBlock}
 
-{:ok, _} = Db.start_link(nil)
+{:ok, _} = Store.Db.start_link(nil)
+{:ok, _} = Store.Blocks.start_link(nil)
 Cache.initialize_cache()
 
 # NOTE: this slot must be at the beginning of an epoch (i.e. a multiple of 32)
diff --git a/lib/lambda_ethereum_consensus/application.ex b/lib/lambda_ethereum_consensus/application.ex
index 7bdf4fca2..342a89eb5 100644
--- a/lib/lambda_ethereum_consensus/application.ex
+++ b/lib/lambda_ethereum_consensus/application.ex
@@ -14,11 +14,12 @@ defmodule LambdaEthereumConsensus.Application do
       ]
 
     children = [
-      {LambdaEthereumConsensus.Telemetry, []},
-      {LambdaEthereumConsensus.Store.Db, []},
+      LambdaEthereumConsensus.Telemetry,
+      LambdaEthereumConsensus.Store.Db,
+      LambdaEthereumConsensus.Store.Blocks,
       {LambdaEthereumConsensus.Beacon.BeaconNode, [checkpoint_sync]},
-      {LambdaEthereumConsensus.P2P.Metadata, []},
-      {BeaconApi.Endpoint, []}
+      LambdaEthereumConsensus.P2P.Metadata,
+      BeaconApi.Endpoint
     ]
 
     # See https://hexdocs.pm/elixir/Supervisor.html
diff --git a/lib/lambda_ethereum_consensus/beacon/beacon_node.ex b/lib/lambda_ethereum_consensus/beacon/beacon_node.ex
index 64e4b55d4..a9aaa8d9e 100644
--- a/lib/lambda_ethereum_consensus/beacon/beacon_node.ex
+++ b/lib/lambda_ethereum_consensus/beacon/beacon_node.ex
@@ -70,11 +70,11 @@ defmodule LambdaEthereumConsensus.Beacon.BeaconNode do
       {LambdaEthereumConsensus.Beacon.BeaconChain, {anchor_state, time}},
       {LambdaEthereumConsensus.ForkChoice, {anchor_state, anchor_block, time}},
       {LambdaEthereumConsensus.Libp2pPort, libp2p_args},
-      {LambdaEthereumConsensus.P2P.Peerbook, []},
-      {LambdaEthereumConsensus.P2P.IncomingRequests, []},
-      {LambdaEthereumConsensus.Beacon.PendingBlocks, []},
-      {LambdaEthereumConsensus.Beacon.SyncBlocks, []},
-      {LambdaEthereumConsensus.P2P.GossipSub, []}
+      LambdaEthereumConsensus.P2P.Peerbook,
+      LambdaEthereumConsensus.P2P.IncomingRequests,
+      LambdaEthereumConsensus.Beacon.PendingBlocks,
+      LambdaEthereumConsensus.Beacon.SyncBlocks,
+      LambdaEthereumConsensus.P2P.GossipSub
     ]
 
     Supervisor.init(children, strategy: :one_for_all)
diff --git a/lib/lambda_ethereum_consensus/fork_choice/helpers.ex b/lib/lambda_ethereum_consensus/fork_choice/helpers.ex
index 15436495b..a7c8e4e79 100644
--- a/lib/lambda_ethereum_consensus/fork_choice/helpers.ex
+++ b/lib/lambda_ethereum_consensus/fork_choice/helpers.ex
@@ -134,23 +134,23 @@ defmodule LambdaEthereumConsensus.ForkChoice.Helpers do
   # whose leaf state's justified/finalized info agrees with that in ``store``.
   defp get_filtered_block_tree(%Store{} = store) do
     base = store.justified_checkpoint.root
-    {_, blocks} = filter_block_tree(store, base, %{})
+    block = Store.get_block!(store, base)
+    {_, blocks} = filter_block_tree(store, base, block, %{})
     blocks
   end
 
-  defp filter_block_tree(%Store{} = store, block_root, blocks) do
-    block = Store.get_block!(store, block_root)
-
+  defp filter_block_tree(%Store{} = store, block_root, block, blocks) do
     # TODO: this is highly inefficient. We should move to `ForkChoice.Tree` ASAP
     children =
       Store.get_blocks(store)
-      |> Stream.filter(fn {_, block} -> block.parent_root == block_root end)
-      |> Enum.map(fn {root, _} -> root end)
+      |> Enum.filter(fn {_, block} -> block.parent_root == block_root end)
 
     # If any children branches contain expected finalized/justified checkpoints,
     # add to filtered block-tree and signal viability to parent.
     {filter_block_tree_result, new_blocks} =
-      Enum.map_reduce(children, blocks, fn root, acc -> filter_block_tree(store, root, acc) end)
+      Enum.map_reduce(children, blocks, fn {root, block}, acc ->
+        filter_block_tree(store, root, block, acc)
+      end)
 
     cond do
       Enum.any?(filter_block_tree_result) ->
diff --git a/lib/lambda_ethereum_consensus/p2p/incoming_requests/incoming_requests.ex b/lib/lambda_ethereum_consensus/p2p/incoming_requests/incoming_requests.ex
index 9ed845a97..a8cc03308 100644
--- a/lib/lambda_ethereum_consensus/p2p/incoming_requests/incoming_requests.ex
+++ b/lib/lambda_ethereum_consensus/p2p/incoming_requests/incoming_requests.ex
@@ -14,7 +14,7 @@ defmodule LambdaEthereumConsensus.P2P.IncomingRequests do
   def init(_opts) do
     children = [
       {Task.Supervisor, name: IncomingRequests.Handler},
-      {IncomingRequests.Receiver, []}
+      IncomingRequests.Receiver
     ]
 
     Supervisor.init(children, strategy: :one_for_one)
diff --git a/lib/lambda_ethereum_consensus/store/blocks.ex b/lib/lambda_ethereum_consensus/store/blocks.ex
new file mode 100644
index 000000000..ec13048a0
--- /dev/null
+++ b/lib/lambda_ethereum_consensus/store/blocks.ex
@@ -0,0 +1,74 @@
+defmodule LambdaEthereumConsensus.Store.Blocks do
+  @moduledoc false
+  alias LambdaEthereumConsensus.Store.BlockStore
+
+  use GenServer
+
+  @ets_block_by_hash __MODULE__
+
+  ##########################
+  ### Public API
+  ##########################
+
+  def start_link(_opts) do
+    GenServer.start_link(__MODULE__, [], name: __MODULE__)
+  end
+
+  def store_block(block_root, signed_block) do
+    cache_block(block_root, signed_block)
+    GenServer.cast(__MODULE__, {:store_block, block_root, signed_block})
+  end
+
+  def get_block(block_root), do: lookup(block_root)
+
+  @spec clear() :: any()
+  def clear, do: :ets.delete_all_objects(@ets_block_by_hash)
+
+  ##########################
+  ### GenServer Callbacks
+  ##########################
+
+  @impl GenServer
+  def init(_) do
+    :ets.new(@ets_block_by_hash, [:set, :public, :named_table])
+    {:ok, nil}
+  end
+
+  @impl GenServer
+  def handle_cast({:store_block, block_root, signed_block}, state) do
+    BlockStore.store_block(signed_block, block_root)
+    # TODO: remove old blocks from cache
+    {:noreply, state}
+  end
+
+  ##########################
+  ### Private Functions
+  ##########################
+
+  defp lookup(block_root) do
+    case :ets.lookup_element(@ets_block_by_hash, block_root, 2, nil) do
+      nil -> cache_miss(block_root)
+      block -> block
+    end
+  end
+
+  defp cache_miss(block_root) do
+    case fetch_block(block_root) do
+      nil -> nil
+      block -> cache_block(block_root, block)
+    end
+  end
+
+  defp fetch_block(block_root) do
+    case BlockStore.get_block(block_root) do
+      {:ok, signed_block} -> signed_block
+      :not_found -> nil
+      # TODO: handle this somehow?
+      {:error, error} -> raise "database error #{inspect(error)}"
+    end
+  end
+
+  defp cache_block(block_root, signed_block) do
+    :ets.insert_new(@ets_block_by_hash, {block_root, signed_block})
+  end
+end
diff --git a/lib/types/store.ex b/lib/types/store.ex
index 008b0d704..a3cc0162c 100644
--- a/lib/types/store.ex
+++ b/lib/types/store.ex
@@ -35,6 +35,7 @@ defmodule Types.Store do
         }
 
   alias LambdaEthereumConsensus.StateTransition.Misc
+  alias LambdaEthereumConsensus.Store.Blocks
   alias LambdaEthereumConsensus.Store.BlockStore
   alias LambdaEthereumConsensus.Store.StateStore
   alias Types.BeaconState
@@ -104,16 +105,16 @@ defmodule Types.Store do
 
   @spec get_block(t(), Types.root()) :: Types.BeaconBlock.t() | nil
   def get_block(%__MODULE__{}, block_root) do
-    case BlockStore.get_block(block_root) do
-      {:ok, signed_block} -> signed_block.message
-      _ -> nil
+    case Blocks.get_block(block_root) do
+      nil -> nil
+      signed_block -> signed_block.message
     end
   end
 
   @spec get_block!(t(), Types.root()) :: Types.BeaconBlock.t()
   def get_block!(store, block_root) do
     case get_block(store, block_root) do
-      nil -> raise "Block not found: #{block_root}"
+      nil -> raise "Block not found: 0x#{Base.encode16(block_root)}"
       v -> v
     end
   end
@@ -131,7 +132,7 @@ defmodule Types.Store do
 
   @spec store_block(t(), Types.root(), SignedBeaconBlock.t()) :: t()
   def store_block(%__MODULE__{} = store, block_root, %SignedBeaconBlock{} = signed_block) do
-    BlockStore.store_block(signed_block, block_root)
+    Blocks.store_block(block_root, signed_block)
     store
   end
 end