Created the new ValidatorPool and moved everything to work with it instead of Setup

Author: rodrigo-o

config/runtime.exs

@@ -171,7 +171,7 @@ if keystore_pass_dir != nil and not File.dir?(keystore_pass_dir) do
   System.halt(2)
 end
 
-config :lambda_ethereum_consensus, LambdaEthereumConsensus.Validator.Setup,
+config :lambda_ethereum_consensus, LambdaEthereumConsensus.ValidatorPool,
   keystore_dir: keystore_dir,
   keystore_pass_dir: keystore_pass_dir
 

lib/lambda_ethereum_consensus/beacon/beacon_node.ex

@@ -8,7 +8,7 @@ defmodule LambdaEthereumConsensus.Beacon.BeaconNode do
   alias LambdaEthereumConsensus.ForkChoice
   alias LambdaEthereumConsensus.StateTransition.Cache
   alias LambdaEthereumConsensus.Store.BlockStates
-  alias LambdaEthereumConsensus.Validator
+  alias LambdaEthereumConsensus.ValidatorPool
   alias Types.BeaconState
 
   def start_link(opts) do
@@ -30,7 +30,7 @@ defmodule LambdaEthereumConsensus.Beacon.BeaconNode do
 
     init_execution_chain(deposit_tree_snapshot, store.head_root)
 
-    validators = Validator.Setup.init(store.head_slot, store.head_root)
+    validators = ValidatorPool.init(store.head_slot, store.head_root)
 
     libp2p_args = [genesis_time: store.genesis_time, validators: validators] ++ get_libp2p_args()
 

lib/lambda_ethereum_consensus/validator/setup.ex

@@ -0,0 +1,150 @@
+defmodule LambdaEthereumConsensus.ValidatorPool do
+  @moduledoc """
+  Module that holds the pool of validators and their states,
+  it also manages the validator's duties as bitmaps to
+  simplify the delegation of work.
+  """
+
+  defstruct epoch: nil, slot: nil, head_root: nil, validators: %{uninitialized: []}
+
+  require Logger
+
+  alias LambdaEthereumConsensus.StateTransition.Misc
+  alias LambdaEthereumConsensus.Validator
+
+  @type validators :: %{atom() => list(Validator.state())}
+  @type t :: %__MODULE__{
+          epoch: Types.epoch(),
+          slot: Types.slot(),
+          head_root: Types.root(),
+          validators: validators()
+        }
+
+  @doc """
+  Initiate the pool of validators, given the slot and head root.
+  """
+  @spec init(Types.slot(), Types.root()) :: t()
+  def init(slot, head_root) do
+    config = Application.get_env(:lambda_ethereum_consensus, __MODULE__, [])
+    keystore_dir = Keyword.get(config, :keystore_dir)
+    keystore_pass_dir = Keyword.get(config, :keystore_pass_dir)
+
+    setup_validators(slot, head_root, keystore_dir, keystore_pass_dir)
+  end
+
+  defp setup_validators(_s, _r, keystore_dir, keystore_pass_dir)
+       when is_nil(keystore_dir) or is_nil(keystore_pass_dir) do
+    Logger.warning(
+      "[Validator] No keystore_dir or keystore_pass_dir provided. Validators won't start."
+    )
+
+    %__MODULE__{}
+  end
+
+  defp setup_validators(slot, head_root, keystore_dir, keystore_pass_dir) do
+    validator_keys = decode_validator_keys(keystore_dir, keystore_pass_dir)
+
+    validators = Enum.map(validator_keys, &Validator.new({slot, head_root, &1}))
+
+    Logger.info("[Validator] Initialized #{Enum.count(validators)} validators")
+
+    %__MODULE__{
+      epoch: Misc.compute_epoch_at_slot(slot),
+      slot: slot,
+      head_root: head_root,
+      validators: %{uninitialized: validators}
+    }
+  end
+
+  @doc """
+  Notify all validators of a new head.
+  """
+  @spec notify_head(t(), Types.slot(), Types.root()) :: t()
+  def notify_head(%{validators: %{uninitialized: validators}} = pool, slot, head_root) do
+    uninitialized_validators =
+      maybe_debug_notify(
+        fn ->
+          Enum.map(validators, &Validator.handle_new_head(slot, head_root, &1))
+        end,
+        {:new_head, slot, head_root}
+      )
+
+    %{pool | validators: %{uninitialized: uninitialized_validators}}
+  end
+
+  @doc """
+  Notify all validators of a new block.
+  """
+  @spec notify_tick(t(), tuple()) :: t()
+  def notify_tick(%{validators: %{uninitialized: validators}} = pool, slot_data) do
+    uninitialized_validators =
+      maybe_debug_notify(
+        fn ->
+          Enum.map(validators, &Validator.handle_tick(slot_data, &1))
+        end,
+        {:on_tick, slot_data}
+      )
+
+    %{pool | validators: %{uninitialized: uninitialized_validators}}
+  end
+
+  defp maybe_debug_notify(fun, data) do
+    if Application.get_env(:logger, :level) == :debug do
+      Logger.debug("[Validator] Notifying all Validators with message: #{inspect(data)}")
+
+      start_time = System.monotonic_time(:millisecond)
+      result = fun.()
+      end_time = System.monotonic_time(:millisecond)
+
+      Logger.debug(
+        "[Validator] #{inspect(data)} notified to all Validators after #{end_time - start_time} ms"
+      )
+
+      result
+    else
+      fun.()
+    end
+  end
+
+  @doc """
+    Get validator keys from the keystore directory.
+    This function expects two files for each validator:
+      - <keystore_dir>/<public_key>.json
+      - <keystore_pass_dir>/<public_key>.txt
+  """
+  @spec decode_validator_keys(binary(), binary()) ::
+          list({Bls.pubkey(), Bls.privkey()})
+  def decode_validator_keys(keystore_dir, keystore_pass_dir)
+      when is_binary(keystore_dir) and is_binary(keystore_pass_dir) do
+    keystore_dir
+    |> File.ls!()
+    |> Enum.flat_map(&paths_from_filename(keystore_dir, keystore_pass_dir, &1, Path.extname(&1)))
+    |> Enum.flat_map(&decode_key/1)
+  end
+
+  defp paths_from_filename(keystore_dir, keystore_pass_dir, filename, ".json") do
+    basename = Path.basename(filename, ".json")
+
+    keystore_file = Path.join(keystore_dir, "#{basename}.json")
+    keystore_pass_file = Path.join(keystore_pass_dir, "#{basename}.txt")
+
+    [{keystore_file, keystore_pass_file}]
+  end
+
+  defp paths_from_filename(_keystore_dir, _keystore_pass_dir, basename, _ext) do
+    Logger.warning("[Validator] Skipping file: #{basename}. Not a json keystore file.")
+    []
+  end
+
+  defp decode_key({keystore_file, keystore_pass_file}) do
+    # TODO: remove `try` and handle errors properly
+    [Keystore.decode_from_files!(keystore_file, keystore_pass_file)]
+  rescue
+    error ->
+      Logger.error(
+        "[Validator] Failed to decode keystore file: #{keystore_file}. Pass file: #{keystore_pass_file} Error: #{inspect(error)}"
+      )
+
+      []
+  end
+end

