Persist modules bytecode to disk (#1521)

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

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)

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

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}]
 
       _ ->

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()),

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