Warn if session is missing after reboot (#2128)

Closes #2075.

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") ||

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 """

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

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 = <<node_part::binary, random_part::binary>>
-    # 16B + 9B = 25B is suitable for base32 encoding without padding
+    random_part = :crypto.strong_rand_bytes(11)
+    binary = <<boot_id::binary, node_part::binary, random_part::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)
-    <<node_part::binary-size(16), _random_part::binary-size(9)>> = binary
+    case Base.decode32(id, case: :lower) do
+      {:ok,
+       <<boot_id::binary-size(3), node_part::binary-size(16), _random_part::binary-size(11)>>} ->
+        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 """

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

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
 

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