From 31c119a633d7bd887819ef33131ba9359b90dca3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonatan=20K=C5=82osko?= Date: Thu, 10 Nov 2022 16:37:57 +0100 Subject: [PATCH] Persist modules bytecode to disk (#1521) --- lib/livebook/runtime/erl_dist/node_manager.ex | 41 +++++++++++++++-- .../runtime/erl_dist/runtime_server.ex | 14 ++++-- lib/livebook/runtime/evaluator.ex | 45 ++++++++++++++++--- lib/livebook/runtime/evaluator/tracer.ex | 3 +- mix.exs | 1 + test/livebook/runtime/evaluator_test.exs | 37 +++++++++++++-- 6 files changed, 123 insertions(+), 18 deletions(-) diff --git a/lib/livebook/runtime/erl_dist/node_manager.ex b/lib/livebook/runtime/erl_dist/node_manager.ex index e8bdc1c1f..4501d93e5 100644 --- a/lib/livebook/runtime/erl_dist/node_manager.ex +++ b/lib/livebook/runtime/erl_dist/node_manager.ex @@ -93,7 +93,7 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do Process.flag(:trap_exit, true) - {:ok, server_supevisor} = DynamicSupervisor.start_link(strategy: :one_for_one) + {:ok, server_supervisor} = DynamicSupervisor.start_link(strategy: :one_for_one) # Register our own standard error IO device that proxies # to sender's group leader. @@ -108,16 +108,24 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do initial_ignore_module_conflict = Code.compiler_options()[:ignore_module_conflict] Code.compiler_options(ignore_module_conflict: true) + tmp_dir = make_tmp_dir() + + if ebin_path = ebin_path(tmp_dir) do + File.mkdir_p!(ebin_path) + Code.prepend_path(ebin_path) + end + {:ok, %{ unload_modules_on_termination: unload_modules_on_termination, auto_termination: auto_termination, - server_supevisor: server_supevisor, + server_supervisor: server_supervisor, runtime_servers: [], initial_ignore_module_conflict: initial_ignore_module_conflict, original_standard_error: original_standard_error, parent_node: parent_node, - capture_orphan_logs: capture_orphan_logs + capture_orphan_logs: capture_orphan_logs, + tmp_dir: tmp_dir }} end @@ -138,6 +146,14 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do Node.disconnect(state.parent_node) end + if ebin_path = ebin_path(state.tmp_dir) do + Code.delete_path(ebin_path) + end + + if tmp_dir = state.tmp_dir do + File.rm_rf!(tmp_dir) + end + :ok end @@ -168,11 +184,28 @@ defmodule Livebook.Runtime.ErlDist.NodeManager do @impl true def handle_call({:start_runtime_server, opts}, _from, state) do + opts = Keyword.put_new(opts, :ebin_path, ebin_path(state.tmp_dir)) + {:ok, server_pid} = - DynamicSupervisor.start_child(state.server_supevisor, {ErlDist.RuntimeServer, opts}) + DynamicSupervisor.start_child(state.server_supervisor, {ErlDist.RuntimeServer, opts}) Process.monitor(server_pid) state = update_in(state.runtime_servers, &[server_pid | &1]) {:reply, server_pid, state} end + + defp make_tmp_dir() do + path = Path.join([System.tmp_dir!(), "livebook_runtime", random_id()]) + + if File.mkdir_p(path) == :ok do + path + end + end + + defp ebin_path(nil), do: nil + defp ebin_path(tmp_dir), do: Path.join(tmp_dir, "ebin") + + defp random_id() do + :crypto.strong_rand_bytes(20) |> Base.encode32(case: :lower) + end end diff --git a/lib/livebook/runtime/erl_dist/runtime_server.ex b/lib/livebook/runtime/erl_dist/runtime_server.ex index b354496f6..03900dd1e 100644 --- a/lib/livebook/runtime/erl_dist/runtime_server.ex +++ b/lib/livebook/runtime/erl_dist/runtime_server.ex @@ -25,7 +25,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do @memory_usage_interval 15_000 @doc """ - Starts the manager. + Starts the runtime server. Note: make sure to call `attach` within #{@await_owner_timeout}ms or the runtime server assumes it's not needed and terminates. @@ -39,6 +39,10 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do * `:extra_smart_cell_definitions` - a list of predefined smart cell definitions, that may be currently be unavailable, but should be reported together with their requirements + + * `:ebin_path` - a directory to write modules bytecode into. When + not specified, modules are not written to disk + """ def start_link(opts \\ []) do GenServer.start_link(__MODULE__, opts) @@ -193,7 +197,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do end @doc """ - Stops the manager. + Stops the runtime server. This results in all Livebook-related modules being unloaded from the runtime node. @@ -229,7 +233,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do extra_smart_cell_definitions: Keyword.get(opts, :extra_smart_cell_definitions, []), memory_timer_ref: nil, last_evaluator: nil, - initial_path: System.get_env("PATH", "") + initial_path: System.get_env("PATH", ""), + ebin_path: Keyword.get(opts, :ebin_path) }} end @@ -542,7 +547,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do state.evaluator_supervisor, send_to: state.owner, runtime_broadcast_to: state.runtime_broadcast_to, - object_tracker: state.object_tracker + object_tracker: state.object_tracker, + ebin_path: state.ebin_path ) Process.monitor(evaluator.pid) diff --git a/lib/livebook/runtime/evaluator.ex b/lib/livebook/runtime/evaluator.ex index 143c8db3e..22007f91b 100644 --- a/lib/livebook/runtime/evaluator.ex +++ b/lib/livebook/runtime/evaluator.ex @@ -63,6 +63,9 @@ defmodule Livebook.Runtime.Evaluator do # would take too much memory @evaluator_info_key :evaluator_info + # We stor the path in process dictionary, so that the tracer can access it + @ebin_path_key :ebin_path + @doc """ Starts an evaluator. @@ -79,6 +82,10 @@ defmodule Livebook.Runtime.Evaluator do * `:formatter` - a module implementing the `Livebook.Runtime.Evaluator.Formatter` behaviour, used for transforming evaluation result before sending it to the client. Defaults to identity + + * `:ebin_path` - a directory to write modules bytecode into. When + not specified, modules are not written to disk + """ @spec start_link(keyword()) :: {:ok, pid(), t()} | {:error, term()} def start_link(opts \\ []) do @@ -259,6 +266,7 @@ defmodule Livebook.Runtime.Evaluator do runtime_broadcast_to = Keyword.get(opts, :runtime_broadcast_to, send_to) object_tracker = Keyword.fetch!(opts, :object_tracker) formatter = Keyword.get(opts, :formatter, Evaluator.IdentityFormatter) + ebin_path = Keyword.get(opts, :ebin_path) {:ok, io_proxy} = Evaluator.IOProxy.start_link(self(), send_to, runtime_broadcast_to, object_tracker) @@ -278,6 +286,8 @@ defmodule Livebook.Runtime.Evaluator do contexts: %{} }) + Process.put(@ebin_path_key, ebin_path) + ignored_pdict_keys = Process.get_keys() |> MapSet.new() state = %{ @@ -330,9 +340,7 @@ defmodule Livebook.Runtime.Evaluator do if old_context = state.contexts[ref] do for module <- old_context.env.context_modules do - # If there is a deleted code for the module, we purge it first - :code.purge(module) - :code.delete(module) + delete_module!(module) end end @@ -397,9 +405,8 @@ defmodule Livebook.Runtime.Evaluator do {context, state} = pop_context(state, ref) for module <- context.env.context_modules do - # If there is a deleted code for the module, we purge it first - :code.purge(module) - :code.delete(module) + delete_module!(module) + # And we immediately purge the newly deleted code :code.purge(module) end @@ -750,4 +757,30 @@ defmodule Livebook.Runtime.Evaluator do |> Kernel.-(started_at) |> System.convert_time_unit(:native, :millisecond) end + + @doc false + def write_module!(module, bytecode) do + if ebin_path = ebin_path() do + ebin_path + |> Path.join("#{module}.beam") + |> File.write!(bytecode) + end + end + + defp delete_module!(module) do + # If there is a deleted code for the module, we purge it first + :code.purge(module) + + :code.delete(module) + + if ebin_path = ebin_path() do + ebin_path + |> Path.join("#{module}.beam") + |> File.rm!() + end + end + + defp ebin_path() do + Process.get(@ebin_path_key) + end end diff --git a/lib/livebook/runtime/evaluator/tracer.ex b/lib/livebook/runtime/evaluator/tracer.ex index aeef78add..82e601f7c 100644 --- a/lib/livebook/runtime/evaluator/tracer.ex +++ b/lib/livebook/runtime/evaluator/tracer.ex @@ -83,9 +83,10 @@ defmodule Livebook.Runtime.Evaluator.Tracer do {:remote_macro, _meta, module, _name, _arity} -> [{:module_used, module}, {:require_used, module}] - {:on_module, _bytecode, _ignore} -> + {:on_module, bytecode, _ignore} -> module = env.module vars = Map.keys(env.versioned_vars) + Evaluator.write_module!(module, bytecode) [{:module_defined, module, vars}, {:alias_used, module}] _ -> diff --git a/mix.exs b/mix.exs index 3a2c0ef58..680089615 100644 --- a/mix.exs +++ b/mix.exs @@ -16,6 +16,7 @@ defmodule Livebook.MixProject do name: "Livebook", description: @description, elixirc_paths: elixirc_paths(Mix.env()), + test_elixirc_options: [docs: true], start_permanent: Mix.env() == :prod, aliases: aliases(), deps: with_lock(target_deps(Mix.target()) ++ deps()), diff --git a/test/livebook/runtime/evaluator_test.exs b/test/livebook/runtime/evaluator_test.exs index fcb02e04d..4b6b85b00 100644 --- a/test/livebook/runtime/evaluator_test.exs +++ b/test/livebook/runtime/evaluator_test.exs @@ -3,13 +3,25 @@ defmodule Livebook.Runtime.EvaluatorTest do alias Livebook.Runtime.Evaluator - setup do + setup ctx do + ebin_path = + if ctx[:with_ebin_path] do + hash = ctx.test |> to_string() |> :erlang.md5() |> Base.encode32(padding: false) + path = ["tmp", inspect(ctx.module), hash, "ebin"] |> Path.join() |> Path.expand() + File.rm_rf!(path) + File.mkdir_p!(path) + Code.append_path(path) + path + end + {:ok, object_tracker} = start_supervised(Evaluator.ObjectTracker) {:ok, _pid, evaluator} = - start_supervised({Evaluator, [send_to: self(), object_tracker: object_tracker]}) + start_supervised( + {Evaluator, [send_to: self(), object_tracker: object_tracker, ebin_path: ebin_path]} + ) - %{evaluator: evaluator, object_tracker: object_tracker} + %{evaluator: evaluator, object_tracker: object_tracker, ebin_path: ebin_path} end defmacrop metadata do @@ -278,6 +290,7 @@ defmodule Livebook.Runtime.EvaluatorTest do assert_receive {:DOWN, ^ref, :process, ^widget_pid1, _reason} end + @tag :with_ebin_path test "raises when redefining a module in a different evaluation", %{evaluator: evaluator} do code = """ defmodule Livebook.Runtime.EvaluatorTest.Redefinition do @@ -303,6 +316,23 @@ defmodule Livebook.Runtime.EvaluatorTest do } }} end + + @tag :with_ebin_path + test "writes module bytecode to disk when :ebin_path is specified", + %{evaluator: evaluator, ebin_path: ebin_path} do + code = """ + defmodule Livebook.Runtime.EvaluatorTest.Disk do + @moduledoc "Test." + end + """ + + Evaluator.evaluate_code(evaluator, code, :code_1, []) + assert_receive {:runtime_evaluation_response, :code_1, {:ok, _}, metadata()} + + assert File.exists?(Path.join(ebin_path, "Elixir.Livebook.Runtime.EvaluatorTest.Disk.beam")) + + assert {:docs_v1, _, _, _, _, _, _} = Code.fetch_docs(Livebook.Runtime.EvaluatorTest.Disk) + end end describe "evaluate_code/6 identifier tracking" do @@ -732,6 +762,7 @@ defmodule Livebook.Runtime.EvaluatorTest do assert_receive {:DOWN, ^ref, :process, ^widget_pid1, _reason} end + @tag :with_ebin_path test "deletes modules defined by the given evaluation", %{evaluator: evaluator} do code = """ defmodule Livebook.Runtime.EvaluatorTest.ForgetEvaluation.Redefinition do