Trigger garbage collection before evaluation (#1593)

lib/livebook/runtime/evaluator.ex

@@ -337,105 +337,11 @@ defmodule Livebook.Runtime.Evaluator do
   end
 
   defp handle_cast({:evaluate_code, code, ref, parent_refs, opts}, state) do
-    Evaluator.ObjectTracker.remove_reference_sync(state.object_tracker, {self(), ref})
+    do_evaluate_code(code, ref, parent_refs, opts, state)
-
-    context = get_context(state, parent_refs)
-    file = Keyword.get(opts, :file, "nofile")
-    context = put_in(context.env.file, file)
-
-    Evaluator.IOProxy.configure(state.io_proxy, ref, file)
-
-    set_pdict(context, state.ignored_pdict_keys)
-
-    if old_context = state.contexts[ref] do
-      for module <- old_context.env.context_modules do
-        delete_module!(module)
-      end
-    end
-
-    start_time = System.monotonic_time()
-    eval_result = eval(code, context.binding, context.env)
-    evaluation_time_ms = time_diff_ms(start_time)
-
-    tracer_info = Evaluator.IOProxy.get_tracer_info(state.io_proxy)
-
-    {new_context, result, code_error, identifiers_used, identifiers_defined} =
-      case eval_result do
-        {:ok, value, binding, env} ->
-          context_id = random_id()
-
-          new_context = %{
-            id: context_id,
-            binding: binding,
-            env: prune_env(env, tracer_info),
-            pdict: current_pdict(state)
-          }
-
-          {identifiers_used, identifiers_defined} =
-            identifier_dependencies(new_context, tracer_info, context)
-
-          result = {:ok, value}
-          {new_context, result, nil, identifiers_used, identifiers_defined}
-
-        {:error, kind, error, stacktrace, code_error} ->
-          for {module, _} <- tracer_info.modules_defined do
-            delete_module!(module)
-          end
-
-          result = {:error, kind, error, stacktrace}
-          identifiers_used = :unknown
-          identifiers_defined = %{}
-          # Empty context
-          new_context = initial_context()
-          {new_context, result, code_error, identifiers_used, identifiers_defined}
-      end
-
-    if ebin_path() do
-      Livebook.Runtime.Evaluator.Doctests.run(new_context.env.context_modules)
-    end
-
-    state = put_context(state, ref, new_context)
-
-    Evaluator.IOProxy.flush(state.io_proxy)
-    Evaluator.IOProxy.clear_input_cache(state.io_proxy)
-
-    output = state.formatter.format_result(result)
-
-    metadata = %{
-      evaluation_time_ms: evaluation_time_ms,
-      memory_usage: memory(),
-      code_error: code_error,
-      identifiers_used: identifiers_used,
-      identifiers_defined: identifiers_defined
-    }
-
-    send(state.send_to, {:runtime_evaluation_response, ref, output, metadata})
-
-    if on_finish = opts[:on_finish] do
-      on_finish.(result)
-    end
-
-    :erlang.garbage_collect(self())
-    {:noreply, state}
   end
 
   defp handle_cast({:forget_evaluation, ref}, state) do
-    {context, state} = pop_context(state, ref)
+    do_forget_evaluation(ref, state)
-
-    if context do
-      for module <- context.env.context_modules do
-        delete_module!(module)
-
-        # And we immediately purge the newly deleted code
-        :code.purge(module)
-      end
-
-      Evaluator.ObjectTracker.remove_reference_sync(state.object_tracker, {self(), ref})
-
-      :erlang.garbage_collect(self())
-    end
-
-    {:noreply, state}
   end
 
   defp handle_cast({:peek_context, parent_refs, fun}, state) do
@@ -483,6 +389,126 @@ defmodule Livebook.Runtime.Evaluator do
     {:reply, result, state}
   end
 
+  defp do_evaluate_code(code, ref, parent_refs, opts, state) do
+    {old_context, state} = pop_in(state.contexts[ref])
+
+    if old_context do
+      for module <- old_context.env.context_modules do
+        delete_module(module)
+      end
+    end
+
+    # We remove the old context from state and jump to a tail-recursive
+    # function. This way we are sure there is no reference to the old
+    # state and we can garbage collect the old context before the evaluation
+    continue_do_evaluate_code(code, ref, parent_refs, opts, state)
+  end
+
+  defp continue_do_evaluate_code(code, ref, parent_refs, opts, state) do
+    :erlang.garbage_collect(self())
+
+    Evaluator.ObjectTracker.remove_reference_sync(state.object_tracker, {self(), ref})
+
+    context = get_context(state, parent_refs)
+    file = Keyword.get(opts, :file, "nofile")
+    context = put_in(context.env.file, file)
+
+    Evaluator.IOProxy.configure(state.io_proxy, ref, file)
+
+    set_pdict(context, state.ignored_pdict_keys)
+
+    start_time = System.monotonic_time()
+    eval_result = eval(code, context.binding, context.env)
+    evaluation_time_ms = time_diff_ms(start_time)
+
+    tracer_info = Evaluator.IOProxy.get_tracer_info(state.io_proxy)
+
+    {new_context, result, code_error, identifiers_used, identifiers_defined} =
+      case eval_result do
+        {:ok, value, binding, env} ->
+          context_id = random_id()
+
+          new_context = %{
+            id: context_id,
+            binding: binding,
+            env: prune_env(env, tracer_info),
+            pdict: current_pdict(state)
+          }
+
+          {identifiers_used, identifiers_defined} =
+            identifier_dependencies(new_context, tracer_info, context)
+
+          result = {:ok, value}
+          {new_context, result, nil, identifiers_used, identifiers_defined}
+
+        {:error, kind, error, stacktrace, code_error} ->
+          for {module, _} <- tracer_info.modules_defined do
+            delete_module(module)
+          end
+
+          result = {:error, kind, error, stacktrace}
+          identifiers_used = :unknown
+          identifiers_defined = %{}
+          # Empty context
+          new_context = initial_context()
+          {new_context, result, code_error, identifiers_used, identifiers_defined}
+      end
+
+    if ebin_path() do
+      Livebook.Runtime.Evaluator.Doctests.run(new_context.env.context_modules)
+    end
+
+    state = put_context(state, ref, new_context)
+
+    Evaluator.IOProxy.flush(state.io_proxy)
+    Evaluator.IOProxy.clear_input_cache(state.io_proxy)
+
+    output = state.formatter.format_result(result)
+
+    metadata = %{
+      evaluation_time_ms: evaluation_time_ms,
+      memory_usage: memory(),
+      code_error: code_error,
+      identifiers_used: identifiers_used,
+      identifiers_defined: identifiers_defined
+    }
+
+    send(state.send_to, {:runtime_evaluation_response, ref, output, metadata})
+
+    if on_finish = opts[:on_finish] do
+      on_finish.(result)
+    end
+
+    :erlang.garbage_collect(self())
+
+    {:noreply, state}
+  end
+
+  defp do_forget_evaluation(ref, state) do
+    {context, state} = pop_context(state, ref)
+
+    if context do
+      for module <- context.env.context_modules do
+        delete_module(module)
+
+        # And we immediately purge the newly deleted code
+        :code.purge(module)
+      end
+
+      Evaluator.ObjectTracker.remove_reference_sync(state.object_tracker, {self(), ref})
+    end
+
+    continue_do_forget_evaluation(context != nil, state)
+  end
+
+  defp continue_do_forget_evaluation(context?, state) do
+    if context? do
+      :erlang.garbage_collect(self())
+    end
+
+    {:noreply, state}
+  end
+
   defp put_context(state, ref, context) do
     update_evaluator_info(fn info ->
       put_in(info.contexts[ref], {context.id, context.env})
@@ -781,7 +807,7 @@ defmodule Livebook.Runtime.Evaluator do
   end
 
   @doc false
-  def delete_module!(module, ebin_path \\ ebin_path()) do
+  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)
 
@@ -790,7 +816,7 @@ defmodule Livebook.Runtime.Evaluator do
     if ebin_path do
       ebin_path
       |> Path.join("#{module}.beam")
-      |> File.rm!()
+      |> File.rm()
     end
   end
 

lib/livebook/runtime/evaluator/io_proxy.ex

@@ -161,7 +161,7 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do
   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)
+      Evaluator.delete_module(module, state.ebin_path)
     end
   end