Improvements to memory tracking (#917)

* Address race condition on cancel timer
* Include memory measurement as part of evaluation metadata
* Move periodic resource computation to a single process
* Have a explicit call out for total memory

lib/livebook/application.ex

@@ -16,6 +16,8 @@ defmodule Livebook.Application do
         LivebookWeb.Telemetry,
         # Start the PubSub system
         {Phoenix.PubSub, name: Livebook.PubSub},
+        # Periodid measurement of system resources
+        Livebook.SystemResources,
         # Start the tracker server on this node
         {Livebook.Tracker, pubsub_server: Livebook.PubSub},
         # Start the supervisor dynamically managing sessions

lib/livebook/evaluator.ex

@@ -68,6 +68,31 @@ defmodule Livebook.Evaluator do
     end
   end
 
+  @doc """
+  Computes the memory usage from this evaluator node.
+  """
+  @spec memory :: Livebook.Runtime.runtime_memory()
+  def memory do
+    %{
+      total: total,
+      processes: processes,
+      atom: atom,
+      binary: binary,
+      code: code,
+      ets: ets
+    } = Map.new(:erlang.memory())
+
+    %{
+      total: total,
+      processes: processes,
+      atom: atom,
+      binary: binary,
+      code: code,
+      ets: ets,
+      other: total - processes - atom - binary - code - ets
+    }
+  end
+
   @doc """
   Asynchronously parses and evaluates the given code.
 
@@ -85,8 +110,6 @@ defmodule Livebook.Evaluator do
     * `:file` - file to which the evaluated code belongs. Most importantly,
       this has an impact on the value of `__DIR__`.
 
-    * `:notify_to` - a pid to be notified when an evaluation is finished.
-       The process should expect a `{:evaluation_finished, ref}` message.
   """
   @spec evaluate_code(t(), pid(), String.t(), ref(), ref() | nil, keyword()) :: :ok
   def evaluate_code(evaluator, send_to, code, ref, prev_ref \\ nil, opts \\ []) when ref != nil do
@@ -235,7 +258,6 @@ defmodule Livebook.Evaluator do
     file = Keyword.get(opts, :file, "nofile")
     context = put_in(context.env.file, file)
     start_time = System.monotonic_time()
-    notify_to = Keyword.get(opts, :notify_to)
 
     {result_context, response} =
       case eval(code, context.binding, context.env) do
@@ -257,9 +279,8 @@ defmodule Livebook.Evaluator do
     Evaluator.IOProxy.clear_input_cache(state.io_proxy)
 
     output = state.formatter.format_response(response)
-    metadata = %{evaluation_time_ms: evaluation_time_ms}
+    metadata = %{evaluation_time_ms: evaluation_time_ms, memory_usage: memory()}
     send(send_to, {:evaluation_response, ref, output, metadata})
-    if notify_to, do: send(notify_to, {:evaluation_finished, ref})
 
     :erlang.garbage_collect(self())
     {:noreply, state}

lib/livebook/runtime.ex

@@ -127,7 +127,6 @@ defprotocol Livebook.Runtime do
           ets: non_neg_integer(),
           other: non_neg_integer(),
           processes: non_neg_integer(),
-          system: non_neg_integer(),
           total: non_neg_integer()
         }
 
@@ -174,7 +173,7 @@ defprotocol Livebook.Runtime do
 
     * `{:evaluation_response, ref, output, metadata}` - final
       result of the evaluation. Recognised metadata entries
-      are: `evaluation_time_ms`
+      are: `evaluation_time_ms` and `memory_usage`
 
   The output may include input fields. The evaluation may then
   request the current value of a previously rendered input by

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -179,14 +179,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     end
   end
 
-  def handle_info({:evaluation_finished, _ref}, state) do
-    send_memory_usage(state)
-    Process.cancel_timer(state.memory_timer_ref)
-    {:noreply, schedule_memory_usage(state)}
-  end
-
   def handle_info(:memory_usage, state) do
-    send_memory_usage(state)
+    send(state.owner, {:memory_usage, Evaluator.memory()})
     {:noreply, schedule_memory_usage(state)}
   end
 
@@ -204,7 +198,6 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
         state
       ) do
     state = ensure_evaluator(state, container_ref)
-    opts = Keyword.put(opts, :notify_to, self())
 
     prev_evaluation_ref =
       case prev_locator do
