Replace lock() do-method with @lock

Should make the codebase more resistant to closure capture bugs.

Author: JamesWrigley

src/cluster.jl

@@ -157,7 +157,7 @@ mutable struct Worker
 end
 
 function set_worker_state(w, state)
-    lock(w.c_state) do
+    @lock w.c_state begin
         @atomic w.state = state
         notify(w.c_state; all=true)
     end
@@ -696,11 +696,9 @@ function create_worker(manager::ClusterManager, wconfig::WorkerConfig)
         # require the value of config.connect_at which is set only upon connection completion
         for jw in CTX[].pgrp.workers
             if (jw.id != 1) && (jw.id < w.id)
-                lock(jw.c_state) do
+                @lock jw.c_state if (@atomic jw.state) === WorkerState_created
                     # wait for wl to join
-                    if (@atomic jw.state) === WorkerState_created
-                        wait(jw.c_state)
-                    end
+                    wait(jw.c_state)
                 end
                 push!(join_list, jw)
             end
@@ -724,11 +722,9 @@ function create_worker(manager::ClusterManager, wconfig::WorkerConfig)
         end
 
         for wl in wlist
-            lock(wl.c_state) do
-                if (@atomic wl.state) === WorkerState_created
-                    # wait for wl to join
-                    wait(wl.c_state)
-                end
+            @lock wl.c_state if (@atomic wl.state) === WorkerState_created
+                # wait for wl to join
+                wait(wl.c_state)
             end
             push!(join_list, wl)
         end
@@ -749,9 +745,7 @@ function create_worker(manager::ClusterManager, wconfig::WorkerConfig)
     if timedwait(() -> isready(rr_ntfy_join), timeout) === :timed_out
         error("worker did not connect within $timeout seconds")
     end
-    lock(CTX[].client_refs) do
-        delete!(CTX[].pgrp.refs, ntfy_oid)
-    end
+    @lock CTX[].client_refs delete!(CTX[].pgrp.refs, ntfy_oid)
 
     return w.id
 end
@@ -1399,8 +1393,7 @@ function rmprocs(pids...; waitfor=typemax(Int), callback_timeout=10)
 end
 
 function _rmprocs(pids, waitfor, callback_timeout)
-    lock(CTX[].worker_lock)
-    try
+    @lock CTX[].worker_lock begin
         # Run the callbacks
         callback_tasks = Tuple{Any, Task}[]
         for pid in pids
@@ -1440,8 +1433,6 @@ function _rmprocs(pids, waitfor, callback_timeout)
             estr = string("rmprocs: pids ", unremoved, " not terminated after ", waitfor, " seconds.")
             throw(ErrorException(estr))
         end
-    finally
-        unlock(CTX[].worker_lock)
     end
 end
 
@@ -1547,7 +1538,7 @@ function deregister_worker(pg, pid)
     # delete this worker from our remote reference client sets
     ids = []
     tonotify = []
-    lock(CTX[].client_refs) do
+    @lock CTX[].client_refs begin
         for (id, rv) in pg.refs
             if in(pid, rv.clientset)
                 push!(ids, id)

src/messages.jl

@@ -174,8 +174,8 @@ function send_msg_(w::Worker, header, msg, now::Bool)
         wait(w.initialized)
     end
     io = w.w_stream
-    lock(io)
-    try
+
+    @lock io begin
         reset_state(w.w_serializer)
         serialize_hdr_raw(io, header)
         invokelatest(serialize_msg, w.w_serializer, msg)  # io is wrapped in w_serializer
@@ -186,8 +186,6 @@ function send_msg_(w::Worker, header, msg, now::Bool)
         else
             flush(io)
         end
-    finally
-        unlock(io)
     end
 end
 

src/remotecall.jl

