Add embedded runtime for evaluating code in the Livebook VM (#266)

* Add embedded runtime for evaluating code in the Livebook VM

* Update lib/livebook_web/live/session_live/embedded_live.ex

Co-authored-by: José Valim <jose.valim@dashbit.co>

* Use standard error proxy globally in the Livebook node

* Add configuration env variable for setting the default runtime

* Increase evaluation response assertion timeouts

Co-authored-by: José Valim <jose.valim@dashbit.co>

README.md

@@ -96,6 +96,11 @@ The following environment variables configure Livebook:
   * LIVEBOOK_COOKIE - sets the cookie for running Livebook in a cluster.
     Defaults to a random string that is generated on boot.
 
+  * LIVEBOOK_DEFAULT_RUNTIME - sets the runtime type that is used
+    by default when none is started explicitly for the given notebook.
+    Must be either "standalone" (Elixir standalone) or "embedded" (Embedded).
+    Defaults to "standalone".
+
   * LIVEBOOK_IP - sets the ip address to start the web application on. Must be a valid IPv4 or IPv6 address.
 
   * LIVEBOOK_PASSWORD - sets a password that must be used to access Livebook. Must be at least 12 characters. Defaults to token authentication.

assets/css/components.css

@@ -139,3 +139,9 @@
 .menu__item {
   @apply flex space-x-3 px-5 py-2 items-center hover:bg-gray-50 whitespace-nowrap;
 }
+
+/* Boxes */
+
+.error-box {
+  @apply mb-3 rounded-lg px-4 py-2 bg-red-100 text-red-400 font-medium;
+}

config/config.exs

@@ -17,6 +17,14 @@ config :phoenix, :json_library, Jason
 # Sets the default authentication mode to token
 config :livebook, :authentication_mode, :token
 
+# Sets the default runtime to ElixirStandalone.
+# This is the desired default most of the time,
+# but in some specific use cases you may want
+# to configure that to the Embedded runtime instead.
+# Also make sure the configured runtime has
+# a synchronous `init/0` method.
+config :livebook, :default_runtime, Livebook.Runtime.ElixirStandalone
+
 # Import environment specific config. This must remain at the bottom
 # of this file so it overrides the configuration defined above.
 import_config "#{Mix.env()}.exs"

config/runtime.exs

@@ -25,3 +25,8 @@ end
 config :livebook,
        :cookie,
        Livebook.Config.cookie!("LIVEBOOK_COOKIE") || Livebook.Utils.random_cookie()
+
+config :livebook,
+       :default_runtime,
+       Livebook.Config.default_runtime!("LIVEBOOK_DEFAULT_RUNTIME") ||
+         Livebook.Runtime.ElixirStandalone

config/test.exs

@@ -12,6 +12,11 @@ config :logger, level: :warn
 # Disable authentication mode during test
 config :livebook, :authentication_mode, :disabled
 
+# Use the embedded runtime in tests by default, so they
+# are cheaper to run. Other runtimes can be tested by starting
+# and setting them explicitly
+config :livebook, :default_runtime, Livebook.Runtime.Embedded
+
 # Use longnames when running tests in CI, so that no host resolution is required,
 # see https://github.com/elixir-nx/livebook/pull/173#issuecomment-819468549
 if System.get_env("CI") == "true" do

lib/livebook/application.ex

@@ -9,11 +9,18 @@ defmodule Livebook.Application do
     ensure_distribution!()
     set_cookie()
 
+    # We register our own :standard_error below
+    Process.unregister(:standard_error)
+
     children = [
       # Start the Telemetry supervisor
       LivebookWeb.Telemetry,
       # Start the PubSub system
       {Phoenix.PubSub, name: Livebook.PubSub},
+      # Start the our own :standard_error handler (standard error -> group leader)
+      # This way we can run multiple embedded runtimes without worrying
+      # about restoring :standard_error to a valid process when terminating
+      {Livebook.Runtime.ErlDist.IOForwardGL, name: :standard_error},
       # Start the supervisor dynamically managing sessions
       Livebook.SessionSupervisor,
       # Start the server responsible for associating files with sessions

lib/livebook/config.ex

@@ -15,6 +15,14 @@ defmodule Livebook.Config do
     end
   end
 
+  @doc """
+  Returns the runtime module to be used by default.
+  """
+  @spec default_runtime() :: Livebook.Runtime
+  def default_runtime() do
+    Application.fetch_env!(:livebook, :default_runtime)
+  end
+
   @doc """
   Returns the authentication mode.
   """
@@ -128,6 +136,26 @@ defmodule Livebook.Config do
     end
   end
 
+  @doc """
+  Parses and validates default runtime from env.
+  """
+  def default_runtime!(env) do
+    if runtime = System.get_env(env) do
+      case runtime do
+        "standalone" ->
+          Livebook.Runtime.ElixirStandalone
+
+        "embedded" ->
+          Livebook.Runtime.Embedded
+
+        other ->
+          abort!(
+            ~s{expected #{env} to be either "standalone" or "embedded", got: #{inspect(other)}}
+          )
+      end
+    end
+  end
+
   @doc """
   Aborts booting due to a configuration error.
   """

lib/livebook/runtime/embedded.ex

@@ -0,0 +1,98 @@
+defmodule Livebook.Runtime.Embedded do
+  @moduledoc false
+
+  # A runtime backed by the same node Livebook is running in.
+  #
+  # This runtime is reserved for specific use cases,
+  # where there is no option of starting a separate
+  # Elixir runtime.
+
+  defstruct [:node, :manager_pid]
+
+  @type t :: %__MODULE__{
+          node: node(),
+          manager_pid: pid()
+        }
+
+  alias Livebook.Runtime.ErlDist
+
+  @doc """
+  Initializes new runtime by starting the necessary
+  processes within the current node.
+  """
+  @spec init() :: {:ok, t()} | {:error, :failure}
+  def init() do
+    # As we run in the Livebook node, all the necessary modules
+    # are in place, so we just start the manager process.
+    # We make it anonymous, so that multiple embedded runtimes
+    # can be started (for different notebooks).
+    # We also disable cleanup, as we don't want to unload any
+    # modules or revert the configuration (because other runtimes
+    # may rely on it). If someone uses embedded runtimes,
+    # this cleanup is not particularly important anyway.
+    # We tell manager to not override :standard_error,
+    # as we already do it for the Livebook application globally
+    # (see Livebook.Application.start/2).
+    case ErlDist.Manager.start(
+           anonymous: true,
+           cleanup_on_termination: false,
+           register_standard_error_proxy: false
+         ) do
+      {:ok, pid} ->
+        {:ok, %__MODULE__{node: node(), manager_pid: pid}}
+
+      _ ->
+        {:error, :failure}
+    end
+  end
+end
+
+defimpl Livebook.Runtime, for: Livebook.Runtime.Embedded do
+  alias Livebook.Runtime.ErlDist
+
+  def connect(runtime) do
+    ErlDist.Manager.set_owner(runtime.manager_pid, self())
+    Process.monitor(runtime.manager_pid)
+  end
+
+  def disconnect(runtime) do
+    ErlDist.Manager.stop(runtime.manager_pid)
+  end
+
+  def evaluate_code(
+        runtime,
+        code,
+        container_ref,
+        evaluation_ref,
+        prev_evaluation_ref,
+        opts \\ []
+      ) do
+    ErlDist.Manager.evaluate_code(
+      runtime.manager_pid,
+      code,
+      container_ref,
+      evaluation_ref,
+      prev_evaluation_ref,
+      opts
+    )
+  end
+
+  def forget_evaluation(runtime, container_ref, evaluation_ref) do
+    ErlDist.Manager.forget_evaluation(runtime.manager_pid, container_ref, evaluation_ref)
+  end
+
+  def drop_container(runtime, container_ref) do
+    ErlDist.Manager.drop_container(runtime.manager_pid, container_ref)
+  end
+
+  def request_completion_items(runtime, send_to, ref, hint, container_ref, evaluation_ref) do
+    ErlDist.Manager.request_completion_items(
+      runtime.manager_pid,
+      send_to,
+      ref,
+      hint,
+      container_ref,
+      evaluation_ref
+    )
+  end
+end

lib/livebook/runtime/erl_dist/io_forward_gl.ex

@@ -15,9 +15,18 @@ defmodule Livebook.Runtime.ErlDist.IOForwardGL do
 
   ## API
 
+  @doc """
+  Starts the IO device.
+
+  ## Options
+
+    * `:name` - the name to regsiter the process under. Optional.
+  """
   @spec start_link() :: GenServer.on_start()
   def start_link(opts \\ []) do
-    GenServer.start_link(__MODULE__, opts)
+    {gen_opts, opts} = Keyword.split(opts, [:name])
+
+    GenServer.start_link(__MODULE__, opts, gen_opts)
   end
 
   ## Callbacks

lib/livebook/runtime/erl_dist/manager.ex

@@ -20,11 +20,34 @@ defmodule Livebook.Runtime.ErlDist.Manager do
   @doc """
   Starts the manager.
 
-  Note: make sure to `set_owner` within `@await_owner_timeout`
+  Note: make sure to call `set_owner` within `@await_owner_timeout`
   or the manager assumes it's not needed and terminates.
+
+  ## Options
+
+    * `:anonymous` - configures whether manager should
+      be registered under a global name or not.
+      In most cases we enforce a single manager per node
+      and identify it by a name, but this can be opted-out
+      from using this option. Defaults to `false`.
+
+    * `:cleanup_on_termination` - configures whether
+      manager should cleanup any global configuration
+      it altered and unload Livebook-specific modules
+      from the node. Defaults to `true`.
+
+    * `:register_standard_error_proxy` - configures whether
+      manager should start an IOForwardGL process and set
+      it as the group leader. Defaults to `true`.
   """
   def start(opts \\ []) do
-    GenServer.start(__MODULE__, opts, name: @name)
+    {anonymous?, opts} = Keyword.pop(opts, :anonymous, false)
+
+    gen_opts = [
+      name: if(anonymous?, do: nil, else: @name)
+    ]
+
+    GenServer.start(__MODULE__, opts, gen_opts)
   end
 
   @doc """
@@ -34,9 +57,9 @@ defmodule Livebook.Runtime.ErlDist.Manager do
   the manager also terminates. All the evaluation results are
   send directly to the owner.
   """
-  @spec set_owner(node(), pid()) :: :ok
-  def set_owner(node, owner) do
-    GenServer.cast({@name, node}, {:set_owner, owner})
+  @spec set_owner(node() | pid(), pid()) :: :ok
+  def set_owner(node_or_pid, owner) do
+    GenServer.cast(server(node_or_pid), {:set_owner, owner})
   end
 
   @doc """
@@ -50,16 +73,23 @@ defmodule Livebook.Runtime.ErlDist.Manager do
   See `Evaluator` for more details.
   """
   @spec evaluate_code(
-          node(),
+          node() | pid(),
           String.t(),
           Evaluator.ref(),
           Evaluator.ref(),
           Evaluator.ref() | nil,
           keyword()
         ) :: :ok
-  def evaluate_code(node, code, container_ref, evaluation_ref, prev_evaluation_ref, opts \\ []) do
+  def evaluate_code(
+        node_or_pid,
+        code,
+        container_ref,
+        evaluation_ref,
+        prev_evaluation_ref,
+        opts \\ []
+      ) do
     GenServer.cast(
-      {@name, node},
+      server(node_or_pid),
       {:evaluate_code, code, container_ref, evaluation_ref, prev_evaluation_ref, opts}
     )
   end
@@ -69,17 +99,17 @@ defmodule Livebook.Runtime.ErlDist.Manager do
 
   See `Evaluator` for more details.
   """
-  @spec forget_evaluation(node(), Evaluator.ref(), Evaluator.ref()) :: :ok
-  def forget_evaluation(node, container_ref, evaluation_ref) do
-    GenServer.cast({@name, node}, {:forget_evaluation, container_ref, evaluation_ref})
+  @spec forget_evaluation(node() | pid(), Evaluator.ref(), Evaluator.ref()) :: :ok
+  def forget_evaluation(node_or_pid, container_ref, evaluation_ref) do
+    GenServer.cast(server(node_or_pid), {:forget_evaluation, container_ref, evaluation_ref})
   end
 
   @doc """
   Terminates the `Evaluator` process belonging to the given container.
   """
-  @spec drop_container(node(), Evaluator.ref()) :: :ok
-  def drop_container(node, container_ref) do
-    GenServer.cast({@name, node}, {:drop_container, container_ref})
+  @spec drop_container(node() | pid(), Evaluator.ref()) :: :ok
+  def drop_container(node_or_pid, container_ref) do
+    GenServer.cast(server(node_or_pid), {:drop_container, container_ref})
   end
 
   @doc """
@@ -93,16 +123,16 @@ defmodule Livebook.Runtime.ErlDist.Manager do
   See `Livebook.Runtime` for more details.
   """
   @spec request_completion_items(
-          node(),
+          node() | pid(),
           pid(),
           term(),
           String.t(),
           Evaluator.ref(),
           Evaluator.ref()
         ) :: :ok
-  def request_completion_items(node, send_to, ref, hint, container_ref, evaluation_ref) do
+  def request_completion_items(node_or_pid, send_to, ref, hint, container_ref, evaluation_ref) do
     GenServer.cast(
-      {@name, node},
+      server(node_or_pid),
       {:request_completion_items, send_to, ref, hint, container_ref, evaluation_ref}
     )
   end
@@ -112,13 +142,19 @@ defmodule Livebook.Runtime.ErlDist.Manager do
 
   This results in all Livebook-related modules being unloaded from this node.
   """
-  @spec stop(node()) :: :ok
-  def stop(node) do
-    GenServer.stop({@name, node})
+  @spec stop(node() | pid()) :: :ok
+  def stop(node_or_pid) do
+    GenServer.stop(server(node_or_pid))
   end
 
+  defp server(pid) when is_pid(pid), do: pid
+  defp server(node) when is_atom(node), do: {@name, node}
+
   @impl true
-  def init(_opts) do
+  def init(opts) do
+    cleanup_on_termination = Keyword.get(opts, :cleanup_on_termination, true)
+    register_standard_error_proxy = Keyword.get(opts, :register_standard_error_proxy, true)
+
     Process.send_after(self(), :check_owner, @await_owner_timeout)
 
     ## Initialize the node
@@ -126,21 +162,28 @@ defmodule Livebook.Runtime.ErlDist.Manager do
     Process.flag(:trap_exit, true)
 
     {:ok, evaluator_supervisor} = ErlDist.EvaluatorSupervisor.start_link()
-    {:ok, io_forward_gl_pid} = ErlDist.IOForwardGL.start_link()
     {:ok, completion_supervisor} = Task.Supervisor.start_link()
 
+    # Register our own standard error IO device that proxies
+    # to sender's group leader.
+
+    original_standard_error = Process.whereis(:standard_error)
+
+    if register_standard_error_proxy do
+      {:ok, io_forward_gl_pid} = ErlDist.IOForwardGL.start_link()
+
+      Process.unregister(:standard_error)
+      Process.register(io_forward_gl_pid, :standard_error)
+    end
+
     # Set `ignore_module_conflict` only for the Manager lifetime.
     initial_ignore_module_conflict = Code.compiler_options()[:ignore_module_conflict]
     Code.compiler_options(ignore_module_conflict: true)
 
-    # Register our own standard error IO devices that proxies
-    # to sender's group leader.
-    original_standard_error = Process.whereis(:standard_error)
-    Process.unregister(:standard_error)
-    Process.register(io_forward_gl_pid, :standard_error)
-
     {:ok,
      %{
+       cleanup_on_termination: cleanup_on_termination,
+       register_standard_error_proxy: register_standard_error_proxy,
        owner: nil,
        evaluators: %{},
        evaluator_supervisor: evaluator_supervisor,
@@ -152,12 +195,16 @@ defmodule Livebook.Runtime.ErlDist.Manager do
 
   @impl true
   def terminate(_reason, state) do
+    if state.cleanup_on_termination do
       Code.compiler_options(ignore_module_conflict: state.initial_ignore_module_conflict)
 
+      if state.register_standard_error_proxy do
         Process.unregister(:standard_error)
         Process.register(state.original_standard_error, :standard_error)
+      end
 
       ErlDist.unload_required_modules()
+    end
 
     :ok
   end

lib/livebook/session.ex

@@ -692,7 +692,7 @@ defmodule Livebook.Session do
   # Checks if a runtime already set, and if that's not the case
   # starts a new standalone one.
   defp ensure_runtime(%{data: %{runtime: nil}} = state) do
-    with {:ok, runtime} <- Runtime.ElixirStandalone.init() do
+    with {:ok, runtime} <- Livebook.Config.default_runtime().init() do
       runtime_monitor_ref = Runtime.connect(runtime)
 
       {:ok,

lib/livebook_web/live/home_live/import_url_component.ex

@@ -13,7 +13,7 @@ defmodule LivebookWeb.HomeLive.ImportUrlComponent do
     ~L"""
     <div class="flex-col space-y-5">
       <%= if @error_message do %>
-        <div class="mb-3 rounded-lg px-4 py-2 bg-red-100 text-red-400 font-medium">
+        <div class="error-box">
           <%= @error_message %>
         </div>
       <% end %>

lib/livebook_web/live/session_live/attached_live.ex

@@ -18,7 +18,7 @@ defmodule LivebookWeb.SessionLive.AttachedLive do
     ~L"""
     <div class="flex-col space-y-5">
       <%= if @error_message do %>
-        <div class="mb-3 rounded-lg px-4 py-2 bg-red-100 text-red-400 font-medium">
+        <div class="error-box">
           <%= @error_message %>
         </div>
       <% end %>

lib/livebook_web/live/session_live/cell_component.ex

@@ -245,7 +245,9 @@ defmodule LivebookWeb.SessionLive.CellComponent do
     ~L"""
     <div id="<%= @id %>" phx-hook="VirtualizedLines" data-max-height="300" data-follow="true">
       <div data-template class="hidden"><%= for line <- @lines do %><div><%= line %></div><% end %></div>
-      <div data-content phx-update="ignore" class="overflow-auto whitespace-pre text-gray-500 tiny-scrollbar"></div>
+      <div data-content class="overflow-auto whitespace-pre text-gray-500 tiny-scrollbar"
+        id="<%= @id %>-content"
+        phx-update="ignore"></div>
     </div>
     """
   end
@@ -257,7 +259,9 @@ defmodule LivebookWeb.SessionLive.CellComponent do
     ~L"""
     <div id="<%= @id %>" phx-hook="VirtualizedLines" data-max-height="300" data-follow="false">
       <div data-template class="hidden"><%= for line <- @lines do %><div><%= line %></div><% end %></div>
-      <div data-content phx-update="ignore" class="overflow-auto whitespace-pre text-gray-500 tiny-scrollbar"></div>
+      <div data-content class="overflow-auto whitespace-pre text-gray-500 tiny-scrollbar"
+        id="<%= @id %>-content"
+        phx-update="ignore"></div>
     </div>
     """
   end

lib/livebook_web/live/session_live/cell_upload_component.ex

@@ -16,7 +16,7 @@ defmodule LivebookWeb.SessionLive.CellUploadComponent do
         Insert image
       </h3>
       <%= if @uploads.cell_image.errors != [] do %>
-        <div class="mb-3 rounded-lg px-4 py-2 bg-red-100 text-red-400 font-medium">
+        <div class="error-box">
           Invalid image file. The image must be either GIF, JPEG, or PNG and cannot exceed 5MB in size.
         </div>
       <% end %>

lib/livebook_web/live/session_live/elixir_standalone_live.ex

@@ -9,13 +9,19 @@ defmodule LivebookWeb.SessionLive.ElixirStandaloneLive do
       Phoenix.PubSub.subscribe(Livebook.PubSub, "sessions:#{session_id}")
     end
 
-    {:ok, assign(socket, session_id: session_id, output: nil, current_runtime: current_runtime)}
+    {:ok,
+     assign(socket, session_id: session_id, current_runtime: current_runtime, error_message: nil)}
   end
 
   @impl true
   def render(assigns) do
     ~L"""
     <div class="flex-col space-y-5">
+      <%= if @error_message do %>
+        <div class="error-box">
+          <%= @error_message %>
+        </div>
+      <% end %>
       <p class="text-gray-700">
         Start a new local node to handle code evaluation.
         This is the default runtime and is started automatically
@@ -24,24 +30,23 @@ defmodule LivebookWeb.SessionLive.ElixirStandaloneLive do
       <button class="button button-blue" phx-click="init">
         <%= if(matching_runtime?(@current_runtime), do: "Reconnect", else: "Connect") %>
       </button>
-      <%= if @output do %>
-        <div class="markdown max-h-20 overflow-y-auto tiny-scrollbar">
-          <pre><code><%= @output %></code></pre>
-        </div>
-      <% end %>
     </div>
     """
   end
 
   defp matching_runtime?(%Runtime.ElixirStandalone{}), do: true
-
   defp matching_runtime?(_runtime), do: false
 
   @impl true
   def handle_event("init", _params, socket) do
-    {:ok, runtime} = Runtime.ElixirStandalone.init()
+    case Runtime.ElixirStandalone.init() do
+      {:ok, runtime} ->
         Session.connect_runtime(socket.assigns.session_id, runtime)
-    {:noreply, socket}
+        {:noreply, assign(socket, error_message: nil)}
+
+      {:error, message} ->
+        {:noreply, assign(socket, error_message: message)}
+    end
   end
 
   @impl true

lib/livebook_web/live/session_live/embedded_live.ex

@@ -0,0 +1,41 @@
+defmodule LivebookWeb.SessionLive.EmbeddedLive do
+  use LivebookWeb, :live_view
+
+  alias Livebook.{Session, Runtime}
+
+  @impl true
+  def mount(_params, %{"session_id" => session_id}, socket) do
+    {:ok, assign(socket, session_id: session_id)}
+  end
+
+  @impl true
+  def render(assigns) do
+    ~L"""
+    <div class="flex-col space-y-5">
+      <p class="text-gray-700">
+        Run the notebook code within the Livebook node itself.
+        This is reserved for specific cases where there is no option
+        of starting a separate Elixir runtime (for example, on embedded
+        devices or cases where the amount of memory available is
+        limited). Prefer the "Elixir standalone" runtime whenever possible.
+      </p>
+      <p class="text-gray-700">
+        <span class="font-semibold">Warning:</span>
+        any module that you define will be defined globally until
+        you restart Livebook. Furthermore, code in one notebook
+        may interfere with code from another notebook.
+      </p>
+      <button class="button button-blue" phx-click="init">
+        Connect
+      </button>
+    </div>
+    """
+  end
+
+  @impl true
+  def handle_event("init", _params, socket) do
+    {:ok, runtime} = Runtime.Embedded.init()
+    Session.connect_runtime(socket.assigns.session_id, runtime)
+    {:noreply, socket}
+  end
+end

lib/livebook_web/live/session_live/runtime_component.ex

@@ -85,6 +85,11 @@ defmodule LivebookWeb.SessionLive.RuntimeComponent do
             phx_click: "set_runtime_type",
             phx_value_type: "attached",
             phx_target: @myself %>
+          <%= content_tag :button, "Embedded",
+            class: "choice-button #{if(@type == "embedded", do: "active")}",
+            phx_click: "set_runtime_type",
+            phx_value_type: "embedded",
+            phx_target: @myself %>
         </div>
         <div>
           <%= if @type == "elixir_standalone" do %>
@@ -102,6 +107,11 @@ defmodule LivebookWeb.SessionLive.RuntimeComponent do
               id: :attached_runtime,
               session: %{"session_id" => @session_id, "current_runtime" => @runtime} %>
           <% end %>
+          <%= if @type == "embedded" do %>
+            <%= live_render @socket, LivebookWeb.SessionLive.EmbeddedLive,
+              id: :embedded_runtime,
+              session: %{"session_id" => @session_id, "current_runtime" => @runtime} %>
+          <% end %>
         </div>
       </div>
     </div>
@@ -111,10 +121,12 @@ defmodule LivebookWeb.SessionLive.RuntimeComponent do
   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"
+  defp runtime_type_label(%Runtime.Embedded{}), do: "Embedded"
 
   defp runtime_type(%Runtime.ElixirStandalone{}), do: "elixir_standalone"
   defp runtime_type(%Runtime.MixStandalone{}), do: "mix_standalone"
   defp runtime_type(%Runtime.Attached{}), do: "attached"
+  defp runtime_type(%Runtime.Embedded{}), do: "embedded"
 
   @impl true
   def handle_event("set_runtime_type", %{"type" => type}, socket) do

test/livebook/session/data_test.exs

@@ -2,7 +2,7 @@ defmodule Livebook.Session.DataTest do
   use ExUnit.Case, async: true
 
   alias Livebook.Session.Data
-  alias Livebook.{Delta, Notebook}
+  alias Livebook.{Delta, Notebook, Runtime}
   alias Livebook.Users.User
 
   describe "new/1" do
@@ -1610,7 +1610,8 @@ defmodule Livebook.Session.DataTest do
     test "updates data with the given runtime" do
       data = Data.new()
 
-      {:ok, runtime} = LivebookTest.Runtime.SingleEvaluator.init()
+      {:ok, runtime} = Runtime.Embedded.init()
+      Runtime.connect(runtime)
 
       operation = {:set_runtime, self(), runtime}
 
@@ -1634,7 +1635,8 @@ defmodule Livebook.Session.DataTest do
           {:queue_cell_evaluation, self(), "c4"}
         ])
 
-      {:ok, runtime} = LivebookTest.Runtime.SingleEvaluator.init()
+      {:ok, runtime} = Runtime.Embedded.init()
+      Runtime.connect(runtime)
 
       operation = {:set_runtime, self(), runtime}
 

test/livebook/session_test.exs

@@ -3,6 +3,12 @@ defmodule Livebook.SessionTest do
 
   alias Livebook.{Session, Delta, Runtime, Utils}
 
+  # Note: queueing evaluation in most of the tests below
+  # requires the runtime to synchronously start first,
+  # so we use a longer timeout just to make sure the tests
+  # pass reliably
+  @evaluation_wait_timeout 2_000
+
   setup do
     session_id = start_session()
     %{session_id: session_id}
@@ -63,7 +69,9 @@ defmodule Livebook.SessionTest do
       {_section_id, cell_id} = insert_section_and_cell(session_id)
 
       Session.queue_cell_evaluation(session_id, cell_id)
-      assert_receive {:operation, {:queue_cell_evaluation, ^pid, ^cell_id}}
+
+      assert_receive {:operation, {:queue_cell_evaluation, ^pid, ^cell_id}},
+                     @evaluation_wait_timeout
     end
 
     test "triggers evaluation and sends update operation once it finishes",
@@ -73,7 +81,9 @@ defmodule Livebook.SessionTest do
       {_section_id, cell_id} = insert_section_and_cell(session_id)
 
       Session.queue_cell_evaluation(session_id, cell_id)
-      assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _}}
+
+      assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _}},
+                     @evaluation_wait_timeout
     end
   end
 
