From 5a367e5ea2bfa6d857943e35415dbd0ca38ef9d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Valim?= Date: Tue, 1 Aug 2023 13:57:03 +0200 Subject: [PATCH] Warn if session is missing after reboot, closes #2075 --- lib/livebook.ex | 2 ++ lib/livebook/config.ex | 7 +++++ lib/livebook/sessions.ex | 28 +++++++++++++------ lib/livebook/utils.ex | 22 +++++++++------ .../controllers/session_controller.ex | 4 +-- lib/livebook_web/live/session_live.ex | 12 +++++++- test/livebook/sessions_test.exs | 5 +++- 7 files changed, 58 insertions(+), 22 deletions(-) diff --git a/lib/livebook.ex b/lib/livebook.ex index 45a9ba4d5a4..cc0f5d3bc2c 100644 --- a/lib/livebook.ex +++ b/lib/livebook.ex @@ -86,6 +86,8 @@ defmodule Livebook do def config_runtime do import Config + config :livebook, :random_boot_id, Base.encode64(:crypto.strong_rand_bytes(3)) + config :livebook, LivebookWeb.Endpoint, secret_key_base: Livebook.Config.secret!("LIVEBOOK_SECRET_KEY_BASE") || diff --git a/lib/livebook/config.ex b/lib/livebook/config.ex index c8ff948d9ed..33d7d61b8f2 100644 --- a/lib/livebook/config.ex +++ b/lib/livebook/config.ex @@ -305,6 +305,13 @@ defmodule Livebook.Config do Application.fetch_env!(:livebook, :allowed_uri_schemes) end + @doc """ + Returns a random id set on boot. + """ + def random_boot_id() do + Application.fetch_env!(:livebook, :random_boot_id) + end + ## Parsing @doc """ diff --git a/lib/livebook/sessions.ex b/lib/livebook/sessions.ex index d71353085d6..6842337ad4b 100644 --- a/lib/livebook/sessions.ex +++ b/lib/livebook/sessions.ex @@ -50,21 +50,31 @@ defmodule Livebook.Sessions do @doc """ Returns tracked session with the given id. """ - @spec fetch_session(Session.id()) :: {:ok, Session.t()} | :error + @spec fetch_session(Session.id()) :: {:ok, Session.t()} | {:error, :not_found | :maybe_crashed} def fetch_session(id) do case Livebook.Tracker.fetch_session(id) do {:ok, session} -> {:ok, session} :error -> - # The local tracker server doesn't know about this session, - # but it may not have propagated yet, so we extract the session - # node from id and ask the corresponding tracker directly - with {:ok, other_node} when other_node != node() <- Utils.node_from_node_aware_id(id), - {:ok, session} <- :rpc.call(other_node, Livebook.Tracker, :fetch_session, [id]) do - {:ok, session} - else - _ -> :error + boot_id = Livebook.Config.random_boot_id() + + case Utils.node_from_node_aware_id(id) do + # The local tracker server doesn't know about this session, + # but it may not have propagated yet, so we extract the session + # node from id and ask the corresponding tracker directly + {:ok, other_node, _other_boot_id} when other_node != node() -> + case :rpc.call(other_node, Livebook.Tracker, :fetch_session, [id]) do + {:ok, session} -> {:ok, session} + _ -> {:error, :not_found} + end + + {:ok, other_node, other_boot_id} + when other_node == node() and other_boot_id != boot_id -> + {:error, :maybe_crashed} + + _ -> + {:error, :not_found} end end end diff --git a/lib/livebook/utils.ex b/lib/livebook/utils.ex index 24a561f0b2f..29563572c44 100644 --- a/lib/livebook/utils.ex +++ b/lib/livebook/utils.ex @@ -47,7 +47,7 @@ defmodule Livebook.Utils do random_part = :crypto.strong_rand_bytes(9) binary = <> # 16B + 9B = 25B is suitable for base32 encoding without padding - Base.encode32(binary, case: :lower) + Livebook.Config.random_boot_id() <> Base.encode32(binary, case: :lower) end # Note: the result is always 16 bytes long @@ -63,14 +63,18 @@ defmodule Livebook.Utils do """ @spec node_from_node_aware_id(id()) :: {:ok, node()} | :error def node_from_node_aware_id(id) do - binary = Base.decode32!(id, case: :lower) - <> = binary - - known_nodes = [node() | Node.list()] - - Enum.find_value(known_nodes, :error, fn node -> - node_hash(node) == node_part && {:ok, node} - end) + with <> <- id, + {:ok, binary} <- Base.decode32(id, case: :lower), + <> <- + binary do + known_nodes = [node() | Node.list()] + + Enum.find_value(known_nodes, :error, fn node -> + node_hash(node) == node_part && {:ok, node, boot_id} + end) + else + _ -> :error + end end @doc """ diff --git a/lib/livebook_web/controllers/session_controller.ex b/lib/livebook_web/controllers/session_controller.ex index 6d8e6f0d66c..fffa581be4e 100644 --- a/lib/livebook_web/controllers/session_controller.ex +++ b/lib/livebook_web/controllers/session_controller.ex @@ -34,7 +34,7 @@ defmodule LivebookWeb.SessionController do file = FileSystem.File.resolve(images_dir, name) serve_static(conn, file) - :error -> + {:error, _} -> send_resp(conn, 404, "Not found") end end @@ -57,7 +57,7 @@ defmodule LivebookWeb.SessionController do send_notebook_source(conn, notebook, file_name, format) - :error -> + {:error, _} -> send_resp(conn, 404, "Not found") end end diff --git a/lib/livebook_web/live/session_live.ex b/lib/livebook_web/live/session_live.ex index eca503b3d27..77b4287f062 100644 --- a/lib/livebook_web/live/session_live.ex +++ b/lib/livebook_web/live/session_live.ex @@ -94,8 +94,18 @@ defmodule LivebookWeb.SessionLive do |> prune_outputs() |> prune_cell_sources()} - :error -> + {:error, :not_found} -> {:ok, redirect(socket, to: ~p"/")} + + {:error, :maybe_crashed} -> + {:ok, + socket + |> put_flash( + :error, + "Could not find notebook session because Livebook has rebooted. " <> + "This may happen if Livebook runs out of memory while installing dependencies or executing code." + ) + |> redirect(to: ~p"/")} end end diff --git a/test/livebook/sessions_test.exs b/test/livebook/sessions_test.exs index 9e533899c0f..0f6052328b4 100644 --- a/test/livebook/sessions_test.exs +++ b/test/livebook/sessions_test.exs @@ -32,7 +32,7 @@ defmodule Livebook.SessionsTest do describe "fetch_session/1" do test "returns an error if no session with the given id exists" do id = Livebook.Utils.random_node_aware_id() - assert :error = Sessions.fetch_session(id) + assert Sessions.fetch_session(id) == {:error, :not_found} end test "returns session matching the given id" do @@ -40,6 +40,9 @@ defmodule Livebook.SessionsTest do assert {:ok, ^session} = Sessions.fetch_session(session.id) Session.close(session.pid) + assert Sessions.fetch_session(session.id) == {:error, :not_found} + Application.put_env(:livebook, :random_boot_id, "^^^^") + assert Sessions.fetch_session(session.id) == {:error, :maybe_crashed} end end