From d488e95b0b78535f3f1cbe15e0b671b875310830 Mon Sep 17 00:00:00 2001
From: garrison <garrison@corporate.fm>
Date: Sun, 1 Feb 2026 17:17:18 -0500
Subject: [PATCH] Add trap_exit flag

---
 lib/trinity/scheduler.ex   | 70 ++++++++++++++++++++++++--------------
 lib/trinity/sim_process.ex | 20 +++++++++--
 2 files changed, 62 insertions(+), 28 deletions(-)

diff --git a/lib/trinity/scheduler.ex b/lib/trinity/scheduler.ex
index 76c37a6..26cc190 100644
--- a/lib/trinity/scheduler.ex
+++ b/lib/trinity/scheduler.ex
@@ -44,6 +44,17 @@ defmodule Trinity.Scheduler do
     defstruct @enforce_keys
   end
 
+  defmodule ProcInfo do
+    @type t :: %__MODULE__{
+      trap_exit: boolean,
+      links: [pid],
+    }
+    defstruct [
+      trap_exit: false,
+      links: [],
+    ]
+  end
+
   defmacro simulation_key, do: :_trinity_simulation
   defmacro sim_node_key, do: :_trinity_simulation_node
 
@@ -194,22 +205,22 @@ defmodule Trinity.Scheduler do
     %Simulation{proc_links: proc_links} = get_sim()
     from_pid = self()
 
-    from_links =
+    from_info =
       case :ets.lookup(proc_links, from_pid) do
-        [{^from_pid, links}] -> links
-        _ -> []
+        [{^from_pid, info}] -> info
+        _ -> %ProcInfo{}
       end
 
-    to_links =
+    to_info =
       case :ets.lookup(proc_links, to_pid) do
-        [{^to_pid, links}] -> links
-        _ -> []
+        [{^to_pid, info}] -> info
+        _ -> %ProcInfo{}
       end
 
-    case to_pid in from_links do
+    case to_pid in from_info.links do
       false ->
-        :ets.insert(proc_links, {from_pid, [to_pid | from_links]})
-        :ets.insert(proc_links, {to_pid, [from_pid | to_links]})
+        :ets.insert(proc_links, {from_pid, %{from_info | links: [to_pid | from_info.links]}})
+        :ets.insert(proc_links, {to_pid, %{to_info | links: [from_pid | to_info.links]}})
       _ -> :noop
     end
 
@@ -235,27 +246,43 @@ defmodule Trinity.Scheduler do
   end
 
   def handle_down(pid, reason) do
+    # Should always be converted to :killed by the time it reaches the supervisor
+    assert reason != :kill
+
     sim = get_sim()
     destroy_or_message_linked(sim, pid, reason, MapSet.new([pid]))
+    destroy_process(sim, pid)
     :ok
   end
 
   defp destroy_or_message_linked(%Simulation{} = sim, for_pid, reason, visited) do
-    linked = get_linked(sim.proc_links, for_pid)
-    linked = Enum.reject(linked, &MapSet.member?(visited, &1))
+    linked =
+      case :ets.lookup(sim.proc_links, for_pid) do
+        [{_key, %ProcInfo{links: links}}] -> links
+        [] -> []
+      end
+      |> Enum.reject(&MapSet.member?(visited, &1))
 
     visited = Enum.reduce(linked, visited, fn p, visited -> MapSet.put(visited, p) end)
     Enum.reduce(linked, visited, fn p, visited ->
       visited = destroy_or_message_linked(sim, p, reason, visited)
-      destroy_or_message(sim, p, reason)
+      destroy_or_message(sim, p, for_pid, reason)
       visited
     end)
   end
 
-  defp destroy_or_message(%Simulation{} = sim, pid, reason) do
-    # TODO: trap_exit
-    trap_exit = false
+  defp destroy_or_message(%Simulation{} = sim, pid, from_pid, reason) do
+    trap_exit =
+      case :ets.lookup(sim.proc_links, pid) do
+        [{_key, %ProcInfo{trap_exit: trap_exit}}] -> trap_exit
+        [] -> false
+      end
+
     case trap_exit do
+      true ->
+        send pid, {:EXIT, from_pid, reason}
+        handle_sent(pid)
+
       false ->
         destroy_process(sim, pid)
 
@@ -267,13 +294,6 @@ defmodule Trinity.Scheduler do
     end
   end
 
-  defp get_linked(proc_links, pid) do
-    case :ets.lookup(proc_links, pid) do
-      [{_pid, links}] -> links
-      [] -> []
-    end
-  end
-
   def handle_sent(dest_pid) when is_pid(dest_pid) do
     %{queue: queue, proc_queue_keys: proc_queue_keys, now: now} = get_sim()
 
@@ -395,11 +415,11 @@ defmodule Trinity.Scheduler do
 
   defp destroy_links(proc_links, from_pid) do
     case :ets.lookup(proc_links, from_pid) do
-      [{^from_pid, from_links}] ->
+      [{^from_pid, %ProcInfo{links: 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)})
+          [{^to_pid, %ProcInfo{} = to_info}] = :ets.lookup(proc_links, to_pid)
+          :ets.insert(proc_links, {to_pid, %{to_info | links: List.delete(to_info.links, from_pid)}})
         end)
 
         # Delete `from_pid`'s links entry
diff --git a/lib/trinity/sim_process.ex b/lib/trinity/sim_process.ex
index 2c82948..f14d850 100644
--- a/lib/trinity/sim_process.ex
+++ b/lib/trinity/sim_process.ex
@@ -1,6 +1,6 @@
 defmodule Trinity.SimProcess do
   alias Trinity.Scheduler
-  alias Trinity.Scheduler.Simulation
+  alias Trinity.Scheduler.{Simulation, ProcInfo}
 
   import Trinity.Scheduler, only: [simulation_key: 0, sim_node_key: 0]
 
@@ -55,8 +55,7 @@ defmodule Trinity.SimProcess do
   def flag(flag, value) do
     case get_sim() do
       nil -> Process.flag(flag, value)
-      # TODO
-      _sim -> :noop
+      _sim -> sim_flag(flag, value)
     end
   end
 
@@ -136,6 +135,21 @@ defmodule Trinity.SimProcess do
     nil
   end
 
+  defp sim_flag(:trap_exit = _flag, value) when is_boolean(value) do
+    %{proc_links: proc_links} = get_sim()
+    pid = self()
+
+    %ProcInfo{} = info =
+      case :ets.lookup(proc_links, pid) do
+        [{_key, info}] -> info
+        [] -> %ProcInfo{}
+      end
+    old_value = info.trap_exit
+
+    :ets.insert(proc_links, {pid, %{info | trap_exit: value}})
+    old_value
+  end
+
   defp sim_alias do
     %Simulation{proc_aliases: proc_aliases} = get_sim()
 
-- 
2.51.2