From ec07e4749a67f94887801cc6c33fee7fe941b1b1 Mon Sep 17 00:00:00 2001 From: Benjamin Philip Date: Fri, 7 May 2021 12:13:36 +0530 Subject: [PATCH] Use a node pool for node names (#256) --- lib/livebook/application.ex | 2 + lib/livebook/runtime/node_pool.ex | 105 +++++++++++++++++++++++ lib/livebook/runtime/standalone_init.ex | 4 +- test/livebook/runtime/node_pool_test.exs | 76 ++++++++++++++++ 4 files changed, 185 insertions(+), 2 deletions(-) create mode 100644 lib/livebook/runtime/node_pool.ex create mode 100644 test/livebook/runtime/node_pool_test.exs diff --git a/lib/livebook/application.ex b/lib/livebook/application.ex index c49ef5d26..1490c229f 100644 --- a/lib/livebook/application.ex +++ b/lib/livebook/application.ex @@ -18,6 +18,8 @@ defmodule Livebook.Application do Livebook.SessionSupervisor, # Start the server responsible for associating files with sessions Livebook.Session.FileGuard, + # Start the Node Pool for managing node names + Livebook.Runtime.NodePool, # Start the Endpoint (http/https) LivebookWeb.Endpoint ] diff --git a/lib/livebook/runtime/node_pool.ex b/lib/livebook/runtime/node_pool.ex new file mode 100644 index 000000000..7305f2e01 --- /dev/null +++ b/lib/livebook/runtime/node_pool.ex @@ -0,0 +1,105 @@ +defmodule Livebook.Runtime.NodePool do + use GenServer + + @moduledoc false + + # A pool for reusing child node names. + # + # `free_name` refers to the list of unused names. + # `generated_names` refers to the list of names ever generated. + # + # `buffer_time` refers to time the pool waits before + # adding a name to `pool`, which by default is 1 minute. + + @default_time 60_000 + + # Client interface + + @doc """ + Starts the GenServer from a Supervision tree + + ## Options + + - `:name` - The name the NodePool is locally registered as. By default, it is `Livebook.Runtime.NodePool` + - `:buffer_time` - The time that is spent before a disconnected node's name is added to pool. The default is 1 minute. + """ + def start_link(opts) do + name = opts[:name] || __MODULE__ + buffer_time = opts[:buffer_time] || @default_time + + GenServer.start_link( + __MODULE__, + %{buffer_time: buffer_time}, + name: name + ) + end + + @doc """ + Retuns a node name. + + Generates a new name if pool is empty, or takes one from pool. + """ + def get_name(server \\ __MODULE__, basename) do + GenServer.call(server, {:get_name, basename}) + end + + # Server side code + + @impl GenServer + def init(opts) do + :net_kernel.monitor_nodes(true, node_type: :all) + {:ok, %{buffer_time: opts.buffer_time, generated_names: MapSet.new(), free_names: []}} + end + + @impl GenServer + def handle_call({:get_name, basename}, _, state) do + {name, new_state} = name(state, basename) + {:reply, name, new_state} + end + + @impl GenServer + def handle_info({:nodedown, node, _info}, state) do + _ = Process.send_after(self(), {:add_node, node}, state.buffer_time) + {:noreply, state} + end + + @impl GenServer + def handle_info({:nodeup, _node, _info}, state) do + {:noreply, state} + end + + @impl GenServer + def handle_info({:add_node, node}, state) do + {:noreply, add_node(state, node)} + end + + # Helper functions + + defp name(state, basename) do + if Enum.empty?(state.free_names) do + generate_name(state, basename) + else + get_existing_name(state) + end + end + + defp generate_name(state, basename) do + new_name = :"#{Livebook.Utils.random_short_id()}-#{basename}" + generated_names = MapSet.put(state.generated_names, new_name) + {new_name, %{state | generated_names: generated_names}} + end + + defp get_existing_name(state) do + [name | free_names] = state.free_names + {name, %{state | free_names: free_names}} + end + + defp add_node(state, node) do + if MapSet.member?(state.generated_names, node) do + free_names = [node | state.free_names] + %{state | free_names: free_names} + else + state + end + end +end diff --git a/lib/livebook/runtime/standalone_init.ex b/lib/livebook/runtime/standalone_init.ex index 9a44c9552..af934c0cd 100644 --- a/lib/livebook/runtime/standalone_init.ex +++ b/lib/livebook/runtime/standalone_init.ex @@ -5,15 +5,15 @@ defmodule Livebook.Runtime.StandaloneInit do # a new Elixir system process. It's used by both ElixirStandalone # and MixStandalone runtimes. - alias Livebook.Utils alias Livebook.Utils.Emitter + alias Livebook.Runtime.NodePool @doc """ Returns a random name for a dynamically spawned node. """ @spec child_node_name(atom()) :: atom() def child_node_name(parent) do - :"#{Utils.random_short_id()}-#{parent}" + NodePool.get_name(parent) end @doc """ diff --git a/test/livebook/runtime/node_pool_test.exs b/test/livebook/runtime/node_pool_test.exs new file mode 100644 index 000000000..1a61e41de --- /dev/null +++ b/test/livebook/runtime/node_pool_test.exs @@ -0,0 +1,76 @@ +defmodule Livebook.Runtime.NodePoolTest do + use ExUnit.Case, async: true + + alias Livebook.Runtime.NodePool + + # Tests for Livebook.Runtime.NodePool + # + # Note: + # + # We do not spawn actual nodes as it can be time + # intensive (on low spec machines) and is generally + # complicated. + + describe "start_link" do + test "correctly starts a registered GenServer", config do + start_supervised!({NodePool, name: config.test}) + + # Verify Process is running + assert Process.whereis(config.test) + end + end + + describe "get_name/2" do + test "creates a new node name if pool is empty", config do + start_supervised!({NodePool, name: config.test}) + + # Assert that we get a result and that it is an atom + result = NodePool.get_name(config.test, node()) + assert result + assert is_atom(result) + end + + test "returns an existing name if pool is not empty", config do + start_supervised!({NodePool, name: config.test, buffer_time: 0}) + + name = NodePool.get_name(config.test, node()) + send(config.test, {:nodedown, name, {}}) + + # Since we want the `:add_node` message processed first + # before we call `get_name`, we wait + Process.sleep(1) + + assert NodePool.get_name(config.test, node()) == name + end + + test "removes an existing name when used", config do + start_supervised!({NodePool, name: config.test, buffer_time: 0}) + + name = NodePool.get_name(config.test, node()) + send(config.test, {:nodedown, name, {}}) + + # Since we want the `:add_node` message processed first + # before we call `get_name`, we wait + Process.sleep(1) + + name = NodePool.get_name(config.test, node()) + assert NodePool.get_name(config.test, node()) != name + end + end + + describe "on nodedown" do + test "does not add node name to pool if not in generated_names", config do + start_supervised!({NodePool, name: config.test, buffer_time: 0}) + + # Mock a nodedown + send(config.test, {:nodedown, :some_foo, {}}) + + # Since we want the `:add_node` message processed first + # before we call `get_name`, we wait + Process.sleep(1) + + # Verify that name is not in pool, by calling get_name/2 + assert NodePool.get_name(config.test, node()) != :some_foo + end + end +end