refactor: libp2pPort gossip as callbacks instead of pids (#1111)

Co-authored-by: Tomás Grüner <[email protected]>

Author: avilagaston9

lib/lambda_ethereum_consensus/p2p/gossip/attestation.ex

@@ -7,8 +7,11 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.Attestation do
   alias LambdaEthereumConsensus.Beacon.BeaconChain
   alias LambdaEthereumConsensus.Libp2pPort
   alias LambdaEthereumConsensus.P2P
+  alias LambdaEthereumConsensus.P2P.Gossip.Handler
   alias LambdaEthereumConsensus.StateTransition.Misc
 
+  @behaviour Handler
+
   def start_link(init_arg) do
     GenServer.start_link(__MODULE__, init_arg, name: __MODULE__)
   end
@@ -22,6 +25,11 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.Attestation do
     update_enr()
   end
 
+  @impl true
+  def handle_gossip_message(topic, msg_id, message) do
+    GenServer.cast(__MODULE__, {:gossipsub, {topic, msg_id, message}})
+  end
+
   @spec leave(non_neg_integer()) :: :ok
   def leave(subnet_id) do
     topic = get_topic_name(subnet_id)
@@ -98,7 +106,7 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.Attestation do
     attestations = Map.put(state.attestations, subnet_id, [attestation])
     attnets = Map.put(state.attnets, subnet_id, attestation.data)
     new_state = %{state | attnets: attnets, attestations: attestations}
-    get_topic_name(subnet_id) |> Libp2pPort.subscribe_to_topic()
+    get_topic_name(subnet_id) |> Libp2pPort.subscribe_to_topic(__MODULE__)
     {:reply, :ok, new_state}
   end
 
@@ -113,7 +121,7 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.Attestation do
   end
 
   @impl true
-  def handle_info({:gossipsub, {topic, msg_id, message}}, state) do
+  def handle_cast({:gossipsub, {topic, msg_id, message}}, state) do
     subnet_id = extract_subnet_id(topic)
 
     with {:ok, uncompressed} <- :snappyer.decompress(message),

lib/lambda_ethereum_consensus/p2p/gossip/beacon_block.ex

@@ -5,11 +5,13 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.BeaconBlock do
   alias LambdaEthereumConsensus.Beacon.BeaconChain
   alias LambdaEthereumConsensus.Beacon.PendingBlocks
   alias LambdaEthereumConsensus.Libp2pPort
+  alias LambdaEthereumConsensus.P2P.Gossip.Handler
   alias Types.SignedBeaconBlock
 
   use GenServer
 
   require Logger
+  @behaviour Handler
 
   @type state :: %{topic: String.t(), slot: Types.slot()}
 
@@ -30,6 +32,11 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.BeaconBlock do
     GenServer.cast(__MODULE__, {:slot_transition, slot})
   end
 
+  @impl true
+  def handle_gossip_message(topic, msg_id, message) do
+    GenServer.cast(__MODULE__, {:gossipsub, {topic, msg_id, message}})
+  end
+
   ##########################
   ### GenServer Callbacks
   ##########################
@@ -47,7 +54,7 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.BeaconBlock do
 
   @impl true
   def handle_call(:start, _from, %{topic: topic_name} = state) do
-    Libp2pPort.subscribe_to_topic(topic_name)
+    Libp2pPort.subscribe_to_topic(topic_name, __MODULE__)
     {:reply, :ok, state}
   end
 
@@ -57,7 +64,7 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.BeaconBlock do
   end
 
   @impl true
-  def handle_info({:gossipsub, {_topic, msg_id, message}}, %{slot: slot} = state) do
+  def handle_cast({:gossipsub, {_topic, msg_id, message}}, %{slot: slot} = state) do
     with {:ok, uncompressed} <- :snappyer.decompress(message),
          {:ok, signed_block} <- Ssz.from_ssz(uncompressed, SignedBeaconBlock),
          :ok <- validate(signed_block, slot) do

lib/lambda_ethereum_consensus/p2p/gossip/blob_sidecar.ex

@@ -4,12 +4,15 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.BlobSideCar do
   """
   alias LambdaEthereumConsensus.Beacon.BeaconChain
   alias LambdaEthereumConsensus.Libp2pPort
+  alias LambdaEthereumConsensus.P2P.Gossip.Handler
   alias LambdaEthereumConsensus.Store.BlobDb
 
   use GenServer
 
   require Logger
 
+  @behaviour Handler
+
   @type topics :: [String.t()]
 
   ##########################
@@ -24,6 +27,11 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.BlobSideCar do
     GenServer.call(__MODULE__, :start)
   end
 
+  @impl true
+  def handle_gossip_message(topic, msg_id, message) do
+    GenServer.cast(__MODULE__, {:gossipsub, {topic, msg_id, message}})
+  end
+
   ##########################
   ### GenServer Callbacks
   ##########################
@@ -47,12 +55,12 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.BlobSideCar do
 
   @impl true
   def handle_call(:start, _from, topics) do
-    Enum.each(topics, &Libp2pPort.subscribe_to_topic/1)
+    Enum.each(topics, fn topic -> Libp2pPort.subscribe_to_topic(topic, __MODULE__) end)
     {:reply, :ok, topics}
   end
 
   @impl true
-  def handle_info({:gossipsub, {_topic, msg_id, message}}, topics) do
+  def handle_cast({:gossipsub, {_topic, msg_id, message}}, topics) do
     with {:ok, uncompressed} <- :snappyer.decompress(message),
          {:ok, %Types.BlobSidecar{index: blob_index} = blob} <-
            Ssz.from_ssz(uncompressed, Types.BlobSidecar) do

lib/lambda_ethereum_consensus/p2p/gossip/handler.ex

@@ -0,0 +1,7 @@
+defmodule LambdaEthereumConsensus.P2P.Gossip.Handler do
+  @moduledoc """
+  Gossip handler behaviour
+  """
+
+  @callback handle_gossip_message(binary(), binary(), iodata()) :: :ok | {:error, any}
+end

lib/lambda_ethereum_consensus/p2p/gossip/operations_collector.ex

@@ -6,6 +6,7 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.OperationsCollector do
 
   alias LambdaEthereumConsensus.Beacon.BeaconChain
   alias LambdaEthereumConsensus.Libp2pPort
+  alias LambdaEthereumConsensus.P2P.Gossip.Handler
   alias LambdaEthereumConsensus.StateTransition.Misc
   alias LambdaEthereumConsensus.Utils.BitField
   alias Types.Attestation
@@ -17,6 +18,8 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.OperationsCollector do
 
   require Logger
 
+  @behaviour Handler
+
   @operations [
     :bls_to_execution_change,
     :attester_slashing,
@@ -77,6 +80,11 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.OperationsCollector do
     GenServer.cast(__MODULE__, {:new_block, block.slot, operations})
   end
 
+  @impl true
+  def handle_gossip_message(topic, msg_id, message) do
+    GenServer.cast(__MODULE__, {:gossipsub, {topic, msg_id, message}})
+  end
+
   @impl GenServer
   def init(_init_arg) do
     topics = get_topic_names()
@@ -88,7 +96,7 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.OperationsCollector do
 
   @impl true
   def handle_call(:start, _from, %{topics: topics} = state) do
-    Enum.each(topics, &Libp2pPort.subscribe_to_topic/1)
+    Enum.each(topics, fn topic -> Libp2pPort.subscribe_to_topic(topic, __MODULE__) end)
     {:reply, :ok, state}
   end
 
@@ -108,7 +116,7 @@ defmodule LambdaEthereumConsensus.P2P.Gossip.OperationsCollector do
   end
 
   @impl true
-  def handle_info(
+  def handle_cast(
         {:gossipsub,
          {<<_::binary-size(15)>> <> "beacon_aggregate_and_proof" <> _, _msg_id, message}},
         state

lib/libp2p_port.ex

@@ -187,13 +187,15 @@ defmodule LambdaEthereumConsensus.Libp2pPort do
   Subscribes to the given topic. After this, messages published to the topic
   will be received by `self()`.
   """
-  @spec subscribe_to_topic(GenServer.server(), String.t()) :: :ok | {:error, String.t()}
-  def subscribe_to_topic(pid \\ __MODULE__, topic_name) do
+  @spec subscribe_to_topic(GenServer.server(), String.t(), module()) :: :ok | {:error, String.t()}
+  def subscribe_to_topic(pid \\ __MODULE__, topic_name, module) do
     :telemetry.execute([:port, :message], %{}, %{
       function: "subscribe_to_topic",
       direction: "elixir->"
     })
 
+    GenServer.cast(pid, {:new_subscriptor, topic_name, module})
+
     call_command(pid, {:subscribe, %SubscribeToTopic{name: topic_name}})
   end
 
@@ -279,7 +281,13 @@ defmodule LambdaEthereumConsensus.Libp2pPort do
     |> InitArgs.encode()
     |> then(&send_data(port, &1))
 
-    {:ok, %{port: port, new_peer_handler: new_peer_handler}}
+    {:ok, %{port: port, new_peer_handler: new_peer_handler, subscriptors: %{}}}
+  end
+
+  @impl GenServer
+  def handle_cast({:new_subscriptor, topic, module}, %{subscriptors: subscriptors} = state) do
+    new_subscriptors = Map.put(subscriptors, topic, module)
+    {:noreply, %{state | subscriptors: new_subscriptors}}
   end
 
   @impl GenServer
@@ -316,10 +324,13 @@ defmodule LambdaEthereumConsensus.Libp2pPort do
   ### PRIVATE FUNCTIONS
   ######################
 
-  defp handle_notification(%GossipSub{} = gs, _state) do
+  defp handle_notification(%GossipSub{} = gs, %{subscriptors: subscriptors}) do
     :telemetry.execute([:port, :message], %{}, %{function: "gossipsub", direction: "->elixir"})
-    handler_pid = :erlang.binary_to_term(gs.handler)
-    send(handler_pid, {:gossipsub, {gs.topic, gs.msg_id, gs.message}})
+
+    case Map.fetch(subscriptors, gs.topic) do
+      {:ok, module} -> module.handle_gossip_message(gs.topic, gs.msg_id, gs.message)
+      :error -> Logger.error("[Gossip] Received gossip from unknown topic: #{gs.topic}.")
+    end
   end
 
   defp handle_notification(

test/unit/libp2p_port_test.exs

@@ -4,6 +4,7 @@ defmodule Unit.Libp2pPortTest do
 
   alias LambdaEthereumConsensus.Beacon.BeaconChain
   alias LambdaEthereumConsensus.Libp2pPort
+  alias LambdaEthereumConsensus.P2P.Gossip.Handler
 
   doctest Libp2pPort
 
@@ -86,37 +87,30 @@ defmodule Unit.Libp2pPortTest do
     start_port(:publisher)
     start_port(:gossiper, listen_addr: gossiper_addr)
 
+    # Send the PID in the message, so that we can receive a notification later.
+    message = self() |> :erlang.term_to_binary()
     topic = "/test/gossipping"
-    message = "hello world!"
 
     # Connect the two peers
     %{peer_id: id} = Libp2pPort.get_node_identity(:gossiper)
     :ok = Libp2pPort.add_peer(:publisher, id, gossiper_addr, 999_999_999_999)
 
-    pid = self()
-
-    spawn_link(fn ->
-      # Subscribe to the topic
-      :ok = Libp2pPort.subscribe_to_topic(:gossiper, topic)
-      send(pid, :subscribed)
-
-      # Receive the message
-      assert {^topic, message_id, ^message} = Libp2pPort.receive_gossip()
-
-      Libp2pPort.validate_message(message_id, :accept)
-
-      # Send acknowledgement
-      send(pid, :received)
-    end)
-
-    # Give a head start to the other process
-    assert_receive :subscribed, 100
+    # Subscribe to the topic
+    :ok = Libp2pPort.subscribe_to_topic(:gossiper, topic, __MODULE__)
 
     # Publish message
     :ok = Libp2pPort.publish(:publisher, topic, message)
 
-    # Receive acknowledgement
-    assert_receive :received, 100
+    # Receive the message
+    assert {^topic, message_id, ^message} = Libp2pPort.receive_gossip()
+
+    Libp2pPort.validate_message(message_id, :accept)
+  end
+
+  @behaviour Handler
+  def handle_gossip_message(topic, msg_id, message) do
+    # Decode the PID from the message and send a notification.
+    send(:erlang.binary_to_term(message), {:gossipsub, {topic, msg_id, message}})
   end
 
   defp retry_test(f, retries) do
@@ -137,7 +131,7 @@ defmodule Unit.Libp2pPortTest do
     port = start_port(:some, listen_addr: ["/ip4/127.0.0.1/tcp/48790"])
     topic = "test"
 
-    Libp2pPort.subscribe_to_topic(port, topic)
+    Libp2pPort.subscribe_to_topic(port, topic, __MODULE__)
     Libp2pPort.leave_topic(port, topic)
     Libp2pPort.join_topic(port, topic)
   end