lib/lambda_ethereum_consensus/validator/validator_pool.ex

@@ -23,7 +23,7 @@ defmodule LambdaEthereumConsensus.Libp2pPort do
   alias LambdaEthereumConsensus.P2p.Requests
   alias LambdaEthereumConsensus.StateTransition.Misc
   alias LambdaEthereumConsensus.Utils.BitVector
-  alias LambdaEthereumConsensus.Validator
+  alias LambdaEthereumConsensus.ValidatorPool
   alias Libp2pProto.AddPeer
   alias Libp2pProto.Command
   alias Libp2pProto.Enr
@@ -507,7 +507,7 @@ defmodule LambdaEthereumConsensus.Libp2pPort do
   @impl GenServer
   def handle_info({:new_head, slot, head_root}, %{validators: validators} = state) do
     updated_validators =
-      Validator.Setup.notify_validators(validators, {:new_head, slot, head_root})
+      ValidatorPool.notify_head(validators, slot, head_root)
 
     {:noreply, %{state | validators: updated_validators}}
   end
@@ -748,7 +748,7 @@ defmodule LambdaEthereumConsensus.Libp2pPort do
         state
       else
         updated_validators =
-          Validator.Setup.notify_validators(state.validators, {:on_tick, new_slot_data})
+          ValidatorPool.notify_tick(state.validators, new_slot_data)
 
         %{state | slot_data: new_slot_data, validators: updated_validators}
       end

lib/libp2p_port.ex

@@ -17,7 +17,8 @@ defmodule Unit.Libp2pPortTest do
   end
 
   defp start_port(name \\ Libp2pPort, init_args \\ []) do
-    start_link_supervised!({Libp2pPort, [opts: [name: name], genesis_time: :os.system_time(:second)] ++ init_args},
+    start_link_supervised!(
+      {Libp2pPort, [opts: [name: name], genesis_time: :os.system_time(:second)] ++ init_args},
       id: name
     )
   end