@@ -160,9 +160,7 @@ A low-level API which returns the backing `AbstractChannel` for an `id` returned
 The call is valid only on the node where the backing channel exists.
 """
 function channel_from_id(id)
-    rv = lock(CTX[].client_refs) do
-        return get(CTX[].pgrp.refs, id, false)
-    end
+    rv = @lock CTX[].client_refs get(CTX[].pgrp.refs, id, false)
     if rv === false
         throw(ErrorException("Local instance of remote reference not found"))
     end
@@ -171,16 +169,18 @@ end
 
 lookup_ref(rrid::RRID, f=def_rv_channel) = lookup_ref(CTX[].pgrp, rrid, f)
 function lookup_ref(pg, rrid, f)
-    return lock(CTX[].client_refs) do
+    rv = @lock CTX[].client_refs begin
         rv = get(pg.refs, rrid, false)
         if rv === false
             # first we've heard of this ref
             rv = RemoteValue(invokelatest(f))
             pg.refs[rrid] = rv
             push!(rv.clientset, rrid.whence)
         end
-        return rv
+        rv
     end::RemoteValue
+
+    return rv
 end
 
 """
@@ -232,9 +232,7 @@ del_client(rr::AbstractRemoteRef) = del_client(remoteref_id(rr), myid())
 
 del_client(id, client) = del_client(CTX[].pgrp, id, client)
 function del_client(pg, id, client)
-    lock(CTX[].client_refs) do
-        _del_client(pg, id, client)
-    end
+    @lock CTX[].client_refs _del_client(pg, id, client)
     nothing
 end
 
@@ -297,13 +295,11 @@ function send_del_client_no_lock(rr)
 end
 
 function publish_del_msg!(w::Worker, msg)
-    lock(w.msg_lock) do
+    @lock w.msg_lock begin
         push!(w.del_msgs, msg)
         @atomic w.gcflag = true
     end
-    lock(CTX[].any_gc_flag) do
-        notify(CTX[].any_gc_flag)
-    end
+    @lock CTX[].any_gc_flag notify(CTX[].any_gc_flag)
 end
 
 function process_worker(rr)
@@ -320,7 +316,7 @@ function process_worker(rr)
 end
 
 function add_client(id, client)
-    lock(CTX[].client_refs) do
+    @lock CTX[].client_refs begin
         rv = lookup_ref(id)
         push!(rv.clientset, client)
     end
@@ -341,13 +337,11 @@ function send_add_client(rr::AbstractRemoteRef, i)
         # to the processor that owns the remote ref. it will add_client
         # itself inside deserialize().
         w = worker_from_id(rr.where)
-        lock(w.msg_lock) do
+        @lock w.msg_lock begin
             push!(w.add_msgs, (remoteref_id(rr), i))
             @atomic w.gcflag = true
         end
-        lock(CTX[].any_gc_flag) do
-            notify(CTX[].any_gc_flag)
-        end
+        @lock CTX[].any_gc_flag notify(CTX[].any_gc_flag)
     end
 end
 
@@ -448,9 +442,7 @@ function remotecall_fetch(f, w::Worker, args...; kwargs...)
     rv.waitingfor = w.id
     send_msg(w, MsgHeader(RRID(0,0), oid), CallMsg{:call_fetch}(f, args, kwargs))
     v = take!(rv)
-    lock(CTX[].client_refs) do
-        delete!(CTX[].pgrp.refs, oid)
-    end
+    @lock CTX[].client_refs delete!(CTX[].pgrp.refs, oid)
     return isa(v, RemoteException) ? throw(v) : v
 end
 
@@ -490,9 +482,7 @@ function remotecall_wait(f, w::Worker, args...; kwargs...)
     rr = Future(w)
     send_msg(w, MsgHeader(remoteref_id(rr), prid), CallWaitMsg(f, args, kwargs))
     v = fetch(rv.c)
-    lock(CTX[].client_refs) do
-        delete!(CTX[].pgrp.refs, prid)
-    end
+    @lock CTX[].client_refs delete!(CTX[].pgrp.refs, prid)
     isa(v, RemoteException) && throw(v)
     return rr
 end