Restore Mix.install/2 dirs across session runtimes (#2499)

lib/livebook/runtime.ex

@@ -727,6 +727,25 @@ defprotocol Livebook.Runtime do
   """
   @type file_ref :: {:file, id :: String.t()}
 
+  @typedoc """
+  A state that can optionally be passed from one runtime to another.
+
+  To report a new transition state, the runtime may send:
+
+      {:runtime_transition_state, transition_state()}
+
+  The runtime owner can then use `restore_transient_state/2` when
+  starting another instance of this runtime.
+
+  The state should be considered complementary, it is not guaranteed
+  that any future runtime will receive it. Therefore, a transient state
+  should never point to resources with the expectation that a future
+  runtime will clean them. One valid use case is for the transient state
+  to point to some global cache, that is not managed by the runtime
+  itself.
+  """
+  @type transient_state :: %{atom() => term()}
+
   @doc """
   Returns relevant information about the runtime.
 
@@ -1045,4 +1064,12 @@ defprotocol Livebook.Runtime do
   """
   @spec delete_system_envs(t(), list(String.t())) :: :ok
   def delete_system_envs(runtime, names)
+
+  @doc """
+  Restores information from a past runtime.
+
+  See `t:transient_state/0` for details.
+  """
+  @spec restore_transient_state(t(), transient_state()) :: :ok
+  def restore_transient_state(runtime, transient_state)
 end

lib/livebook/runtime/attached.ex

@@ -192,4 +192,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Attached do
   def delete_system_envs(runtime, names) do
     RuntimeServer.delete_system_envs(runtime.server_pid, names)
   end
+
+  def restore_transient_state(runtime, transient_state) do
+    RuntimeServer.restore_transient_state(runtime.server_pid, transient_state)
+  end
 end

lib/livebook/runtime/elixir_standalone.ex

@@ -190,4 +190,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.ElixirStandalone do
   def delete_system_envs(runtime, names) do
     RuntimeServer.delete_system_envs(runtime.server_pid, names)
   end
+
+  def restore_transient_state(runtime, transient_state) do
+    RuntimeServer.restore_transient_state(runtime.server_pid, transient_state)
+  end
 end

lib/livebook/runtime/embedded.ex

@@ -159,6 +159,10 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Embedded do
     RuntimeServer.delete_system_envs(runtime.server_pid, names)
   end
 
+  def restore_transient_state(runtime, transient_state) do
+    RuntimeServer.restore_transient_state(runtime.server_pid, transient_state)
+  end
+
   defp config() do
     Application.get_env(:livebook, Livebook.Runtime.Embedded, [])
   end

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -293,6 +293,14 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     GenServer.cast(pid, {:delete_system_envs, names})
   end
 
+  @doc """
+  Restores information from a past runtime.
+  """
+  @spec restore_transient_state(pid(), Runtime.transient_state()) :: :ok
+  def restore_transient_state(pid, transient_state) do
+    GenServer.cast(pid, {:restore_transient_state, transient_state})
+  end
+
   @doc """
   Stops the runtime server.
 
@@ -335,7 +343,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
          Keyword.get_lazy(opts, :base_env_path, fn -> System.get_env("PATH", "") end),
        ebin_path: Keyword.get(opts, :ebin_path),
        io_proxy_registry: Keyword.get(opts, :io_proxy_registry),
-       tmp_dir: Keyword.get(opts, :tmp_dir)
+       tmp_dir: Keyword.get(opts, :tmp_dir),
+       mix_install_project_dir: nil
      }}
   end
 
@@ -371,6 +380,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {:noreply,
      state
      |> report_smart_cell_definitions()
+     |> report_transient_state()
      |> scan_binding_after_evaluation(locator)}
   end
 
@@ -638,6 +648,14 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {:noreply, state}
   end
 
+  def handle_cast({:restore_transient_state, transient_state}, state) do
+    if dir = transient_state[:mix_install_project_dir] do
+      System.put_env("MIX_INSTALL_RESTORE_PROJECT_DIR", dir)
+    end
+
+    {:noreply, state}
+  end
+
   def handle_cast({:relabel_file, file_id, new_file_id}, state) do
     path = file_path(state, file_id)
     new_path = file_path(state, new_file_id)
@@ -779,6 +797,28 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     end
   end
 
+  defp report_transient_state(state) do
+    # We propagate Mix.install/2 project dir in the transient state,
+    # so that future runtimes can set it as the starting point for
+    # Mix.install/2
+    if dir = state.mix_install_project_dir == nil && install_project_dir() do
+      send(state.owner, {:runtime_transient_state, %{mix_install_project_dir: dir}})
+      %{state | mix_install_project_dir: dir}
+    else
+      state
+    end
+  end
+
+  # TODO: remove once CI runs Elixir v1.16.2
+  @compile {:no_warn_undefined, {Mix, :install_project_dir, 0}}
+
+  defp install_project_dir() do
+    # TODO: remove the check once we require Elixir v1.16.2
+    if Code.ensure_loaded?(Mix) && function_exported?(Mix, :install_project_dir, 0) do
+      Mix.install_project_dir()
+    end
+  end
+
   defp scan_binding_async(_ref, %{scan_binding: nil} = info, _state), do: info
 
   # We wait for the current scanning to finish, this way we avoid

