Cleanup modules when evaluator terminates (#1582)

This commit is contained in:
Jonatan Kłosko 2022-12-13 20:19:29 +01:00 committed by GitHub
parent 600666fd46
commit ae7fbca0ba
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 90 additions and 20 deletions

View file

@ -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

View file

@ -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!()

View file

@ -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

View file

@ -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

View file

@ -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'''