Propagate exit through links

lib/trinity/scheduler.ex

···
   defmodule Simulation do
     @enforce_keys [
       :queue,
+      :proc_queue_keys,
       :proc_links,
       :now,
       :supervisor_pid,
···
   defmodule SimulationSupervisor do
     use GenServer
 
+    defmodule State do
+      @enforce_keys [:sim, :monitors]
+      defstruct @enforce_keys
+    end
+
     def start_link(%Simulation{} = sim) do
-      GenServer.start_link(__MODULE__, sim)
+      GenServer.start(__MODULE__, sim)
     end
 
     def register_self(server) do
···
     end
 
     def init(%Simulation{} = sim) do
-      {:ok, sim}
+      state = %State{
+        sim: sim,
+        monitors: %{},
+      }
+      {:ok, state}
     end
 
-    def handle_cast({:register, pid}, state) do
-      Process.monitor(pid)
+    def handle_cast({:register, pid}, %State{} = state) do
+      mref = Process.monitor(pid)
+      state = %{state | monitors: Map.put(state.monitors, pid, mref)}
       {:noreply, state}
     end
 
-    def handle_info({:DOWN, _ref, :process, _pid, _reason}, state) do
-      Trinity.Scheduler.perform_next_task(state)
+    def handle_info({:DOWN, _ref, :process, pid, reason}, %State{} = state) do
+      killed = Trinity.Scheduler.handle_down(state.sim, pid, reason)
+      killed = [pid | killed]
+
+      monitors =
+        Enum.reduce(killed, state.monitors, fn p, acc ->
+          {mref, acc} = Map.pop!(acc, p)
+          Process.demonitor(mref, [:flush])
+          acc
+        end)
+
+      state = %{state | monitors: monitors}
       {:noreply, state}
     end
   end
···
   def start do
     sim = %Simulation{
       queue: :ets.new(__MODULE__, [:ordered_set, :public]),
+      proc_queue_keys: :ets.new(__MODULE__, [:set, :public]),
       proc_links: :ets.new(__MODULE__, [:set, :public]),
       now: :atomics.new(1, signed: false),
       supervisor_pid: nil,
···
 
   @spec yield(non_neg_integer) :: :ok
   def yield(delay \\ 0) do
-    %Simulation{queue: queue, now: now} = get_sim()
-    yield_ref = enqueue_self(queue, now, delay)
+    %Simulation{queue: queue, proc_queue_keys: proc_queue_keys, now: now} = get_sim()
+    yield_ref = enqueue_self(queue, proc_queue_keys, now, delay)
 
     perform_next(queue, now)
 
···
       # We then enqueue the new process and unsuspend the parent
       # so that it can register any links before yielding back
       # to us (eventually, if there are others in queue at `time=now`)
-      yield_ref = enqueue_self(sim.queue, sim.now, 0)
+      yield_ref = enqueue_self(sim.queue, sim.proc_queue_keys, sim.now, 0)
       send parent_pid, spawn_ref
       receive do
         ^yield_ref -> :noop
···
     :ok
   end
 
+  def handle_down(%Simulation{} = sim, _pid, :normal) do
+    perform_next(sim.queue, sim.now)
+    []
+  end
+
+  def handle_down(%Simulation{} = sim, pid, reason) do
+    %Simulation{proc_links: proc_links} = sim
+    linked = gather_linked(pid, proc_links)
+
+    Enum.each(linked, fn pid ->
+      destroy_process(sim, pid)
+      Process.exit(pid, reason)
+    end)
+
+    perform_next(sim.queue, sim.now)
+    linked
+  end
+
   defp perform_next(queue, now) do
     case pop_next(queue) do
       {time, {:resume, pid, ref}} ->
···
     end
   end
 
-  defp enqueue_self(queue, now, delay) do
+  defp enqueue_self(queue, proc_queue_keys, now, delay) do
     ref = make_ref()
     time = read_now(now) + delay
 
-    entry = {:resume, self(), ref}
+    pid = self()
+    entry = {:resume, pid, ref}
 
     i = case :ets.prev(queue, {time, :infinity}) do
-      [{^time, prev}] -> prev + 1
+      {^time, prev} -> prev + 1
       _ -> 0
     end
+    queue_key = {time, i}
     :ets.insert(queue, {{time, i}, entry})
+    :ets.insert(proc_queue_keys, {pid, queue_key})
 
     ref
   end
 
-  defp read_now(now), do: :atomics.get(now, 1)
-  defp set_now(now, time), do: :atomics.put(now, 1, time)
+  defp destroy_process(%Simulation{} = sim, pid) do
+    %Simulation{
+      queue: queue,
+      proc_links: proc_links,
+      proc_queue_keys: proc_queue_keys,
+    } = sim
+
+    destroy_links(proc_links, pid)
 
-  # Used by SimulationSupervisor
-  @doc false
-  def perform_next_task(%Simulation{queue: queue, now: now}) do
-    perform_next(queue, now)
+    case :ets.lookup(proc_queue_keys, pid) do
+      [{^pid, qk}] ->
+        :ets.delete(queue, qk)
+        :ets.delete(proc_queue_keys, pid)
+      _ -> :noop
+    end
   end
+
+  defp destroy_links(proc_links, from_pid) do
+    case :ets.lookup(proc_links, from_pid) do
+      [{^from_pid, from_links}] ->
+        # Remove `from_pid` from each of its links' links
+        Enum.each(from_links, fn to_pid ->
+          [{^to_pid, to_links}] = :ets.lookup(proc_links, to_pid)
+          :ets.insert(proc_links, {to_pid, List.delete(to_links, from_pid)})
+        end)
+
+        # Delete `from_pid`'s links entry
+        :ets.delete(proc_links, from_pid)
+
+      # No links, do nothing
+      _ -> :noop
+    end
+  end
+
+  defp gather_linked(pid, proc_links) do
+    # TODO: there must be a nicer way to structure this
+    # (especially the uniq)
+    gather_linked(pid, proc_links, [pid])
+    |> Enum.uniq()
+    |> List.delete(pid)
+  end
+
+  defp gather_linked(pid, proc_links, acc) do
+    case :ets.lookup(proc_links, pid) do
+      [{^pid, links}] ->
+        links = Enum.filter(links, &(&1 not in acc))
+        acc = links ++ acc
+
+        links_of_links =
+          Enum.reduce(links, acc, fn p, acc ->
+            gather_linked(p, proc_links, acc)
+            |> Enum.uniq()
+          end)
+
+        links_of_links ++ acc
+
+      _ -> acc
+    end
+  end
+
+  defp read_now(now), do: :atomics.get(now, 1)
+  defp set_now(now, time), do: :atomics.put(now, 1, time)
 
   @doc false
   def dump do

test/trinity_test.exs

···
   test "scheduler" do
     Scheduler.start()
 
-    pid1 = Scheduler.spawn_and_yield(fn -> dbg "hello 1" end, true)
-    pid2 = Scheduler.spawn_and_yield(fn -> dbg "hello 2" end, false)
-    pid3 = Scheduler.spawn_and_yield(fn -> dbg "hello 3" end, true)
-    dbg {pid1, pid2, pid3}
+    Scheduler.spawn_and_yield(fn ->
+      Enum.each(1..100, fn i ->
+        Scheduler.spawn_and_yield(fn ->
+          dbg {self(), "proc #{i}"}
+          Scheduler.yield(1_000_000)
+        end, true)
+      end)
+      Scheduler.yield(1000)
+      raise "foobar"
+    end, false)
 
-    Scheduler.yield(100)
+    Scheduler.yield(2000)
     dbg Scheduler.dump()
   end
 end