Merge pull request #149 from JuliaParallel/amitm/darr_leak

fix leak - make registry entry a WeakRef on the node releasing the darray

Author: andreasnoack

src/core.jl

@@ -44,7 +44,7 @@ function d_closeall()
     crefs = copy(refs)
     for id in crefs
         if id[1] ==  myid() # sanity check
-            haskey(registry, id) && close(registry[id])
+            haskey(registry, id) && close(d_from_weakref_or_d(id))
             yield()
         end
     end

src/darray.jl

@@ -39,12 +39,14 @@ type DArray{T,N,A} <: AbstractArray{T,N}
         end
         release = (myid() == id[1])
 
-        haskey(registry, id) && return registry[id]
+        d = d_from_weakref_or_d(id)
+        if d === nothing
+            d = new(id, dims, pids, indexes, cuts, lp, release)
+        end
 
-        d = new(id, dims, pids, indexes, cuts, lp, release)
         if release
             push!(refs, id)
-            registry[id] = d
+            registry[id] = WeakRef(d)
 
 #            println("Installing finalizer for : ", d.id, ", : ", object_id(d), ", isbits: ", isbits(d))
             finalizer(d, close)
@@ -55,6 +57,12 @@ type DArray{T,N,A} <: AbstractArray{T,N}
     DArray{T,N,A}() where {T,N,A} = new()
 end
 
+function d_from_weakref_or_d(id)
+    d = get(registry, id, nothing)
+    isa(d, WeakRef) && return d.value
+    return d
+end
+
 eltype{T}(::Type{DArray{T}}) = T
 empty_localpart(T,N,A) = A(Array{T}(ntuple(zero, N)))
 
@@ -90,6 +98,7 @@ function DArray(id, init, dims, pids, idxs, cuts)
     A = take!(r)
     if myid() in pids
         d = registry[id]
+        d = isa(d, WeakRef) ? d.value : d
     else
         T = eltype(A)
         N = length(dims)
@@ -137,6 +146,7 @@ function ddata(;T::Type=Any, init::Function=I->nothing, pids=workers(), data::Ve
 
     if myid() in pids
         d = registry[id]
+        d = isa(d, WeakRef) ? d.value : d
     else
         d = DArray{T,1,T}(id, (npids,), pids, idxs, cuts, Nullable{T}())
     end
@@ -320,7 +330,7 @@ function localpart{T,N,A}(d::DArray{T,N,A})
         return empty_localpart(T,N,A)::A
     end
 
-    return get(registry[d.id].localpart)::A
+    return get(d.localpart)::A
 end
 
 localpart(d::DArray, localidx...) = localpart(d)[localidx...]

src/serialize.jl

@@ -20,15 +20,14 @@ function Base.deserialize{DT<:DArray}(S::AbstractSerializer, t::Type{DT})
     id = what[2]
 
     if id_only
-        if haskey(registry, id)
-            return registry[id]
-        else
+        d = d_from_weakref_or_d(id)
+        if d === nothing
             # access to fields will throw an error, at least the deserialization process will not
             # result in worker death
             d = DT()
             d.id = id
-            return d
         end
+        return d
     else
         # We are not a participating worker, deser fields and instantiate locally.
         dims = deserialize(S)