@@ -83,10 +93,12 @@ defmodule Livebook.SessionTest do
       pid = self()
 
       {_section_id, cell_id} = insert_section_and_cell(session_id)
-      queue_evaluation(session_id, cell_id)
+      Session.queue_cell_evaluation(session_id, cell_id)
 
       Session.cancel_cell_evaluation(session_id, cell_id)
-      assert_receive {:operation, {:cancel_cell_evaluation, ^pid, ^cell_id}}
+
+      assert_receive {:operation, {:cancel_cell_evaluation, ^pid, ^cell_id}},
+                     @evaluation_wait_timeout
     end
   end
 
@@ -158,7 +170,7 @@ defmodule Livebook.SessionTest do
       Phoenix.PubSub.subscribe(Livebook.PubSub, "sessions:#{session_id}")
       pid = self()
 
-      {:ok, runtime} = LivebookTest.Runtime.SingleEvaluator.init()
+      {:ok, runtime} = Livebook.Runtime.Embedded.init()
       Session.connect_runtime(session_id, runtime)
 
       assert_receive {:operation, {:set_runtime, ^pid, ^runtime}}
@@ -170,6 +182,9 @@ defmodule Livebook.SessionTest do
       Phoenix.PubSub.subscribe(Livebook.PubSub, "sessions:#{session_id}")
       pid = self()
 
