Track memory usage - visualization (#898)

* Utils to fetch and format system and node memory usage

* Order by memory

* Track memory on session

* Show memory usage on runtime sidebar

* Shows memory usage percentage on home

* Layout adjustments

* Sidebar design adjustments to match Figma

* Home design adjustments to show the memory information

* Move memory calculations to utils

* Shows disconnected notebooks as consuming 0mb on home

* Simplifies the data structure of memory usage

* Node memory tracker on runtime

* Clean up

* Renames node memory to runtime memory

* Standardizes the data structure of memory usage

* Sends evaluation_finished to the runtime to update the memory usage after an evaluation

* Fix: The evalutor does not notify when there is no notify_to option

* Adds a test with the notify_to option to the evaluator

* Documents the notify_to option

* Minor fixes on runtime and runtime_server

* Minor fixes on sessions

* Minor adjustments

* Updates docs and specs on Utils

* Minor adjustments on session_live

* Fix total memory used by sessions on home

* Put duplicated functions on helpers

* Better filter by memory

* Fix the tooltip text for memory information on sidebar

* Minor alignment adjustment on home

lib/livebook/evaluator.ex

@@ -84,6 +84,9 @@ 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
@@ -232,6 +235,7 @@ 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
@@ -255,6 +259,7 @@ defmodule Livebook.Evaluator do
     output = state.formatter.format_response(response)
     metadata = %{evaluation_time_ms: evaluation_time_ms}
     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

@@ -115,6 +115,22 @@ defprotocol Livebook.Runtime do
           code: String.t()
         }
 
+  @typedoc """
+  The runtime memory usage for each type in bytes.
+
+  The runtime may periodically send messages of type {:memory_usage, runtime_memory()}
+  """
+  @type runtime_memory :: %{
+          atom: non_neg_integer(),
+          binary: non_neg_integer(),
+          code: non_neg_integer(),
+          ets: non_neg_integer(),
+          other: non_neg_integer(),
+          processes: non_neg_integer(),
+          system: non_neg_integer(),
+          total: non_neg_integer()
+        }
+
   @doc """
   Sets the caller as runtime owner.
 

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -20,6 +20,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   alias Livebook.Runtime.ErlDist
 
   @await_owner_timeout 5_000
+  @memory_usage_interval 15_000
 
   @doc """
   Starts the manager.
@@ -142,7 +143,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
        evaluators: %{},
        evaluator_supervisor: evaluator_supervisor,
        task_supervisor: task_supervisor,
-       object_tracker: object_tracker
+       object_tracker: object_tracker,
+       memory_timer_ref: nil
      }}
   end
 
@@ -177,12 +179,23 @@ 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)
+    {:noreply, schedule_memory_usage(state)}
+  end
+
   def handle_info(_message, state), do: {:noreply, state}
 
   @impl true
   def handle_cast({:set_owner, owner}, state) do
     Process.monitor(owner)
-
+    send(self(), :memory_usage)
     {:noreply, %{state | owner: owner}}
   end
 
@@ -191,6 +204,7 @@ 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
@@ -297,4 +311,22 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
         state
     end
   end
+
+  defp schedule_memory_usage(state) 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

@@ -46,7 +46,7 @@ defmodule Livebook.Session do
   # The struct holds the basic session information that we track
   # and pass around. The notebook and evaluation state is kept
   # within the process state.
-  defstruct [:id, :pid, :origin, :notebook_name, :file, :images_dir, :created_at]
+  defstruct [:id, :pid, :origin, :notebook_name, :file, :images_dir, :created_at, :memory_usage]
 
   use GenServer, restart: :temporary
 
@@ -55,6 +55,8 @@ 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(),
@@ -62,7 +64,8 @@ defmodule Livebook.Session do
           notebook_name: String.t(),
           file: FileSystem.File.t() | nil,
           images_dir: FileSystem.File.t(),
-          created_at: DateTime.t()
+          created_at: DateTime.t(),
+          memory_usage: memory_usage()
         }
 
   @type state :: %{
@@ -72,9 +75,17 @@ defmodule Livebook.Session do
           runtime_monitor_ref: reference() | nil,
           autosave_timer_ref: reference() | nil,
           save_task_pid: pid() | nil,
-          saved_default_file: FileSystem.File.t() | 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()
+          }
+
   @typedoc """
   An id assigned to every running session process.
   """
@@ -447,7 +458,7 @@ defmodule Livebook.Session do
              do: dump_images(state, images),
              else: :ok
            ) do
-      state = schedule_autosave(state)
+      state = state |> schedule_autosave() |> schedule_system_memory_update()
       {:ok, state}
     else
       {:error, error} ->
