Fix race condition in app auto shutdown test (#2887)

lib/livebook/runtime/evaluator.ex

@@ -91,10 +91,7 @@ defmodule Livebook.Runtime.Evaluator do
   """
   @spec start_link(keyword()) :: {:ok, pid(), t()} | {:error, term()}
   def start_link(opts \\ []) do
-    case :proc_lib.start_link(__MODULE__, :init, [opts]) do
-      {:error, error} -> {:error, error}
-      evaluator -> {:ok, evaluator.pid, evaluator}
-    end
+    :proc_lib.start_link(__MODULE__, :init, [opts])
   end
 
   @doc """
@@ -322,7 +319,7 @@ defmodule Livebook.Runtime.Evaluator do
       tmp_dir: tmp_dir
     }
 
-    :proc_lib.init_ack(evaluator)
+    :proc_lib.init_ack({:ok, evaluator.pid, evaluator})
 
     loop(state)
   end

lib/livebook/session.ex

@@ -199,9 +199,13 @@ defmodule Livebook.Session do
       session belongs to. This is only relevant for app sessions
 
   """
-  @spec start_link(keyword()) :: {:ok, pid} | {:error, any()}
+  @spec start_link(keyword()) :: {:ok, pid, t()} | {:error, any()}
   def start_link(opts) do
-    GenServer.start_link(__MODULE__, opts)
+    with {:ok, pid} <- GenServer.start_link(__MODULE__, {self(), opts}) do
+      receive do
+        {:started, ^pid, session} -> {:ok, pid, session}
+      end
+    end
   end
 
   @doc """
@@ -873,7 +877,7 @@ defmodule Livebook.Session do
   ## Callbacks
 
   @impl true
-  def init(opts) do
+  def init({caller_pid, opts}) do
     Livebook.Settings.subscribe()
     Livebook.Hubs.Broadcasts.subscribe([:crud, :secrets])
 
@@ -901,6 +905,8 @@ defmodule Livebook.Session do
       session = self_from_state(state)
 
       with :ok <- Livebook.Tracker.track_session(session) do
+        send(caller_pid, {:started, self(), session})
+
         if state.data.mode == :app do
           {:ok, state, {:continue, :app_init}}
         else

lib/livebook/sessions.ex

@@ -18,8 +18,8 @@ defmodule Livebook.Sessions do
     opts = Keyword.put(opts, :id, id)
 
     case DynamicSupervisor.start_child(Livebook.SessionSupervisor, {Session, opts}) do
-      {:ok, pid} ->
-        {:ok, Session.get_by_pid(pid)}
+      {:ok, _pid, session} ->
+        {:ok, session}
 
       {:error, reason} ->
         {:error, reason}

test/livebook/app_test.exs

@@ -221,7 +221,7 @@ defmodule Livebook.AppTest do
         Notebook.AppSettings.new()
         | slug: slug,
           multi_session: true,
-          auto_shutdown_ms: 10
+          auto_shutdown_ms: 1
       }
 
       notebook = %{Notebook.new() | app_settings: app_settings}