Repository: livebook-dev/kino
Pull request #269 (2023-03-17 12:49:39)

Only terminate evaluator on listen exit with non-normal reasons


Patch:
diff --git a/lib/kino.ex b/lib/kino.ex
index 0f39dad..2772881 100644
--- a/lib/kino.ex
+++ b/lib/kino.ex
@@ -435,13 +435,12 @@ defmodule Kino do
   end
 
   defp async(fun) do
-    {:ok, pid} = Kino.start_child({Task, fun})
-    # In case the stream breaks we want to propagate the failure
-    try do
-      Process.link(pid)
-    catch
-      :error, :noproc -> :ok
-    end
+    {:ok, _pid} =
+      Kino.start_child(%{
+        id: Task,
+        start: {Kino.Terminator, :start_task, [self(), fun]},
+        restart: :temporary
+      })
 
     :ok
   end
@@ -519,37 +518,7 @@ defmodule Kino do
     %{start: start} = child_spec = Supervisor.child_spec(child_spec, [])
     parent = self()
     gl = Process.group_leader()
-    child_spec = %{child_spec | start: {Kino, :__start_override__, [start, parent, gl]}}
+    child_spec = %{child_spec | start: {Kino.Terminator, :start_child, [start, parent, gl]}}
     DynamicSupervisor.start_child(Kino.DynamicSupervisor, child_spec)
   end
-
-  @doc false
-  def __start_override__({mod, fun, args}, parent, gl) do
-    # We switch the group leader, so that the newly started
-    # process gets the same group leader as the caller
-    initial_gl = Process.group_leader()
-
-    Process.group_leader(self(), gl)
-
-    try do
-      {resp, pid} =
-        case apply(mod, fun, args) do
-          {:ok, pid} = resp -> {resp, pid}
-          {:ok, pid, _info} = resp -> {resp, pid}
-          resp -> {resp, nil}
-        end
-
-      if pid do
-        Kino.Bridge.reference_object(pid, parent)
-
-        Kino.Bridge.monitor_object(pid, Kino.Terminator.cross_node_name(), {:terminate, pid},
-          ack?: true
-        )
-      end
-
-      resp
-    after
-      Process.group_leader(self(), initial_gl)
-    end
-  end
 end
diff --git a/lib/kino/terminator.ex b/lib/kino/terminator.ex
index 487378f..901f79a 100644
--- a/lib/kino/terminator.ex
+++ b/lib/kino/terminator.ex
@@ -11,6 +11,45 @@ defmodule Kino.Terminator do
     {@name, node()}
   end
 
+  @doc """
+  Starts a Kino process to be shutdown by the terminator.
+  """
+  def start_child({mod, fun, args}, parent, gl) do
+    # We switch the group leader, so that the newly started
+    # process gets the same group leader as the caller
+    initial_gl = Process.group_leader()
+
+    Process.group_leader(self(), gl)
+
+    try do
+      {resp, pid} =
+        case apply(mod, fun, args) do
+          {:ok, pid} = resp -> {resp, pid}
+          {:ok, pid, _info} = resp -> {resp, pid}
+          resp -> {resp, nil}
+        end
+
+      if pid do
+        Kino.Bridge.reference_object(pid, parent)
+        Kino.Bridge.monitor_object(pid, cross_node_name(), {:terminate, pid}, ack?: true)
+      end
+
+      resp
+    after
+      Process.group_leader(self(), initial_gl)
+    end
+  end
+
+  @doc """
+  Starts a task that will terminate the parent in case of crashes.
+  """
+  def start_task(parent, fun) do
+    Task.start_link(fn ->
+      GenServer.call(@name, {:monitor, self(), parent})
+      fun.()
+    end)
+  end
+
   @doc """
   Starts the terminator.
   """
@@ -23,10 +62,32 @@ defmodule Kino.Terminator do
     {:ok, %{}}
   end
 
+  @impl true
+  def handle_call({:monitor, pid, parent}, _from, state) do
+    _ref = Process.monitor(pid)
+    {:reply, :ok, Map.put(state, pid, parent)}
+  end
+
   @impl true
   def handle_info({{:terminate, pid}, reply_to, reply_as}, state) do
     DynamicSupervisor.terminate_child(Kino.DynamicSupervisor, pid)
     send(reply_to, reply_as)
+    {:noreply, Map.delete(state, pid)}
+  end
+
+  @impl true
+  def handle_info({:DOWN, _, _, pid, reason}, state) do
+    {parent, state} = Map.pop(state, pid)
+
+    if is_pid(parent) and abnormal?(reason) do
+      Process.exit(parent, reason)
+    end
+
     {:noreply, state}
   end
+
+  defp abnormal?(:normal), do: false
+  defp abnormal?(:shutdown), do: false
+  defp abnormal?({:shutdown, _}), do: false
+  defp abnormal?(_), do: true
 end


Test patch:
diff --git a/test/kino_test.exs b/test/kino_test.exs
index 6c36721..5c8ea87 100644
--- a/test/kino_test.exs
+++ b/test/kino_test.exs
@@ -137,6 +137,17 @@ defmodule KinoTest do
       assert log =~ "Kino.listen"
       assert log =~ "** (MatchError) no match of right hand side value: 0"
     end
+
+    @tag :capture_log
+    test "kills parent if stream itself fails" do
+      Process.flag(:trap_exit, true)
+
+      1..10
+      |> Stream.map(fn _ -> exit(:oops) end)
+      |> Kino.listen(&IO.inspect/1)
+
+      assert_receive {:EXIT, _, :oops}
+    end
   end
 
   describe "listen/3" do