@@ -316,17 +309,4 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     ref = Process.send_after(self(), :memory_usage, @memory_usage_interval)
     %{state | memory_timer_ref: ref}
   end
-
-  defp send_memory_usage(state) do
-    memory =
-      :erlang.memory()
-      |> Enum.into(%{})
-      |> Map.drop([:processes_used, :atom_used])
-
-    other =
-      memory.total - memory.processes - memory.atom - memory.binary - memory.code - memory.ets
-
-    memory = Map.put(memory, :other, other)
-    send(state.owner, {:memory_usage, memory})
-  end
 end

lib/livebook/session.ex

@@ -55,8 +55,6 @@ defmodule Livebook.Session do
   alias Livebook.Users.User
   alias Livebook.Notebook.{Cell, Section}
 
-  @memory_usage_interval 15_000
-
   @type t :: %__MODULE__{
           id: id(),
           pid: pid(),
@@ -76,14 +74,13 @@ defmodule Livebook.Session do
           autosave_timer_ref: reference() | nil,
           save_task_pid: pid() | nil,
           saved_default_file: FileSystem.File.t() | nil,
-          system_memory_timer_ref: reference() | nil,
           memory_usage: memory_usage()
         }
 
   @type memory_usage ::
           %{
             runtime: Livebook.Runtime.runtime_memory() | nil,
-            system: Livebook.Utils.system_memory()
+            system: Livebook.SystemResources.memory()
           }
 
   @typedoc """
@@ -458,7 +455,7 @@ defmodule Livebook.Session do
              do: dump_images(state, images),
              else: :ok
            ) do
-      state = state |> schedule_autosave() |> schedule_system_memory_update()
+      state = schedule_autosave(state)
       {:ok, state}
     else
       {:error, error} ->
@@ -479,8 +476,7 @@ defmodule Livebook.Session do
         autosave_path: opts[:autosave_path],
         save_task_pid: nil,
         saved_default_file: nil,
-        system_memory_timer_ref: nil,
+        memory_usage: %{runtime: nil, system: Livebook.SystemResources.memory()}
-        memory_usage: %{runtime: nil, system: Utils.fetch_system_memory()}
       }
 
       {:ok, state}
@@ -521,11 +517,6 @@ defmodule Livebook.Session do
     end
   end
 
-  defp schedule_system_memory_update(state) do
-    ref = Process.send_after(self(), :system_memory, @memory_usage_interval)
-    %{state | system_memory_timer_ref: ref}
-  end
-
   @impl true
   def handle_call(:describe_self, _from, state) do
     {:reply, self_from_state(state), state}
@@ -781,8 +772,14 @@ defmodule Livebook.Session do
   end
 
   def handle_info({:evaluation_response, cell_id, response, metadata}, state) do
+    {memory_usage, metadata} = Map.pop(metadata, :memory_usage)
     operation = {:add_cell_evaluation_response, self(), cell_id, response, metadata}
-    {:noreply, handle_operation(state, operation)}
+
+    {:noreply,
+     state
+     |> put_memory_usage(memory_usage)
+     |> handle_operation(operation)
+     |> notify_update()}
   end
 
   def handle_info({:evaluation_input, cell_id, reply_to, input_id}, state) do
@@ -832,17 +829,8 @@ defmodule Livebook.Session do
     {:noreply, handle_save_finished(state, result, file, default?)}
   end
 
-  def handle_info(:system_memory, state) do
-    {:noreply, state |> update_system_memory_usage() |> schedule_system_memory_update()}
-  end
-
   def handle_info({:memory_usage, runtime_memory}, state) do
-    Process.cancel_timer(state.system_memory_timer_ref)
+    {:noreply, state |> put_memory_usage(runtime_memory) |> notify_update()}
-    system_memory = Utils.fetch_system_memory()
-    memory = %{runtime: runtime_memory, system: system_memory}
-    state = %{state | memory_usage: memory}
-    notify_update(state)
-    {:noreply, state}
   end
 
   def handle_info(_message, state), do: {:noreply, state}
@@ -1010,16 +998,15 @@ defmodule Livebook.Session do
 
   defp after_operation(state, _prev_state, {:set_notebook_name, _pid, _name}) do
     notify_update(state)
-    state
   end
 
   defp after_operation(state, _prev_state, {:set_runtime, _pid, runtime}) do
     if runtime do
       state
     else
-      put_in(state.memory_usage.runtime, nil)
+      state
-      |> update_system_memory_usage()
+      |> put_memory_usage(nil)
-      |> schedule_system_memory_update()
+      |> notify_update()
     end
   end
 
@@ -1040,8 +1027,6 @@ defmodule Livebook.Session do
     end
 
     notify_update(state)
-
-    state
   end
 
   defp after_operation(
@@ -1166,10 +1151,15 @@ defmodule Livebook.Session do
     Phoenix.PubSub.broadcast(Livebook.PubSub, "sessions:#{session_id}", message)
   end
 
+  defp put_memory_usage(state, runtime) do
+    put_in(state.memory_usage, %{runtime: runtime, system: Livebook.SystemResources.memory()})
+  end
+
   defp notify_update(state) do
     session = self_from_state(state)
     Livebook.Sessions.update_session(session)
     broadcast_message(state.session_id, {:session_updated, session})
+    state
   end
 
   defp maybe_save_notebook_async(state) do
@@ -1310,10 +1300,4 @@ defmodule Livebook.Session do
 
   defp container_ref_for_section(%{parent_id: nil}), do: :main_flow
   defp container_ref_for_section(section), do: section.id
-
-  defp update_system_memory_usage(state) do
-    state = put_in(state.memory_usage.system, Utils.fetch_system_memory())
-    notify_update(state)
-    state
-  end
 end

lib/livebook/system_resources.ex

@@ -0,0 +1,37 @@
+defmodule Livebook.SystemResources do
+  # Periodically compute system resource usage.
+  @moduledoc false
+  @type memory :: %{total: non_neg_integer(), free: non_neg_integer()}
+
+  use GenServer
+  @name __MODULE__
+
+  @spec memory() :: memory()
+  def memory do
+    :ets.lookup_element(@name, :memory, 2)
+  end
+
+  @doc false
+  def start_link(_opts) do
+    GenServer.start_link(__MODULE__, :ok, name: @name)
+  end
+
+  @impl true
+  def init(:ok) do
+    :ets.new(@name, [:set, :named_table, :protected])
+    measure_and_schedule()
+    {:ok, %{}}
+  end
+
+  @impl true
+  def handle_info(:measure, state) do
+    measure_and_schedule()
+    {:noreply, state}
+  end
+
+  defp measure_and_schedule() do
+    memory = :memsup.get_system_memory_data()
+    :ets.insert(@name, {:memory, %{total: memory[:total_memory], free: memory[:free_memory]}})
+    Process.send_after(self(), :measure, 15000)
+  end
+end

lib/livebook/utils.ex

@@ -3,8 +3,6 @@ defmodule Livebook.Utils do
 
   @type id :: binary()
 
-  @type system_memory :: %{total: non_neg_integer(), free: non_neg_integer()}
-
   @doc """
   Generates a random binary id.
   """
@@ -371,15 +369,6 @@ defmodule Livebook.Utils do
     |> Enum.join("\n")
   end
 
-  @doc """
-  Fetches the total and free memory of the system
-  """
-  @spec fetch_system_memory() :: system_memory()
-  def fetch_system_memory() do
-    memory = :memsup.get_system_memory_data()
-    %{total: memory[:total_memory], free: memory[:free_memory]}
-  end
-
   def format_bytes(bytes) when is_integer(bytes) do
     cond do
       bytes >= memory_unit(:TB) -> format_bytes(bytes, :TB)

lib/livebook_web/live/home_live/session_list_component.ex

@@ -1,7 +1,7 @@
 defmodule LivebookWeb.HomeLive.SessionListComponent do
   use LivebookWeb, :live_component
 
-  import Livebook.Utils, only: [format_bytes: 1, fetch_system_memory: 0]
+  import Livebook.Utils, only: [format_bytes: 1]
   import LivebookWeb.SessionHelpers, only: [uses_memory?: 1]
 
   @impl true
@@ -24,8 +24,11 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
     socket =
       socket
       |> assign(assigns)
-      |> assign(sessions: sessions, show_autosave_note?: show_autosave_note?)
+      |> assign(
-      |> assign(memory: memory_info(sessions))
+        sessions: sessions,
+        show_autosave_note?: show_autosave_note?,
+        memory: memory_info()
+      )
 
     {:ok, socket}
   end
@@ -38,15 +41,15 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
         <h2 class="uppercase font-semibold text-gray-500">
           Running sessions (<%= length(@sessions) %>)
         </h2>
-        <span class="tooltip top" data-tooltip={"This machine has #{format_bytes(@memory.system.total)}"}>
+        <span class="tooltip top" data-tooltip={"#{format_bytes(@memory.free)} available"}>
-        <div class="text-md text-gray-500 font-medium">
+          <div class="text-md text-gray-500 font-medium">
-          <span> <%= format_bytes(@memory.sessions) %> / <%= format_bytes(@memory.system.free) %></span>
+            <span> <%= format_bytes(@memory.used) %> / <%= format_bytes(@memory.total) %></span>
-            <div class="w-64 h-4 bg-gray-200">
+              <div class="w-64 h-4 bg-gray-200">
-            <div class="h-4 bg-blue-600"
+              <div class="h-4 bg-blue-600"
-            style={"width: #{@memory.percentage}%"}>
+              style={"width: #{@memory.percentage}%"}>
+              </div>
             </div>
           </div>
-        </div>
         </span>
         <.menu id="sessions-order-menu">
           <:toggle>
@@ -187,16 +190,11 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
     Enum.sort_by(sessions, &total_runtime_memory/1, :desc)
   end
 
-  defp memory_info(sessions) do
+  defp memory_info() do
-    sessions_memory =
+    %{free: free, total: total} = Livebook.SystemResources.memory()
-      sessions
+    used = total - free
-      |> Enum.map(&total_runtime_memory/1)
+    percentage = Float.round(used / total * 100, 2)
-      |> Enum.sum()
+    %{free: free, used: used, total: total, percentage: percentage}
-
-    system_memory = fetch_system_memory()
-    percentage = Float.round(sessions_memory / system_memory.free * 100, 2)
-
-    %{sessions: sessions_memory, system: system_memory, percentage: percentage}
   end
 
   defp total_runtime_memory(%{memory_usage: %{runtime: nil}}), do: 0

lib/livebook_web/live/session_live.ex

@@ -395,7 +395,7 @@ defmodule LivebookWeb.SessionLive do
           <.remix_icon icon="settings-3-line text-xl" />
         <% end %>
       </div>
-      <div class="flex flex-col mt-4 space-y-4">
+      <div class="flex flex-col mt-2 space-y-4">
         <%= if @data_view.runtime do %>
           <div class="flex flex-col space-y-3">
             <.labeled_text label="Type" text={runtime_type_label(@data_view.runtime)} />
@@ -430,41 +430,15 @@ defmodule LivebookWeb.SessionLive do
           </div>
         <% end %>
         <%= if uses_memory?(@session.memory_usage) do %>
-          <div class="py-6 flex flex-col justify-center relative overflow-hidden">
+          <.memory_info memory_usage={@session.memory_usage} />
-            <div class="mb-1 text-sm font-semibold text-gray-800 flex flex-row justify-between">
-              <span class="text-gray-500 uppercase">Memory:</span>
-              <div class="basis-3/4">
-                <span class="tooltip bottom"
-                data-tooltip={"This machine has #{format_bytes(@session.memory_usage.system.total)}"}>
-                  <span class="w-full text-right">
-                    <%= format_bytes(@session.memory_usage.runtime.total) %>
-                    /
-                    <%= format_bytes(@session.memory_usage.system.free) %>
-                  </span>
-                </span>
-              </div>
-            </div>
-            <div class="w-full h-6 flex flex-row gap-1">
-              <%= for {type, memory} <- runtime_memory(@session.memory_usage) do %>
-                <div class={"h-6 #{memory_color(type)}"} style={"width: #{memory.percentage}%"}></div>
-              <% end %>
-            </div>
-            <div class="flex flex-col py-4">
-              <%= for {type, memory} <- runtime_memory(@session.memory_usage) do %>
-                <div class="flex flex-row items-center">
-                  <span class={"w-4 h-4 mr-2 rounded #{memory_color(type)}"}></span>
-                  <span class="capitalize text-gray-700"><%= type %></span>
-                  <span class="text-gray-500 ml-auto"><%= memory.unit %></span>
-                </div>
-              <% end %>
-            </div>
-          </div>
         <% else %>
-          <div class="mb-1 text-sm font-semibold text-gray-800 py-4 flex flex-col">
+          <div class="mb-1 text-sm text-gray-800 py-6 flex flex-col">
-            <span class="w-full uppercase text-gray-500">Memory</span>
+            <span class="w-full uppercase font-semibold text-gray-500">Memory</span>
-            <%= format_bytes(@session.memory_usage.system.free) %>
+            <p class="py-1">
-            available out of
+              <%= format_bytes(@session.memory_usage.system.free) %>
-            <%= format_bytes(@session.memory_usage.system.total) %>
+              available out of
+              <%= format_bytes(@session.memory_usage.system.total) %>
+            </p>
           </div>
         <% end %>
       </div>
@@ -472,6 +446,38 @@ defmodule LivebookWeb.SessionLive do
     """
   end
 
