Improve reproducability of module definitions (#1518)

lib/livebook/runtime/evaluator.ex

@@ -328,17 +328,25 @@ defmodule Livebook.Runtime.Evaluator do
 
     set_pdict(context, state.ignored_pdict_keys)
 
+    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)
+      end
+    end
+
     start_time = System.monotonic_time()
     eval_result = eval(code, context.binding, context.env)
     evaluation_time_ms = time_diff_ms(start_time)
 
-    {result_context, result, code_error, identifiers_used, identifiers_defined} =
+    {new_context, result, code_error, identifiers_used, identifiers_defined} =
       case eval_result do
         {:ok, value, binding, env} ->
           tracer_info = Evaluator.IOProxy.get_tracer_info(state.io_proxy)
           context_id = random_id()
 
-          result_context = %{
+          new_context = %{
             id: context_id,
             binding: binding,
             env: prune_env(env, tracer_info),
@@ -346,21 +354,21 @@ defmodule Livebook.Runtime.Evaluator do
           }
 
           {identifiers_used, identifiers_defined} =
-            identifier_dependencies(result_context, tracer_info, context)
+            identifier_dependencies(new_context, tracer_info, context)
 
           result = {:ok, value}
-          {result_context, result, nil, identifiers_used, identifiers_defined}
+          {new_context, result, nil, identifiers_used, identifiers_defined}
 
         {:error, kind, error, stacktrace, code_error} ->
           result = {:error, kind, error, stacktrace}
           identifiers_used = :unknown
           identifiers_defined = %{}
           # Empty context
-          result_context = initial_context()
-          {result_context, result, code_error, identifiers_used, identifiers_defined}
+          new_context = initial_context()
+          {new_context, result, code_error, identifiers_used, identifiers_defined}
       end
 
-    state = put_context(state, ref, result_context)
+    state = put_context(state, ref, new_context)
 
     Evaluator.IOProxy.flush(state.io_proxy)
     Evaluator.IOProxy.clear_input_cache(state.io_proxy)
@@ -386,7 +394,16 @@ defmodule Livebook.Runtime.Evaluator do
   end
 
   defp handle_cast({:forget_evaluation, ref}, state) do
-    state = delete_context(state, ref)
+    {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)
+      # 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())
@@ -446,14 +463,13 @@ defmodule Livebook.Runtime.Evaluator do
     put_in(state.contexts[ref], context)
   end
 
-  defp delete_context(state, ref) do
+  defp pop_context(state, ref) do
     update_evaluator_info(fn info ->
       {_, info} = pop_in(info.contexts[ref])
       info
     end)
 
-    {_, state} = pop_in(state.contexts[ref])
-    state
+    pop_in(state.contexts[ref])
   end
 
   defp update_evaluator_info(fun) do
@@ -494,6 +510,7 @@ defmodule Livebook.Runtime.Evaluator do
     env
     |> Map.replace!(:aliases, Map.to_list(tracer_info.aliases_defined))
     |> Map.replace!(:requires, MapSet.to_list(tracer_info.requires_defined))
+    |> Map.replace!(:context_modules, Map.keys(tracer_info.modules_defined))
   end
 
   defp merge_context(prev_context, context) do
@@ -523,7 +540,7 @@ defmodule Livebook.Runtime.Evaluator do
     end)
     |> Map.update!(:aliases, &Keyword.merge(prev_env.aliases, &1))
     |> Map.update!(:requires, &:ordsets.union(prev_env.requires, &1))
-    |> Map.replace!(:context_modules, [])
+    |> Map.update!(:context_modules, &(&1 ++ prev_env.context_modules))
   end
 
   @compile {:no_warn_undefined, {Code, :eval_quoted_with_env, 4}}
@@ -590,7 +607,8 @@ defmodule Livebook.Runtime.Evaluator do
           into: identifiers_used
 
     identifiers_defined =
-      for {module, {version, _vars}} <- tracer_info.modules_defined,
+      for {module, _vars} <- tracer_info.modules_defined,
+          version = module.__info__(:md5),
           do: {{:module, module}, version},
           into: identifiers_defined
 
@@ -667,7 +685,7 @@ defmodule Livebook.Runtime.Evaluator do
     # Note that :prune_binding removes variables used by modules
     # (unless used outside), so we get those from the tracer
     module_used_vars =
-      for {_module, {_version, vars}} <- tracer_info.modules_defined,
+      for {_module, vars} <- tracer_info.modules_defined,
           var <- vars,
           into: MapSet.new(),
           do: var
@@ -686,6 +704,8 @@ defmodule Livebook.Runtime.Evaluator do
         do: var
   end
 
+  defp prune_stacktrace([{Livebook.Runtime.Evaluator.Tracer, _fun, _arity, _meta} | _]), do: []
+
   # Adapted from https://github.com/elixir-lang/elixir/blob/1c1654c88adfdbef38ff07fc30f6fbd34a542c07/lib/iex/lib/iex/evaluator.ex#L355-L372
   # TODO: Remove else branch once we depend on the versions below
   if System.otp_release() >= "25" do

lib/livebook/runtime/evaluator/tracer.ex

@@ -21,7 +21,7 @@ defmodule Livebook.Runtime.Evaluator.Tracer do
 
   @doc false
   def trace(event, env) do
-    case to_updates(event, env) do
+    case event_to_updates(event, env) do
       [] ->
         :ok
 
@@ -33,11 +33,21 @@ defmodule Livebook.Runtime.Evaluator.Tracer do
     :ok
   end
 
-  defp to_updates(event, env) do
+  defp event_to_updates(event, env) do
     # Note that import/require/alias/defmodule don't trigger `:alias_reference`
     # for the used alias, so we add it explicitly
 
     case event do
