Fix race conditions in message scheduling (#916)

* Fix race conditions in message scheduling

* Flush already fired timer message

* Move unscheduling to a separate function

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -132,6 +132,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   @impl true
   def init(_opts) do
     Process.send_after(self(), :check_owner, @await_owner_timeout)
+    schedule_memory_usage_report()
 
     {:ok, evaluator_supervisor} = ErlDist.EvaluatorSupervisor.start_link()
     {:ok, task_supervisor} = Task.Supervisor.start_link()
@@ -180,8 +181,9 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   end
 
   def handle_info(:memory_usage, state) do
-    send(state.owner, {:memory_usage, Evaluator.memory()})
-    {:noreply, schedule_memory_usage(state)}
+    report_memory_usage(state)
+    schedule_memory_usage_report()
+    {:noreply, state}
   end
 
   def handle_info(_message, state), do: {:noreply, state}
@@ -189,8 +191,9 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   @impl true
   def handle_cast({:set_owner, owner}, state) do
     Process.monitor(owner)
-    send(self(), :memory_usage)
-    {:noreply, %{state | owner: owner}}
+    state = %{state | owner: owner}
+    report_memory_usage(state)
+    {:noreply, state}
   end
 
   def handle_cast(
@@ -305,8 +308,13 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     end
   end
 
-  defp schedule_memory_usage(state) do
-    ref = Process.send_after(self(), :memory_usage, @memory_usage_interval)
-    %{state | memory_timer_ref: ref}
+  defp schedule_memory_usage_report() do
+    Process.send_after(self(), :memory_usage, @memory_usage_interval)
+  end
+
+  defp report_memory_usage(%{owner: nil}), do: :ok
+
+  defp report_memory_usage(state) do
+    send(state.owner, {:memory_usage, Evaluator.memory()})
   end
 end

lib/livebook/session.ex

@@ -517,6 +517,18 @@ defmodule Livebook.Session do
     end
   end
 
+  defp unschedule_autosave(%{autosave_timer_ref: nil} = state), do: state
+
+  defp unschedule_autosave(state) do
+    if Process.cancel_timer(state.autosave_timer_ref) == false do
+      receive do
+        :autosave -> :ok
+      end
+    end
+
+    %{state | autosave_timer_ref: nil}
+  end
+
   @impl true
   def handle_call(:describe_self, _from, state) do
     {:reply, self_from_state(state), state}
@@ -1034,11 +1046,9 @@ defmodule Livebook.Session do
          _prev_state,
          {:set_notebook_attributes, _client_pid, %{autosave_interval_s: _}}
        ) do
-    if ref = state.autosave_timer_ref do
-      Process.cancel_timer(ref)
-    end
-
-    schedule_autosave(state)
+    state
+    |> unschedule_autosave()
+    |> schedule_autosave()
   end
 
   defp after_operation(state, prev_state, {:client_join, _client_pid, user}) do