+  defp memory_info(assigns) do
+    assigns = assign(assigns, :runtime_memory, runtime_memory(assigns.memory_usage))
+
+    ~H"""
+    <div class="py-6 flex flex-col justify-center">
+      <div class="mb-1 text-sm font-semibold text-gray-800 flex flex-row justify-between">
+        <span class="text-gray-500 uppercase">Memory</span>
+        <span class="text-right">
+          <%= format_bytes(@memory_usage.system.free) %> available
+        </span>
+      </div>
+      <div class="w-full h-8 flex flex-row py-1 gap-0.5">
+        <%= for {type, memory} <- @runtime_memory  do %>
+          <div class={"h-6 #{memory_color(type)}"} style={"width: #{memory.percentage}%"}></div>
+        <% end %>
+      </div>
+      <div class="flex flex-col py-1">
+        <%= for {type, memory} <- @runtime_memory do %>
+          <div class="flex flex-row items-center">
+            <span class={"w-4 h-4 mr-2 rounded #{memory_color(type)}"}></span>
+            <span class="capitalize text-gray-700"><%= type %></span>
+            <span class="text-gray-500 ml-auto"><%= memory.unit %></span>
+          </div>
+        <% end %>
+        <div class="flex rounded justify-center my-2 py-0.5 text-sm text-gray-800 bg-gray-200">
+          Total: <%= format_bytes(@memory_usage.runtime.total) %>
+        </div>
+      </div>
+    </div>
+    """
+  end
+
   defp runtime_type_label(%Runtime.ElixirStandalone{}), do: "Elixir standalone"
   defp runtime_type_label(%Runtime.MixStandalone{}), do: "Mix standalone"
   defp runtime_type_label(%Runtime.Attached{}), do: "Attached"