+      {:ok, runtime} = Livebook.Runtime.Embedded.init()
+      Session.connect_runtime(session_id, runtime)
+
       Session.disconnect_runtime(session_id)
 
       assert_receive {:operation, {:set_runtime, ^pid, nil}}
@@ -337,9 +352,10 @@ defmodule Livebook.SessionTest do
     end
   end
 
-  # For most tests we use the lightweight runtime, so that they are cheap to run.
-  # Here go several integration tests that actually start a separate runtime
-  # to verify session integrates well with it.
+  # For most tests we use the lightweight embedded runtime,
+  # so that they are cheap to run. Here go several integration
+  # tests that actually start a Elixir standalone runtime (default in production)
+  # to verify session integrates well with it properly.
 
   test "starts a standalone runtime upon first evaluation if there was none set explicitly" do
     session_id = Utils.random_id()
@@ -351,7 +367,8 @@ defmodule Livebook.SessionTest do
 
     Session.queue_cell_evaluation(session_id, cell_id)
     # Give it a bit more time as this involves starting a system process.
-    assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _}}, 1000
+    assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _}},
+                   @evaluation_wait_timeout
   end
 
   test "if the runtime node goes down, notifies the subscribers" do
@@ -387,10 +404,6 @@ defmodule Livebook.SessionTest do
   defp start_session(opts \\ []) do
     session_id = Utils.random_id()
     {:ok, _} = Session.start_link(Keyword.merge(opts, id: session_id))
