diff --git a/lib/livebook/runtime.ex b/lib/livebook/runtime.ex index 00ed07ed4..f7a6d4ee8 100644 --- a/lib/livebook/runtime.ex +++ b/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 diff --git a/lib/livebook/runtime/attached.ex b/lib/livebook/runtime/attached.ex index fc3bf1d10..b7b5288b5 100644 --- a/lib/livebook/runtime/attached.ex +++ b/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 diff --git a/lib/livebook/runtime/elixir_standalone.ex b/lib/livebook/runtime/elixir_standalone.ex index 388babc2e..f18c83933 100644 --- a/lib/livebook/runtime/elixir_standalone.ex +++ b/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 diff --git a/lib/livebook/runtime/embedded.ex b/lib/livebook/runtime/embedded.ex index f95098de4..bd094d67e 100644 --- a/lib/livebook/runtime/embedded.ex +++ b/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 diff --git a/lib/livebook/runtime/erl_dist/runtime_server.ex b/lib/livebook/runtime/erl_dist/runtime_server.ex index b0a624235..cb9a7c732 100644 --- a/lib/livebook/runtime/erl_dist/runtime_server.ex +++ b/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 diff --git a/lib/livebook/session.ex b/lib/livebook/session.ex index fffc2daba..62038a53a 100644 --- a/lib/livebook/session.ex +++ b/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 diff --git a/lib/livebook/session/data.ex b/lib/livebook/session/data.ex index 215396eaf..a55465e19 100644 --- a/lib/livebook/session/data.ex +++ b/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() diff --git a/test/livebook/session/data_test.exs b/test/livebook/session/data_test.exs index 022e0533e..b2749ef8d 100644 --- a/test/livebook/session/data_test.exs +++ b/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 = diff --git a/test/livebook/session_test.exs b/test/livebook/session_test.exs index 17d8a8b8b..023072418 100644 --- a/test/livebook/session_test.exs +++ b/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() diff --git a/test/support/noop_runtime.ex b/test/support/noop_runtime.ex index 1b4b6771b..48285a370 100644 --- a/test/support/noop_runtime.ex +++ b/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})