@@ -1525,11 +1531,11 @@ defmodule LivebookWeb.SessionLive do
     "Livebook - #{notebook_name}"
   end
 
-  defp memory_color(:atom), do: "bg-green-500"
+  defp memory_color(:atom), do: "bg-blue-500"
-  defp memory_color(:code), do: "bg-blue-700"
+  defp memory_color(:code), do: "bg-yellow-600"
-  defp memory_color(:processes), do: "bg-red-500"
+  defp memory_color(:processes), do: "bg-blue-700"
-  defp memory_color(:binary), do: "bg-blue-500"
+  defp memory_color(:binary), do: "bg-green-500"
-  defp memory_color(:ets), do: "bg-yellow-600"
+  defp memory_color(:ets), do: "bg-red-500"
   defp memory_color(:other), do: "bg-gray-400"
 
   defp runtime_memory(%{runtime: memory}) do

test/livebook/evaluator_test.exs

@@ -9,6 +9,12 @@ defmodule Livebook.EvaluatorTest do
     %{evaluator: evaluator, object_tracker: object_tracker}
   end
 
+  defmacrop metadata do
+    quote do
+      %{evaluation_time_ms: _, memory_usage: %{}}
+    end
+  end
+
   describe "evaluate_code/6" do
     test "given a valid code returns evaluation result", %{evaluator: evaluator} do
       code = """
@@ -19,12 +25,16 @@ defmodule Livebook.EvaluatorTest do
 
       Evaluator.evaluate_code(evaluator, self(), code, :code_1)
 
-      assert_receive {:evaluation_response, :code_1, {:ok, 3}, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, {:ok, 3}, metadata() = metadata}
+      assert metadata.evaluation_time_ms >= 0
+
+      assert %{atom: _, binary: _, code: _, ets: _, other: _, processes: _, total: _} =
+               metadata.memory_usage
     end
 
     test "given no prev_ref does not see previous evaluation context", %{evaluator: evaluator} do
       Evaluator.evaluate_code(evaluator, self(), "x = 1", :code_1)
-      assert_receive {:evaluation_response, :code_1, _, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, _, metadata()}
 
       ignore_warnings(fn ->
         Evaluator.evaluate_code(evaluator, self(), "x", :code_2)
@@ -33,23 +43,23 @@ defmodule Livebook.EvaluatorTest do
                         {:error, _kind,
                          %CompileError{
                            description: "undefined function x/0 (there is no such import)"
-                         }, _stacktrace}, %{evaluation_time_ms: _time_ms}}
+                         }, _stacktrace}, metadata()}
       end)
     end
 
     test "given prev_ref sees previous evaluation context", %{evaluator: evaluator} do
       Evaluator.evaluate_code(evaluator, self(), "x = 1", :code_1)
-      assert_receive {:evaluation_response, :code_1, _, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, _, metadata()}
 
       Evaluator.evaluate_code(evaluator, self(), "x", :code_2, :code_1)
 
-      assert_receive {:evaluation_response, :code_2, {:ok, 1}, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_2, {:ok, 1}, metadata()}
     end
 
     test "given invalid prev_ref just uses default context", %{evaluator: evaluator} do
       Evaluator.evaluate_code(evaluator, self(), ":hey", :code_1, :code_nonexistent)
 
-      assert_receive {:evaluation_response, :code_1, {:ok, :hey}, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, {:ok, :hey}, metadata()}
     end
 
     test "captures standard output and sends it to the caller", %{evaluator: evaluator} do
@@ -73,8 +83,7 @@ defmodule Livebook.EvaluatorTest do
       assert_receive {:evaluation_input, :code_1, reply_to, "input1"}
       send(reply_to, {:evaluation_input_reply, {:ok, :value}})
 
-      assert_receive {:evaluation_response, :code_1, {:ok, :value},
+      assert_receive {:evaluation_response, :code_1, {:ok, :value}, metadata()}
-                      %{evaluation_time_ms: _time_ms}}
     end
 
     test "returns error along with its kind and stacktrace", %{evaluator: evaluator} do
@@ -86,7 +95,7 @@ defmodule Livebook.EvaluatorTest do
 
       assert_receive {:evaluation_response, :code_1,
                       {:error, :error, %FunctionClauseError{},
-                       [{List, :first, _arity, _location}]}, %{evaluation_time_ms: _time_ms}}
+                       [{List, :first, _arity, _location}]}, metadata()}
     end
 
     test "in case of an error returns only the relevant part of stacktrace",
@@ -119,8 +128,7 @@ defmodule Livebook.EvaluatorTest do
 
         # Note: evaluating module definitions is relatively slow, so we use a higher wait timeout.
         assert_receive {:evaluation_response, :code_1,
-                        {:error, _kind, _error, ^expected_stacktrace},
+                        {:error, _kind, _error, ^expected_stacktrace}, metadata()},
-                        %{evaluation_time_ms: _time_ms}},
                        2_000
       end)
     end
@@ -140,15 +148,14 @@ defmodule Livebook.EvaluatorTest do
       """
 
       Evaluator.evaluate_code(evaluator, self(), code1, :code_1)
