Cleanup modules when evaluator terminates (#1582)

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -689,16 +689,20 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   end
 
   defp finish_scan_binding(ref, state) do
-    update_in(state.smart_cells[ref], fn info ->
-      Process.demonitor(info.scan_binding_monitor_ref, [:flush])
-      info = %{info | scan_binding_monitor_ref: nil}
+    if state.smart_cells[ref] do
+      update_in(state.smart_cells[ref], fn info ->
+        Process.demonitor(info.scan_binding_monitor_ref, [:flush])
+        info = %{info | scan_binding_monitor_ref: nil}
 
-      if info.scan_binding_pending do
-        scan_binding_async(ref, info, state)
-      else
-        info
-      end
-    end)
+        if info.scan_binding_pending do
+          scan_binding_async(ref, info, state)
+        else
+          info
+        end
+      end)
+    else
+      state
+    end
   end
 
   defp scan_binding_after_evaluation(state, locator) do

lib/livebook/runtime/evaluator.ex

@@ -28,6 +28,7 @@ defmodule Livebook.Runtime.Evaluator do
           evaluator_ref: reference(),
           formatter: module(),
           io_proxy: pid(),
+          io_proxy_monitor: reference(),
           send_to: pid(),
           runtime_broadcast_to: pid(),
           object_tracker: pid(),
@@ -271,7 +272,9 @@ defmodule Livebook.Runtime.Evaluator do
     ebin_path = Keyword.get(opts, :ebin_path)
 
     {:ok, io_proxy} =
-      Evaluator.IOProxy.start_link(self(), send_to, runtime_broadcast_to, object_tracker)
+      Evaluator.IOProxy.start(self(), send_to, runtime_broadcast_to, object_tracker, ebin_path)
+
+    io_proxy_monitor = Process.monitor(io_proxy)
 
     # Use the dedicated IO device as the group leader, so that
     # intercepts all :stdio requests and also handles Livebook
@@ -296,6 +299,7 @@ defmodule Livebook.Runtime.Evaluator do
       evaluator_ref: evaluator_ref,
       formatter: formatter,
       io_proxy: io_proxy,
+      io_proxy_monitor: io_proxy_monitor,
       send_to: send_to,
       runtime_broadcast_to: runtime_broadcast_to,
       object_tracker: object_tracker,
@@ -320,6 +324,9 @@ defmodule Livebook.Runtime.Evaluator do
       {:cast, ^evaluator_ref, message} ->
         {:noreply, state} = handle_cast(message, state)
         loop(state)
+
+      {:DOWN, ref, :process, _pid, reason} when ref == state.io_proxy_monitor ->
+        exit(reason)
     end
   end
 
@@ -773,13 +780,14 @@ defmodule Livebook.Runtime.Evaluator do
     end
   end
 
-  defp delete_module!(module) do
+  @doc false
+  def delete_module!(module, ebin_path \\ ebin_path()) 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
+    if ebin_path do
       ebin_path
       |> Path.join("#{module}.beam")
       |> File.rm!()

lib/livebook/runtime/evaluator/io_proxy.ex

@@ -27,9 +27,23 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
   `:send_to` process, so this device serves as a proxy. Make sure
   to also call configure/3` before every evaluation.
   """
-  @spec start_link(pid(), pid(), pid(), pid()) :: GenServer.on_start()
-  def start_link(evaluator, send_to, runtime_broadcast_to, object_tracker) do
-    GenServer.start_link(__MODULE__, {evaluator, send_to, runtime_broadcast_to, object_tracker})
+  @spec start(pid(), pid(), pid(), pid(), String.t() | nil) :: GenServer.on_start()
+  def start(evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path) do
+    GenServer.start(
+      __MODULE__,
+      {evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path}
+    )
+  end
+
+  @doc """
+  Linking version of `start/4`.
+  """
+  @spec start_link(pid(), pid(), pid(), pid(), String.t() | nil) :: GenServer.on_start()
+  def start_link(evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path) do
+    GenServer.start_link(
+      __MODULE__,
+      {evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path}
+    )
   end
 
   @doc """
@@ -77,9 +91,12 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
   end
 
   @impl true
-  def init({evaluator, send_to, runtime_broadcast_to, object_tracker}) do
+  def init({evaluator, send_to, runtime_broadcast_to, object_tracker, ebin_path}) do
+    evaluator_monitor = Process.monitor(evaluator)
+
     {:ok,
      %{
+       evaluator_monitor: evaluator_monitor,
        encoding: :unicode,
        ref: nil,
        file: nil,
@@ -90,7 +107,9 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
        send_to: send_to,
        runtime_broadcast_to: runtime_broadcast_to,
        object_tracker: object_tracker,
-       tracer_info: %Evaluator.Tracer{}
+       ebin_path: ebin_path,
+       tracer_info: %Evaluator.Tracer{},
+       modules_defined: MapSet.new()
      }}
   end
 
@@ -114,7 +133,12 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
   end
 
   def handle_call(:get_tracer_info, _from, state) do
-    {:reply, state.tracer_info, state}
+    modules_defined =
+      state.tracer_info.modules_defined
+      |> Map.keys()
+      |> Enum.into(state.modules_defined)
+
+    {:reply, state.tracer_info, %{state | modules_defined: modules_defined}}
   end
 
   @impl true
@@ -128,6 +152,19 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
     {:noreply, flush_buffer(state)}
   end
 
+  def handle_info({:DOWN, ref, :process, _pid, reason}, state)
+      when ref == state.evaluator_monitor do
+    cleanup(state)
+    {:stop, reason, state}
+  end
+
+  defp cleanup(state) do
+    # Remove all modules defined during evaluation
+    for module <- state.modules_defined, function_exported?(module, :module_info, 1) do
+      Evaluator.delete_module!(module, state.ebin_path)
+    end
+  end
+
   defp io_request({:put_chars, chars} = req, state) do
     put_chars(:latin1, chars, req, state)
   end

test/livebook/runtime/erl_dist/runtime_server_test.exs

@@ -194,15 +194,16 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServerTest do
     end
   end
 
+  @tag capture_log: true
   test "notifies the owner when an evaluator goes down", %{pid: pid} do
     code = """
-    spawn_link(fn -> Process.exit(self(), :kill) end)
+    Task.async(fn -> raise "error" end)
     """
 
     RuntimeServer.evaluate_code(pid, code, {:c1, :e1}, [])
 
     assert_receive {:runtime_container_down, :c1, message}
-    assert message =~ "killed"
+    assert message =~ "(RuntimeError) error"
   end
 
   describe "smart cells" do

test/livebook/runtime/evaluator_test.exs

@@ -348,6 +348,26 @@ defmodule Livebook.Runtime.EvaluatorTest do
       refute Code.ensure_loaded?(Livebook.Runtime.EvaluatorTest.Raised)
     end
 
+    @tag :with_ebin_path
+    @tag capture_log: true
+    test "deletes defined modules on termination", %{evaluator: evaluator} do
+      code = """
+      defmodule Livebook.Runtime.EvaluatorTest.Exited do
+      end
+
+      Task.async(fn -> raise "error" end)
+      """
+
+      {:group_leader, gl} = Process.info(evaluator.pid, :group_leader)
+
+      Evaluator.evaluate_code(evaluator, code, :code_1, [])
+
+      ref = Process.monitor(gl)
+      assert_receive {:DOWN, ^ref, :process, ^gl, _reason}
+
+      refute Code.ensure_loaded?(Livebook.Runtime.EvaluatorTest.Exited)
+    end
+
     @tag :with_ebin_path
     test "runs doctests when a module is defined", %{evaluator: evaluator} do
       code = ~S'''