lib/livebook/session.ex

@@ -1763,6 +1763,11 @@ defmodule Livebook.Session do
     {:noreply, state}
   end
 
+  def handle_info({:runtime_transient_state, transient_state}, state) do
+    operation = {:set_runtime_transient_state, @client_id, transient_state}
+    {:noreply, handle_operation(state, operation)}
+  end
+
   def handle_info({:env_var_set, env_var}, state) do
     if Runtime.connected?(state.data.runtime) do
       Runtime.put_system_envs(state.data.runtime, [{env_var.name, env_var.value}])
@@ -2030,6 +2035,11 @@ defmodule Livebook.Session do
 
   defp own_runtime(runtime, state) do
     runtime_monitor_ref = Runtime.take_ownership(runtime, runtime_broadcast_to: state.worker_pid)
+
+    if state.data.runtime_transient_state != %{} do
+      Runtime.restore_transient_state(runtime, state.data.runtime_transient_state)
+    end
+
     %{state | runtime_monitor_ref: runtime_monitor_ref}
   end
 

lib/livebook/session/data.ex

@@ -27,6 +27,7 @@ defmodule Livebook.Session.Data do
     :input_infos,
     :bin_entries,
     :runtime,
+    :runtime_transient_state,
     :smart_cell_definitions,
     :clients_map,
     :users_map,
@@ -53,6 +54,7 @@ defmodule Livebook.Session.Data do
           input_infos: %{input_id() => input_info()},
           bin_entries: list(cell_bin_entry()),
           runtime: Runtime.t(),
+          runtime_transient_state: Runtime.transient_state(),
           smart_cell_definitions: list(Runtime.smart_cell_definition()),
           clients_map: %{client_id() => User.id()},
           users_map: %{User.id() => User.t()},
@@ -211,6 +213,7 @@ defmodule Livebook.Session.Data do
           | {:set_cell_attributes, client_id(), Cell.id(), map()}
           | {:set_input_value, client_id(), input_id(), value :: term()}
           | {:set_runtime, client_id(), Runtime.t()}
+          | {:set_runtime_transient_state, client_id(), Runtime.transient_state()}
           | {:set_smart_cell_definitions, client_id(), list(Runtime.smart_cell_definition())}
           | {:set_file, client_id(), FileSystem.File.t() | nil}
           | {:set_autosave_interval, client_id(), non_neg_integer() | nil}
@@ -299,6 +302,7 @@ defmodule Livebook.Session.Data do
       input_infos: initial_input_infos(notebook),
       bin_entries: [],
       runtime: default_runtime,
+      runtime_transient_state: %{},
       smart_cell_definitions: [],
       clients_map: %{},
       users_map: %{},
@@ -860,6 +864,13 @@ defmodule Livebook.Session.Data do
     |> wrap_ok()
   end
 
+  def apply_operation(data, {:set_runtime_transient_state, _client_id, transient_state}) do
+    data
+    |> with_actions()
+    |> set!(runtime_transient_state: transient_state)
+    |> wrap_ok()
+  end
+
   def apply_operation(data, {:set_smart_cell_definitions, _client_id, definitions}) do
     data
     |> with_actions()

test/livebook/session/data_test.exs

@@ -3838,6 +3838,21 @@ defmodule Livebook.Session.DataTest do
     end
   end
 
+  describe "apply_operation/2 given :set_runtime_transient_state" do
+    test "sets the definitions and starts dead cells with matching kinds" do
+      data =
+        data_after_operations!([
+          {:set_runtime, @cid, connected_noop_runtime()}
+        ])
+
+      transient_state = %{state: "anything"}
+      operation = {:set_runtime_transient_state, @cid, transient_state}
+
+      assert {:ok, %{runtime_transient_state: ^transient_state}, _actions} =
+               Data.apply_operation(data, operation)
+    end
+  end
+
   describe "apply_operation/2 given :set_smart_cell_definitions" do
     test "sets the definitions and starts dead cells with matching kinds" do
       data =

test/livebook/session_test.exs

@@ -1281,6 +1281,20 @@ defmodule Livebook.SessionTest do
     assert :ok = Session.fetch_assets(session.pid, hash)
   end
 
+  test "restores transient state when restarting runtimes" do
+    session = start_session()
+
+    runtime = connected_noop_runtime(self())
+    Session.set_runtime(session.pid, runtime)
+    transient_state = %{state: "anything"}
+    send(session.pid, {:runtime_transient_state, transient_state})
+
+    runtime = connected_noop_runtime(self())
+    Session.set_runtime(session.pid, runtime)
+
+    assert_receive {:runtime_trace, :restore_transient_state, [^transient_state]}
+  end
+
   describe "deploy_app/1" do
     test "deploys current notebook and keeps track of the deployed app" do
       session = start_session()

test/support/noop_runtime.ex

@@ -66,6 +66,11 @@ defmodule Livebook.Runtime.NoopRuntime do
     def put_system_envs(_, _), do: :ok
     def delete_system_envs(_, _), do: :ok
 
+    def restore_transient_state(runtime, transient_state) do
+      trace(runtime, :restore_transient_state, [transient_state])
+      :ok
+    end
+
     defp trace(runtime, fun, args) do
       if runtime.trace_to do
         send(runtime.trace_to, {:runtime_trace, fun, args})