-      assert_receive {:evaluation_response, :code_1, {:ok, _}, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, {:ok, _}, metadata()}
 
       Evaluator.evaluate_code(evaluator, self(), code2, :code_2, :code_1)
 
-      assert_receive {:evaluation_response, :code_2, {:error, _, _, _},
+      assert_receive {:evaluation_response, :code_2, {:error, _, _, _}, metadata()}
-                      %{evaluation_time_ms: _time_ms}}
 
       Evaluator.evaluate_code(evaluator, self(), code3, :code_3, :code_2)
-      assert_receive {:evaluation_response, :code_3, {:ok, 4}, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_3, {:ok, 4}, metadata()}
     end
 
     test "given file option sets it in evaluation environment", %{evaluator: evaluator} do
@@ -159,18 +166,7 @@ defmodule Livebook.EvaluatorTest do
       opts = [file: "/path/dir/file"]
       Evaluator.evaluate_code(evaluator, self(), code, :code_1, nil, opts)
 
-      assert_receive {:evaluation_response, :code_1, {:ok, "/path/dir"},
+      assert_receive {:evaluation_response, :code_1, {:ok, "/path/dir"}, metadata()}
-                      %{evaluation_time_ms: _time_ms}}
-    end
-
-    test "given a :notify_to option to the evaluator", %{evaluator: evaluator} do
-      code = """
-      IO.puts("CatOps")
-      """
-
-      opts = [notify_to: self()]
-      Evaluator.evaluate_code(evaluator, self(), code, :code_1, nil, opts)
-      assert_receive {:evaluation_finished, :code_1}
     end
 
     test "kills widgets that that no evaluation points to", %{evaluator: evaluator} do