+      :start ->
+        if Code.ensure_loaded?(env.module) do
+          raise CompileError,
+            line: env.line,
+            file: env.file,
+            description: "module #{inspect(env.module)} is already defined"
+        end
+
+        []
+
       {:import, _meta, module, _opts} ->
         if(env.module, do: [], else: [:import_defined]) ++
           [{:module_used, module}, {:alias_used, module}]
@@ -73,11 +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
-        version = :erlang.md5(bytecode)
         vars = Map.keys(env.versioned_vars)
-        [{:module_defined, module, version, vars}, {:alias_used, module}]
+        [{:module_defined, module, vars}, {:alias_used, module}]
 
       _ ->
         []
@@ -96,8 +105,8 @@ defmodule Livebook.Runtime.Evaluator.Tracer do
     update_in(info.modules_used, &MapSet.put(&1, module))
   end
 
-  defp apply_update(info, {:module_defined, module, version, vars}) do
-    put_in(info.modules_defined[module], {version, vars})
+  defp apply_update(info, {:module_defined, module, vars}) do
+    put_in(info.modules_defined[module], vars)
   end
 
   defp apply_update(info, {:alias_used, alias}) do

test/livebook/runtime/evaluator_test.exs

@@ -277,6 +277,32 @@ defmodule Livebook.Runtime.EvaluatorTest do
       ref = Process.monitor(widget_pid1)
       assert_receive {:DOWN, ^ref, :process, ^widget_pid1, _reason}
     end
+
+    test "raises when redefining a module in a different evaluation", %{evaluator: evaluator} do
+      code = """
+      defmodule Livebook.Runtime.EvaluatorTest.Redefinition do
+      end
+      """
+
+      Evaluator.evaluate_code(evaluator, code, :code_1, [])
+      assert_receive {:runtime_evaluation_response, :code_1, {:ok, _}, metadata()}
+
+      # Redefining in the same evaluation works
+      Evaluator.evaluate_code(evaluator, code, :code_1, [])
+      assert_receive {:runtime_evaluation_response, :code_1, {:ok, _}, metadata()}
+
+      Evaluator.evaluate_code(evaluator, code, :code_2, [], file: "file.ex")
+
+      assert_receive {:runtime_evaluation_response, :code_2,
+                      {:error, :error, %CompileError{}, []},
+                      %{
+                        code_error: %{
+                          line: 1,
+                          description:
+                            "module Livebook.Runtime.EvaluatorTest.Redefinition is already defined"
+                        }
+                      }}
+    end
   end
 
   describe "evaluate_code/6 identifier tracking" do
@@ -584,6 +610,9 @@ defmodule Livebook.Runtime.EvaluatorTest do
       Process.put(:x, 1)
       Process.put(:y, 1)
       Process.put(:z, 1)
+
+      defmodule Livebook.Runtime.EvaluatorTest.Identifiers.ContextMergingOne do
+      end
       """
       |> eval(evaluator, 0)
 
@@ -598,6 +627,9 @@ defmodule Livebook.Runtime.EvaluatorTest do
 
       Process.put(:y, 2)
       Process.delete(:z)
+
+      defmodule Livebook.Runtime.EvaluatorTest.Identifiers.ContextMergingTwo do
+      end
       """
       |> eval(evaluator, 1)
 
@@ -616,6 +648,10 @@ defmodule Livebook.Runtime.EvaluatorTest do
 
       assert context.env.versioned_vars == %{{:x, nil} => 0, {:y, nil} => 1}
 
+      assert context.env.context_modules == [
+               Livebook.Runtime.EvaluatorTest.Identifiers.ContextMergingOne
+             ]
+
       assert context.pdict == %{x: 1, y: 1, z: 1}
 
       # Evaluation 1 context
@@ -635,6 +671,10 @@ defmodule Livebook.Runtime.EvaluatorTest do
 
       assert context.env.versioned_vars == %{{:y, nil} => 0}
 
+      assert context.env.context_modules == [
+               Livebook.Runtime.EvaluatorTest.Identifiers.ContextMergingTwo
+             ]
+
       # Process dictionary is not diffed
       assert context.pdict == %{x: 1, y: 2}
 
@@ -654,6 +694,11 @@ defmodule Livebook.Runtime.EvaluatorTest do
 
       assert context.env.versioned_vars == %{{:x, nil} => 0, {:y, nil} => 1}
 
+      assert context.env.context_modules == [
+               Livebook.Runtime.EvaluatorTest.Identifiers.ContextMergingTwo,
+               Livebook.Runtime.EvaluatorTest.Identifiers.ContextMergingOne
+             ]
+
       assert context.pdict == %{x: 1, y: 2}
     end
   end
@@ -686,6 +731,22 @@ defmodule Livebook.Runtime.EvaluatorTest do
 
       assert_receive {:DOWN, ^ref, :process, ^widget_pid1, _reason}
     end
+
+    test "deletes modules defined by the given evaluation", %{evaluator: evaluator} do
+      code = """
+      defmodule Livebook.Runtime.EvaluatorTest.ForgetEvaluation.Redefinition do
+      end
+      """
+
+      Evaluator.evaluate_code(evaluator, code, :code_1, [])
+      assert_receive {:runtime_evaluation_response, :code_1, {:ok, _}, metadata()}
+
+      Evaluator.forget_evaluation(evaluator, :code_1)
+
+      # Define the module in a different evaluation
+      Evaluator.evaluate_code(evaluator, code, :code_2, [])
+      assert_receive {:runtime_evaluation_response, :code_2, {:ok, _}, metadata()}
+    end
   end
 
   describe "initialize_from/3" do