From 80e2cd5c10e25d18268d264152cc3c09953c2b93 Mon Sep 17 00:00:00 2001
From: Amit Murthy <amit.murthy@gmail.com>
Date: Thu, 19 Mar 2015 11:04:04 +0530
Subject: [PATCH] rework

---
 base/util.jl | 66 +++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 45 insertions(+), 21 deletions(-)

diff --git a/base/util.jl b/base/util.jl
index e8621cf558c06..7b0ed5c6e302f 100644
--- a/base/util.jl
+++ b/base/util.jl
@@ -269,36 +269,60 @@ end
 
 julia_exename() = ccall(:jl_is_debugbuild,Cint,())==0 ? "julia" : "julia-debug"
 
+# Advisory reentrant lock
+type ReentrantLock
+    locked_by::Nullable{Task}
+    cond_wait::Condition
+    reentrancy_cnt::Int
+
+    ReentrantLock() = new(nothing, Condition(), 0)
+end
+
 # Lock object during function execution. Recursive calls by the same task is OK.
-const LOCK_MODE_ADVISORY=1
-const adv_locks_map = Dict{UInt64, Tuple}()
-function lock(f::Function, o::Any; mode=LOCK_MODE_ADVISORY)
-    # Currently only advisory locks are supported.
-    t = current_task()
-    release_lock = true
-    oid = object_id(o)
-    while true
-        if !haskey(adv_locks_map, oid)
-            adv_locks_map[oid] = (t, Condition()); break
-        else
-            (locked_by, c) = adv_locks_map[oid]
-            if t == locked_by
-                release_lock = false; break
-            end
-            wait(c)
-        end
+const adv_locks_map = WeakKeyDict{Any, ReentrantLock}()
+function lock(f::Function, o::Any)
+    rl = get(adv_locks_map, o, nothing)
+    if rl == nothing
+        rl = ReentrantLock()
+        adv_locks_map[o] = rl
     end
+    lock(rl)
 
     try
         f(o)
     finally
-        if release_lock
-            (_, c) = pop!(adv_locks_map, oid)
-            notify(c)
-        end
+        unlock(o)
     end
 end
 
+function lock(rl::ReentrantLock)
+    t = current_task()
+    while true
+        if rl.reentrancy_cnt == 0
+            rl.locked_by = t
+            rl.reentrancy_cnt = 1
+            return
+        elseif t == get(rl.locked_by)
+            rl.reentrancy_cnt += 1
+            return
+        end
+        wait(rl.cond_wait)
+    end
+end
 
+function unlock(o::Any)
+    rl = adv_locks_map[o]
+    unlock(rl)
+end
 
+function unlock(rl::ReentrantLock)
+    rl.reentrancy_cnt = rl.reentrancy_cnt - 1
+    if rl.reentrancy_cnt == 0
+        rl.locked_by = nothing
+        notify(rl.cond_wait)
+    elseif rl.reentrancy_cnt < 0
+        AssertionError("unlock count must match lock count")
+    end
+    rl
+end