@@ -467,7 +478,9 @@ defmodule Livebook.Session do
         autosave_timer_ref: nil,
         autosave_path: opts[:autosave_path],
         save_task_pid: nil,
-        saved_default_file: nil
+        saved_default_file: nil,
+        system_memory_timer_ref: nil,
+        memory_usage: %{runtime: nil, system: Utils.fetch_system_memory()}
       }
 
       {:ok, state}
@@ -508,6 +521,11 @@ 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}
@@ -814,6 +832,19 @@ 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)
+    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}
 
   @impl true
@@ -832,7 +863,8 @@ defmodule Livebook.Session do
       notebook_name: state.data.notebook.name,
       file: state.data.file,
       images_dir: images_dir_from_state(state),
-      created_at: state.created_at
+      created_at: state.created_at,
+      memory_usage: state.memory_usage
     }
   end
 
@@ -981,6 +1013,16 @@ defmodule Livebook.Session do
     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)
+      |> update_system_memory_usage()
+      |> schedule_system_memory_update()
+    end
+  end
+
   defp after_operation(state, prev_state, {:set_file, _pid, _file}) do
     prev_images_dir = images_dir_from_state(prev_state)
 
@@ -1268,4 +1310,10 @@ 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/utils.ex

@@ -3,6 +3,8 @@ defmodule Livebook.Utils do
 
   @type id :: binary()
 
+  @type system_memory :: %{total: non_neg_integer(), free: non_neg_integer()}
+
   @doc """
   Generates a random binary id.
   """
@@ -368,4 +370,35 @@ defmodule Livebook.Utils do
     end)
     |> 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)
+      bytes >= memory_unit(:GB) -> format_bytes(bytes, :GB)
+      bytes >= memory_unit(:MB) -> format_bytes(bytes, :MB)
+      bytes >= memory_unit(:KB) -> format_bytes(bytes, :KB)
+      true -> format_bytes(bytes, :B)
+    end
+  end
+
+  defp format_bytes(bytes, :B) when is_integer(bytes), do: "#{bytes} B"
+
+  defp format_bytes(bytes, unit) when is_integer(bytes) do
+    value = bytes / memory_unit(unit)
+    "#{:erlang.float_to_binary(value, decimals: 1)} #{unit}"
+  end
+
+  defp memory_unit(:TB), do: 1024 * 1024 * 1024 * 1024
+  defp memory_unit(:GB), do: 1024 * 1024 * 1024
+  defp memory_unit(:MB), do: 1024 * 1024
+  defp memory_unit(:KB), do: 1024
 end

lib/livebook_web/live/home_live/session_list_component.ex

@@ -1,6 +1,9 @@
 defmodule LivebookWeb.HomeLive.SessionListComponent do
   use LivebookWeb, :live_component
 
+  import Livebook.Utils, only: [format_bytes: 1, fetch_system_memory: 0]
+  import LivebookWeb.SessionHelpers, only: [uses_memory?: 1]
+
   @impl true
   def mount(socket) do
     {:ok, assign(socket, order_by: "date")}
@@ -22,6 +25,7 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
       socket
       |> assign(assigns)
       |> assign(sessions: sessions, show_autosave_note?: show_autosave_note?)
+      |> assign(memory: memory_info(sessions))
 
     {:ok, socket}
   end
