Warn if session is missing after reboot, closes #2075

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

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

lib/livebook/utils.ex

@@ -47,7 +47,7 @@ defmodule Livebook.Utils do
     random_part = :crypto.strong_rand_bytes(9)
     binary = <<node_part::binary, random_part::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)
-    <<node_part::binary-size(16), _random_part::binary-size(9)>> = binary
-
-    known_nodes = [node() | Node.list()]
-
-    Enum.find_value(known_nodes, :error, fn node ->
-      node_hash(node) == node_part && {:ok, node}
-    end)
+    with <<boot_id::binary-size(4), id::binary>> <- id,
+         {:ok, binary} <- Base.decode32(id, case: :lower),
+         <<node_part::binary-size(16), _random_part::binary-size(9)>> <-
+           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 """

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, :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
 

test/livebook/sessions_test.exs

@@ -32,14 +32,17 @@ 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
       {:ok, session} = Sessions.create_session()
       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