@@ -180,15 +176,13 @@ defmodule Livebook.EvaluatorTest do
 
       Evaluator.evaluate_code(evaluator, self(), spawn_widget_code(), :code_1)
 
-      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid1},
+      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid1}, metadata()}
-                      %{evaluation_time_ms: _time_ms}}
 
       ref = Process.monitor(widget_pid1)
 
       Evaluator.evaluate_code(evaluator, self(), spawn_widget_code(), :code_1)
 
-      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid2},
+      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid2}, metadata()}
-                      %{evaluation_time_ms: _time_ms}}
 
       assert_receive {:DOWN, ^ref, :process, ^widget_pid1, _reason}
 
@@ -206,8 +200,7 @@ defmodule Livebook.EvaluatorTest do
         :code_1
       )
 
-      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid1},
+      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid1}, metadata()}
-                      %{evaluation_time_ms: _time_ms}}
 
       refute Process.alive?(widget_pid1)
     end
@@ -216,7 +209,7 @@ defmodule Livebook.EvaluatorTest do
   describe "forget_evaluation/2" do
     test "invalidates the given reference", %{evaluator: evaluator} do
       Evaluator.evaluate_code(evaluator, self(), "x = 1", :code_1)
-      assert_receive {:evaluation_response, :code_1, _, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, _, metadata()}
 
       Evaluator.forget_evaluation(evaluator, :code_1)
 
