From ade85ec9c08b80c97072d4003cedbd413a913387 Mon Sep 17 00:00:00 2001
From: Martin Rehfeld <martin.rehfeld@glnetworks.de>
Date: Tue, 29 Nov 2011 23:18:03 +0100
Subject: [PATCH] trap exits from user processes to make the mutex safe

---
 apps/warmup-semaphore/src/mutex.erl        | 31 ++++++++++++++--
 apps/warmup-semaphore/test/mutex_tests.erl | 43 +++++++++++++++++++++-
 2 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/apps/warmup-semaphore/src/mutex.erl b/apps/warmup-semaphore/src/mutex.erl
index 3d3bf28..581666d 100644
--- a/apps/warmup-semaphore/src/mutex.erl
+++ b/apps/warmup-semaphore/src/mutex.erl
@@ -7,7 +7,7 @@
 
 -export([start/0, stop/0, wait/0, signal/0]).
 
--record(state, {semaphore}).
+-record(state, {semaphore, holder=undefined}).
 
 %% @doc mutex:start() ⇒ ok.
 start() ->
@@ -39,17 +39,40 @@ signal() ->
 
 %% @private
 init() ->
+    process_flag(trap_exit, true),
     loop(#state{semaphore=free}).
 
 %% @private
-loop(#state{semaphore=Semaphore}=S) ->
+loop(#state{semaphore=Semaphore, holder=Holder}=S) ->
     receive
         {Pid, MsgRef, wait} when Semaphore =:= free ->
+            catch link(Pid),
             Pid ! {MsgRef, ok},
-            loop(#state{semaphore=busy});
+            loop(S#state{semaphore=busy, holder=Pid});
         {Pid, MsgRef, signal} when Semaphore =:= busy ->
             Pid ! {MsgRef, ok},
-            loop(#state{semaphore=free});
+            unlink(Pid),
+            %% clean any pending exit messages for Pid from mailbox
+            receive
+                {'EXIT', Pid, _} -> true
+            after 0 ->
+                true
+            end,
+            loop(S#state{semaphore=free, holder=undefined});
+
+        {'EXIT', FromPid, _Reason} ->
+            %% clean any pending messages from Pid from mailbox
+            receive
+                {FromPid, _, wait}   -> true;
+                {FromPid, _, signal} -> true
+            after 0 ->
+                true
+            end,
+            %% free the semaphore when exit came from the current holder
+            case FromPid of
+                Holder -> loop(S#state{semaphore=free});
+                _      -> loop(S)
+            end;
 
         shutdown ->
             exit(shutdown);
diff --git a/apps/warmup-semaphore/test/mutex_tests.erl b/apps/warmup-semaphore/test/mutex_tests.erl
index 544db0f..831305e 100644
--- a/apps/warmup-semaphore/test/mutex_tests.erl
+++ b/apps/warmup-semaphore/test/mutex_tests.erl
@@ -15,7 +15,9 @@ main_test_() ->
      fun setup/0,
      fun cleanup/1,
      [
-      fun test_mutex/1
+      fun test_mutex/1,
+      fun test_semaphore_holder_terminates/1,
+      fun test_waiting_process_terminates/1
      ]}.
 
 %% tests with started mutex
@@ -62,13 +64,50 @@ test_mutex(_) ->
      ?_assert(timer:now_diff(MutexEndA, MutexWaitB) > 0),
      ?_assert(timer:now_diff(MutexStartB, MutexEndA) > 0)].
 
+%% What happens if a process that currently holds the semaphore terminates prior to releasing it?
+test_semaphore_holder_terminates(_) ->
+    PidA = spawn(fun () ->
+            ok = mutex:wait(),
+            timer:sleep(100)
+        end),
+
+    PidB = spawn(fun () ->
+            timer:sleep(200),
+            ok = mutex:wait(),
+            ok = mutex:signal()
+        end),
+
+    ExitReasonPidA = wait_for_exit(PidA),
+    ExitReasonPidB = wait_for_exit(PidB),
+
+    [?_assertEqual(normal, ExitReasonPidA),
+     ?_assertMatch(normal, ExitReasonPidB)].
+
+%% Or what happens if a process waiting to execute is terminated due to an exit signal?
+test_waiting_process_terminates(_) ->
+    PidA = spawn(fun () ->
+            ok = mutex:wait(),
+            timer:sleep(100),
+            ok = mutex:signal()
+        end),
+
+    PidB = spawn(fun () ->
+            ok = mutex:wait()
+        end),
+
+    exit(PidB, kill),
+
+    ExitReasonPidA = wait_for_exit(PidA),
+
+    [?_assertEqual(normal, ExitReasonPidA)].
+
 %% helper
 
 setup() ->
     mutex:start().
 
 cleanup(Pid) ->
-    mutex:stop(),
+    catch mutex:stop(),
     wait_for_exit(Pid),
     timer:sleep(100).