@@ -30,10 +34,20 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
   def render(assigns) do
     ~H"""
     <div>
-      <div class="flex items-center justify-between">
-        <h2 class="mb-4 uppercase font-semibold text-gray-500">
+      <div class="mb-4 flex items-end justify-between">
+        <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)}"}>
+        <div class="text-md text-gray-500 font-medium">
+          <span> <%= format_bytes(@memory.sessions) %> / <%= format_bytes(@memory.system.free) %></span>
+            <div class="w-64 h-4 bg-gray-200">
+            <div class="h-4 bg-blue-600"
+            style={"width: #{@memory.percentage}%"}>
+            </div>
+          </div>
+        </div>
+        </span>
         <.menu id="sessions-order-menu">
           <:toggle>
             <button class="button-base button-outlined-gray px-4 py-1">
@@ -42,7 +56,7 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
             </button>
           </:toggle>
           <:content>
-            <%= for order_by <- ["date", "title"] do %>
+            <%= for order_by <- ["date", "title", "memory"] do %>
               <button class={"menu-item #{if order_by == @order_by, do: "text-gray-900", else: "text-gray-500"}"}
                 role="menuitem"
                 phx-click={JS.push("set_order", value: %{order_by: order_by}, target: @myself)}>
@@ -94,7 +108,14 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
             <div class="text-gray-600 text-sm">
               <%= if session.file, do: session.file.path, else: "No file" %>
             </div>
-            <div class="mt-2 text-gray-600 text-sm">
+            <div class="mt-2 text-gray-600 text-sm flex flex-row items-center">
+            <%= if uses_memory?(session.memory_usage) do %>
+              <div class="h-3 w-3 mr-1 rounded-full bg-green-500"></div>
+              <span class="pr-4"><%= format_bytes(session.memory_usage.runtime.total) %></span>
+            <% else %>
+              <div class="h-3 w-3 mr-1 rounded-full bg-gray-300"></div>
+              <span class="pr-4">0 MB</span>
+            <% end %>
               Created <%= format_creation_date(session.created_at) %>
             </div>
           </div>
@@ -146,9 +167,11 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
 
   defp order_by_label("date"), do: "Date"
   defp order_by_label("title"), do: "Title"
+  defp order_by_label("memory"), do: "Memory"
 
   defp order_by_icon("date"), do: "calendar-2-line"
   defp order_by_icon("title"), do: "text"
+  defp order_by_icon("memory"), do: "cpu-line"
 
   defp sort_sessions(sessions, "date") do
     Enum.sort_by(sessions, & &1.created_at, {:desc, DateTime})
@@ -159,4 +182,23 @@ defmodule LivebookWeb.HomeLive.SessionListComponent do
       {session.notebook_name, -DateTime.to_unix(session.created_at)}
     end)
   end
+
+  defp sort_sessions(sessions, "memory") do
+    Enum.sort_by(sessions, &total_runtime_memory/1, :desc)
+  end
+
+  defp memory_info(sessions) do
+    sessions_memory =
+      sessions
+      |> Enum.map(&total_runtime_memory/1)
+      |> Enum.sum()
+
+    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
+  defp total_runtime_memory(%{memory_usage: %{runtime: %{total: total}}}), do: total
 end

lib/livebook_web/live/session_helpers.ex

@@ -51,4 +51,7 @@ defmodule LivebookWeb.SessionHelpers do
 
     put_flash(socket, :warning, flash)
   end
+
+  def uses_memory?(%{runtime: %{total: total}}) when total > 0, do: true
+  def uses_memory?(_), do: false
 end

lib/livebook_web/live/session_live.ex

@@ -3,7 +3,7 @@ defmodule LivebookWeb.SessionLive do
 
   import LivebookWeb.UserHelpers
   import LivebookWeb.SessionHelpers
-  import Livebook.Utils, only: [access_by_id: 1]
+  import Livebook.Utils, only: [access_by_id: 1, format_bytes: 1]
 
   alias LivebookWeb.SidebarHelpers
   alias Livebook.{Sessions, Session, Delta, Notebook, Runtime, LiveMarkdown}
@@ -429,6 +429,44 @@ defmodule LivebookWeb.SessionLive do
             <% end %>
           </div>
         <% end %>
+        <%= if uses_memory?(@session.memory_usage) do %>
+          <div class="py-6 flex flex-col justify-center relative overflow-hidden">
+            <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">
+            <span class="w-full uppercase text-gray-500">Memory</span>
+            <%= format_bytes(@session.memory_usage.system.free) %>
+            available out of
+            <%= format_bytes(@session.memory_usage.system.total) %>
+          </div>
+        <% end %>
       </div>
     </div>
     """
@@ -1486,4 +1524,24 @@ defmodule LivebookWeb.SessionLive do
   defp get_page_title(notebook_name) do
     "Livebook - #{notebook_name}"
   end
+
+  defp memory_color(:atom), do: "bg-green-500"
+  defp memory_color(:code), do: "bg-blue-700"
+  defp memory_color(:processes), do: "bg-red-500"
+  defp memory_color(:binary), do: "bg-blue-500"
+  defp memory_color(:ets), do: "bg-yellow-600"
+  defp memory_color(:other), do: "bg-gray-400"
+
+  defp runtime_memory(%{runtime: memory}) do
+    memory
+    |> Map.drop([:total, :system])
+    |> Enum.map(fn {type, bytes} ->
+      {type,
+       %{
+         unit: format_bytes(bytes),
+         percentage: Float.round(bytes / memory.total * 100, 2),
+         value: bytes
+       }}
+    end)
+  end
 end

test/livebook/evaluator_test.exs

@@ -163,6 +163,16 @@ defmodule Livebook.EvaluatorTest do
                       %{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
       # Evaluate the code twice, each time a new widget is spawned.
       # The evaluation reference is the same, so the second one overrides