From 7c40ab22e330b830d8ca2cb3ffb6cd123178ce75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Valim?= Date: Tue, 1 Aug 2023 19:12:38 +0200 Subject: [PATCH] Warn if session is missing after reboot (#2128) Closes #2075. --- lib/livebook.ex | 2 ++ lib/livebook/config.ex | 7 +++++ lib/livebook/sessions.ex | 29 +++++++++++++------ lib/livebook/utils.ex | 26 ++++++++++------- .../controllers/session_controller.ex | 4 +-- lib/livebook_web/live/session_live.ex | 12 +++++++- test/livebook/sessions_test.exs | 11 ++++++- 7 files changed, 68 insertions(+), 23 deletions(-) diff --git a/lib/livebook.ex b/lib/livebook.ex index 45a9ba4d5..9dd0d08d1 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, :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 c8ff948d9..33d7d61b8 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 d71353085..e5630b1cc 100644 --- a/lib/livebook/sessions.ex +++ b/lib/livebook/sessions.ex @@ -50,21 +50,32 @@ 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 | :different_boot_id} 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, :different_boot_id} + + _ -> + {:error, :not_found} end end end diff --git a/lib/livebook/utils.ex b/lib/livebook/utils.ex index 24a561f0b..6e2738d1d 100644 --- a/lib/livebook/utils.ex +++ b/lib/livebook/utils.ex @@ -36,6 +36,7 @@ defmodule Livebook.Utils do The id is formed from the following binary parts: + * 3B - random boot id * 16B - hashed node name * 9B - random bytes @@ -43,10 +44,11 @@ defmodule Livebook.Utils do """ @spec random_node_aware_id() :: id() def random_node_aware_id() do + boot_id = Livebook.Config.random_boot_id() node_part = node_hash(node()) - random_part = :crypto.strong_rand_bytes(9) - binary = <> - # 16B + 9B = 25B is suitable for base32 encoding without padding + random_part = :crypto.strong_rand_bytes(11) + binary = <> + # 3B + 16B + 11B = 30B is suitable for base32 encoding without padding Base.encode32(binary, case: :lower) end @@ -61,16 +63,20 @@ defmodule Livebook.Utils do The node in question must be connected, otherwise it won't be found. """ - @spec node_from_node_aware_id(id()) :: {:ok, node()} | :error + @spec node_from_node_aware_id(id()) :: {:ok, node(), boot_id :: binary()} | :error def node_from_node_aware_id(id) do - binary = Base.decode32!(id, case: :lower) - <> = binary + case Base.decode32(id, case: :lower) do + {:ok, + <>} -> + known_nodes = [node() | Node.list()] - known_nodes = [node() | Node.list()] + Enum.find_value(known_nodes, :error, fn node -> + node_hash(node) == node_part && {:ok, node, boot_id} + end) - Enum.find_value(known_nodes, :error, fn node -> - node_hash(node) == node_part && {:ok, node} - end) + _ -> + :error + end end @doc """ diff --git a/lib/livebook_web/controllers/session_controller.ex b/lib/livebook_web/controllers/session_controller.ex index 6d8e6f0d6..fffa581be 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 eca503b3d..c29bac618 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, :different_boot_id} -> + {: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 9e533899c..182d9accd 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 @@ -41,6 +41,15 @@ defmodule Livebook.SessionsTest do Session.close(session.pid) end + + test "returns an error if session comes from a different boot" do + Application.put_env(:livebook, :random_boot_id, "aaa") + id = Livebook.Utils.random_node_aware_id() + + assert Sessions.fetch_session(id) == {:error, :not_found} + Application.put_env(:livebook, :random_boot_id, "bbb") + assert Sessions.fetch_session(id) == {:error, :different_boot_id} + end end describe "update_session/1" do