-    # By default, use the current node for evaluation,
-    # rather than starting a standalone one.
-    {:ok, runtime} = LivebookTest.Runtime.SingleEvaluator.init()
-    Session.connect_runtime(session_id, runtime)
     session_id
   end
 
@@ -402,9 +415,4 @@ defmodule Livebook.SessionTest do
 
     {section_id, cell_id}
   end
-
-  defp queue_evaluation(session_id, cell_id) do
-    Session.queue_cell_evaluation(session_id, cell_id)
-    assert_receive {:operation, {:add_cell_evaluation_response, _, ^cell_id, _}}
-  end
 end

test/livebook_web/live/session_live_test.exs

@@ -235,7 +235,7 @@ defmodule LivebookWeb.SessionLiveTest do
       section_id = insert_section(session_id)
       cell_id = insert_cell(session_id, section_id, :elixir, "Process.sleep(10)")
 
-      {:ok, runtime} = LivebookTest.Runtime.SingleEvaluator.init()
+      {:ok, runtime} = Livebook.Runtime.Embedded.init()
       Session.connect_runtime(session_id, runtime)
 
       {:ok, view, _} = live(conn, "/sessions/#{session_id}")
@@ -367,9 +367,15 @@ defmodule LivebookWeb.SessionLiveTest do
     Session.insert_cell(session_id, section_id, 0, type)
     %{notebook: %{sections: [%{cells: [cell]}]}} = Session.get_data(session_id)
 
