From ae7fbca0ba8ffd833cb4514a783f92e635593882 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonatan=20K=C5=82osko?= Date: Tue, 13 Dec 2022 20:19:29 +0100 Subject: [PATCH] Cleanup modules when evaluator terminates (#1582) --- .../runtime/erl_dist/runtime_server.ex | 22 +++++---- lib/livebook/runtime/evaluator.ex | 14 ++++-- lib/livebook/runtime/evaluator/io_proxy.ex | 49 ++++++++++++++++--- .../runtime/erl_dist/runtime_server_test.exs | 5 +- test/livebook/runtime/evaluator_test.exs | 20 ++++++++ 5 files changed, 90 insertions(+), 20 deletions(-) diff --git a/lib/livebook/runtime/erl_dist/runtime_server.ex b/lib/livebook/runtime/erl_dist/runtime_server.ex index ca7123d0e..27c894314 100644 --- a/lib/livebook/runtime/erl_dist/runtime_server.ex +++ b/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 diff --git a/lib/livebook/runtime/evaluator.ex b/lib/livebook/runtime/evaluator.ex index c0a4582c4..f4a0d186a 100644 --- a/lib/livebook/runtime/evaluator.ex +++ b/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!() diff --git a/lib/livebook/runtime/evaluator/io_proxy.ex b/lib/livebook/runtime/evaluator/io_proxy.ex index ac549568c..7477eb72c 100644 --- a/lib/livebook/runtime/evaluator/io_proxy.ex +++ b/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 diff --git a/test/livebook/runtime/erl_dist/runtime_server_test.exs b/test/livebook/runtime/erl_dist/runtime_server_test.exs index 04dc9285f..3552fcd09 100644 --- a/test/livebook/runtime/erl_dist/runtime_server_test.exs +++ b/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 diff --git a/test/livebook/runtime/evaluator_test.exs b/test/livebook/runtime/evaluator_test.exs index bf67a3314..11eb7113e 100644 --- a/test/livebook/runtime/evaluator_test.exs +++ b/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'''