feat: remove states older than the last finalized checkpoint (#1071)

Author: avilagaston9

lib/lambda_ethereum_consensus/fork_choice/fork_choice.ex

@@ -12,6 +12,7 @@ defmodule LambdaEthereumConsensus.ForkChoice do
   alias LambdaEthereumConsensus.ForkChoice.Head
   alias LambdaEthereumConsensus.P2P.Gossip.OperationsCollector
   alias LambdaEthereumConsensus.Store.Blocks
+  alias LambdaEthereumConsensus.Store.StateDb
   alias LambdaEthereumConsensus.Store.StoreDb
   alias LambdaEthereumConsensus.Validator
   alias Types.Attestation
@@ -72,6 +73,8 @@ defmodule LambdaEthereumConsensus.ForkChoice do
 
     Logger.info("[Fork choice] Adding new block", root: block_root, slot: slot)
 
+    %Store{finalized_checkpoint: last_finalized_checkpoint} = store
+
     result =
       :telemetry.span([:sync, :on_block], %{}, fn ->
         {process_block(block_root, signed_block, store), %{}}
@@ -84,6 +87,15 @@ defmodule LambdaEthereumConsensus.ForkChoice do
 
         Task.async(__MODULE__, :recompute_head, [new_store])
 
+        %Store{finalized_checkpoint: new_finalized_checkpoint} = new_store
+
+        if last_finalized_checkpoint.epoch < new_finalized_checkpoint.epoch do
+          new_finalized_slot =
+            new_finalized_checkpoint.epoch * ChainSpec.get("SLOTS_PER_EPOCH")
+
+          Task.async(StateDb, :prune_states_older_than, [new_finalized_slot])
+        end
+
         GenServer.cast(from, {:block_processed, block_root, true})
         {:noreply, new_store}
 

lib/lambda_ethereum_consensus/store/db.ex

@@ -18,6 +18,12 @@ defmodule LambdaEthereumConsensus.Store.Db do
     Exleveldb.put(ref, key, value)
   end
 
+  @spec delete(binary) :: :ok
+  def delete(key) do
+    ref = GenServer.call(@registered_name, :get_ref)
+    Exleveldb.delete(ref, key)
+  end
+
   @spec get(binary) :: {:ok, binary} | :not_found
   def get(key) do
     ref = GenServer.call(@registered_name, :get_ref)

lib/lambda_ethereum_consensus/store/state_db.ex

@@ -30,6 +30,44 @@ defmodule LambdaEthereumConsensus.Store.StateDb do
     Db.put(slothash_key_block, block_root)
   end
 
+  def prune_states_older_than(slot) do
+    last_finalized_key = slot |> root_by_slot_key()
+
+    with {:ok, it} <- Db.iterate(),
+         {:ok, @stateslot_prefix <> _slot, _value} <-
+           Exleveldb.iterator_move(it, last_finalized_key),
+         {:ok, slots_to_remove} <- get_slots_to_remove(it),
+         :ok <- Exleveldb.iterator_close(it) do
+      slots_to_remove |> Enum.map(&remove_by_slot/1)
+    end
+  end
+
+  defp get_slots_to_remove(slots_to_remove \\ [], iterator) do
+    case Exleveldb.iterator_move(iterator, :prev) do
+      {:ok, @stateslot_prefix <> slot, _root} ->
+        [slot | slots_to_remove] |> get_slots_to_remove(iterator)
+
+      _ ->
+        {:ok, slots_to_remove}
+    end
+  end
+
+  defp remove_by_slot(binary_slot) do
+    slot = :binary.decode_unsigned(binary_slot)
+    key_slot = root_by_slot_key(slot)
+
+    with {:ok, block_root} <- Db.get(key_slot),
+         key_block <- state_key(block_root),
+         {:ok, encoded_state} <- Db.get(key_block) do
+      key_state =
+        encoded_state |> Ssz.from_ssz!(BeaconState) |> Ssz.hash_tree_root!() |> block_key()
+
+      Db.delete(key_slot)
+      Db.delete(key_block)
+      Db.delete(key_state)
+    end
+  end
+
   @spec get_state_by_block_root(Types.root()) ::
           {:ok, BeaconState.t()} | {:error, String.t()} | :not_found
   def get_state_by_block_root(root) do

lib/ssz.ex

@@ -28,6 +28,12 @@ defmodule Ssz do
     |> to_ssz_rs(schema)
   end
 
+  @spec from_ssz!(binary, module) :: struct
+  def from_ssz!(bin, schema) do
+    {:ok, root} = from_ssz(bin, schema)
+    root
+  end
+
   @spec from_ssz(binary, module) :: {:ok, struct} | {:error, String.t()}
   def from_ssz(bin, schema) do
     with {:ok, map} <- from_ssz_rs(bin, schema) do