+    # We need to register ourselves as a client to start submitting cell deltas
+    user = Livebook.Users.User.new()
+    Session.register_client(session_id, self(), user)
+
     delta = Delta.new(insert: content)
     Session.apply_cell_delta(session_id, cell.id, delta, 1)
 
+    wait_for_session_update(session_id)
+
     cell.id
   end
 

test/support/single_evaluator_runtime.ex

@@ -1,63 +0,0 @@
-defmodule LivebookTest.Runtime.SingleEvaluator do
-  @moduledoc false
-
-  # A simple runtime backed by a single evaluator process
-  # running on the local node.
-  #
-  # This allows for working code evaluation without
-  # starting a standalone Elixir runtime, so is neat for testing.
-
-  defstruct [:evaluator]
-
-  def init() do
-    with {:ok, evaluator} <- Livebook.Evaluator.start_link() do
-      {:ok, %__MODULE__{evaluator: evaluator}}
-    end
-  end
-end
-
-defimpl Livebook.Runtime, for: LivebookTest.Runtime.SingleEvaluator do
-  alias Livebook.Evaluator
-
-  def connect(runtime) do
-    Process.monitor(runtime.evaluator)
-  end
-
-  def disconnect(_runtime), do: :ok
-
-  def evaluate_code(
-        runtime,
-        code,
-        _container_ref,
-        evaluation_ref,
-        prev_evaluation_ref,
-        opts \\ []
-      ) do
-    Evaluator.evaluate_code(
-      runtime.evaluator,
-      self(),
-      code,
-      evaluation_ref,
-      prev_evaluation_ref,
-      opts
-    )
-
-    :ok
-  end
-
-  def forget_evaluation(runtime, _container_ref, evaluation_ref) do
-    Evaluator.forget_evaluation(runtime.evaluator, evaluation_ref)
-  end
-
-  def drop_container(_runtime, _container_ref), do: :ok
-
-  def request_completion_items(runtime, send_to, ref, hint, _container_ref, evaluation_ref) do
-    Evaluator.request_completion_items(
-      runtime.evaluator,
-      send_to,
-      ref,
-      hint,
-      evaluation_ref
-    )
-  end
-end