From 75ca4a37347a6723b668168f0a8f52b645384991 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 20 Aug 2024 09:12:39 +0200
Subject: [PATCH] Report connected nodes on nodeup/nodedown events (#2759)

---
 .../runtime/erl_dist/runtime_server.ex        | 31 ++++++++-----------
 1 file changed, 13 insertions(+), 18 deletions(-)

diff --git a/lib/livebook/runtime/erl_dist/runtime_server.ex b/lib/livebook/runtime/erl_dist/runtime_server.ex
index 80f2ac203..80855eda9 100644
--- a/lib/livebook/runtime/erl_dist/runtime_server.ex
+++ b/lib/livebook/runtime/erl_dist/runtime_server.ex
@@ -336,6 +336,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   def init(opts) do
     Process.send_after(self(), :check_owner, @await_owner_timeout)
 
+    :net_kernel.monitor_nodes(true, node_type: :all)
+
     schedule_memory_usage_report()
 
     {:ok, evaluator_supervisor} = ErlDist.EvaluatorSupervisor.start_link()
@@ -361,7 +363,6 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
        smart_cell_definitions_module:
          Keyword.get(opts, :smart_cell_definitions_module, Kino.SmartCell),
        extra_smart_cell_definitions: Keyword.get(opts, :extra_smart_cell_definitions, []),
-       connected_nodes: [],
        memory_timer_ref: nil,
        last_evaluator: nil,
        base_env_path:
@@ -398,7 +399,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   def handle_info({:evaluation_finished, locator}, state) do
     {:noreply,
      state
-     |> report_environment()
+     |> report_smart_cell_definitions()
      |> report_transient_state()
      |> scan_binding_after_evaluation(locator)}
   end
@@ -422,6 +423,11 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {:noreply, state}
   end
 
+  def handle_info({message, _node, _info}, state) when message in [:nodeup, :nodedown] do
+    report_connected_nodes(state)
+    {:noreply, state}
+  end
+
   def handle_info(_message, state), do: {:noreply, state}
 
   defp handle_down_evaluator(state, {:DOWN, _, :process, pid, reason}) do
@@ -470,7 +476,9 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     Process.monitor(owner)
 
     state = %{state | owner: owner, runtime_broadcast_to: opts[:runtime_broadcast_to]}
-    state = report_environment(state)
+
+    state = report_smart_cell_definitions(state)
+    report_connected_nodes(state)
     report_memory_usage(state)
 
     {:ok, smart_cell_supervisor} = DynamicSupervisor.start_link(strategy: :one_for_one)
@@ -696,7 +704,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
 
   def handle_cast({:disconnect_node, node}, state) do
     Node.disconnect(node)
-    {:noreply, report_connected_nodes(state)}
+    {:noreply, state}
   end
 
   @impl true
@@ -805,12 +813,6 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     send(state.owner, {:runtime_memory_usage, Evaluator.memory()})
   end
 
-  defp report_environment(state) do
-    state
-    |> report_smart_cell_definitions()
-    |> report_connected_nodes()
-  end
-
   defp report_smart_cell_definitions(state) do
     smart_cell_definitions = get_smart_cell_definitions(state.smart_cell_definitions_module)
 
@@ -834,14 +836,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   defp report_connected_nodes(state) do
     owner_node = node(state.owner)
     nodes = Node.list(:connected) |> List.delete(owner_node) |> Enum.sort()
-
-    if nodes == state.connected_nodes do
-      state
-    else
-      send(state.owner, {:runtime_connected_nodes, nodes})
-
-      %{state | connected_nodes: nodes}
-    end
+    send(state.owner, {:runtime_connected_nodes, nodes})
   end
 
   defp get_smart_cell_definitions(module) do