@@ -227,15 +220,14 @@ defmodule Livebook.EvaluatorTest do
                         {:error, _kind,
                          %CompileError{
                            description: "undefined function x/0 (there is no such import)"
-                         }, _stacktrace}, %{evaluation_time_ms: _time_ms}}
+                         }, _stacktrace}, metadata()}
       end)
     end
 
     test "kills widgets that no evaluation points to", %{evaluator: evaluator} do
       Evaluator.evaluate_code(evaluator, self(), spawn_widget_code(), :code_1)
 
-      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid1},
+      assert_receive {:evaluation_response, :code_1, {:ok, widget_pid1}, metadata()}
-                      %{evaluation_time_ms: _time_ms}}
 
       ref = Process.monitor(widget_pid1)
       Evaluator.forget_evaluation(evaluator, :code_1)
@@ -260,7 +252,7 @@ defmodule Livebook.EvaluatorTest do
       """
 
       Evaluator.evaluate_code(evaluator, self(), code, :code_1)
-      assert_receive {:evaluation_response, :code_1, _, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, _, metadata()}
 
       request = {:completion, "num"}
       Evaluator.handle_intellisense(evaluator, self(), :ref, request, :code_1)
@@ -287,23 +279,23 @@ defmodule Livebook.EvaluatorTest do
     test "copies the given context and sets as the initial one",
          %{evaluator: evaluator, parent_evaluator: parent_evaluator} do
       Evaluator.evaluate_code(parent_evaluator, self(), "x = 1", :code_1)
-      assert_receive {:evaluation_response, :code_1, _, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, _, metadata()}
 
       Evaluator.initialize_from(evaluator, parent_evaluator, :code_1)
 
       Evaluator.evaluate_code(evaluator, self(), "x", :code_2)
-      assert_receive {:evaluation_response, :code_2, {:ok, 1}, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_2, {:ok, 1}, metadata()}
     end
 
     test "mirrors process dictionary of the given evaluator",
          %{evaluator: evaluator, parent_evaluator: parent_evaluator} do
       Evaluator.evaluate_code(parent_evaluator, self(), "Process.put(:data, 1)", :code_1)
-      assert_receive {:evaluation_response, :code_1, _, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_1, _, metadata()}
 
       Evaluator.initialize_from(evaluator, parent_evaluator, :code_1)
 
       Evaluator.evaluate_code(evaluator, self(), "Process.get(:data)", :code_2)
-      assert_receive {:evaluation_response, :code_2, {:ok, 1}, %{evaluation_time_ms: _time_ms}}
+      assert_receive {:evaluation_response, :code_2, {:ok, 1}, metadata()}
     end
   end