From 5c79554b3f935bc9979b582f6ba5d5d6603e20b6 Mon Sep 17 00:00:00 2001
From: Shashi Gowda <shashi@juliacomputing.com>
Date: Mon, 29 Jan 2018 16:26:56 +0530
Subject: [PATCH] refcounting

---
 src/array/darray.jl | 38 ++++++++++++++++++++++++++++++++++++++
 src/chunks.jl       | 16 ++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/src/array/darray.jl b/src/array/darray.jl
index a36b497cc..91bea2c31 100644
--- a/src/array/darray.jl
+++ b/src/array/darray.jl
@@ -99,6 +99,44 @@ mutable struct DArray{T,N,F} <: ArrayOp{T, N}
     subdomains::AbstractArray{ArrayDomain{N}, N}
     chunks::AbstractArray{Union{Chunk,Thunk}, N}
     concat::F
+    freed::Bool
+    function DArray{T,N,F}(domain, subdomains, chunks, concat) where {T, N,F}
+        A = new(domain, subdomains, chunks, concat, false)
+        refcount_chunks(A.chunks)
+        finalizer(A, free!)
+        A
+    end
+end
+
+function refcount_chunks(chunks)
+    for c in chunks
+        if c isa Chunk{<:Any, DRef}
+            # increment refcount on the master node
+            addrefcount(c.handle, 1)
+        elseif c isa Thunk
+            refcount_chunks(c.inputs)
+        end
+    end
+end
+
+function free_chunks(chunks)
+    @sync for c in chunks
+        if c isa Chunk{<:Any, DRef}
+            # increment refcount on the master node
+            cnt = addrefcount(c.handle, -1)
+            cnt <= 0 && @async free!(c.handle)
+        elseif c isa Thunk
+            free_chunks(c.inputs)
+        end
+    end
+end
+
+function free!(x::DArray)
+    if !x.freed
+        @schedule Dagger.free_chunks(x.chunks)
+        x.freed = true
+    end
+    nothing
 end
 
 # mainly for backwards-compatibility
diff --git a/src/chunks.jl b/src/chunks.jl
index 503695410..1e29a5b1a 100644
--- a/src/chunks.jl
+++ b/src/chunks.jl
@@ -107,6 +107,7 @@ function free!(s::Chunk{X, DRef}; force=true, cache=false) where X
     end
 end
 free!(x; force=true,cache=false) = x # catch-all for non-chunks
+free!(x::DRef) = pooldelete(x)
 
 function savechunk(data, dir, f)
     sz = open(joinpath(dir, f), "w") do io
@@ -116,6 +117,21 @@ function savechunk(data, dir, f)
     Chunk(typeof(data), domain(data), FileRef(f, sz), true)
 end
 
+
+const refcount = Dict{MemPool.DRef, Int}()
+
+function addrefcount(r::DRef, x)
+    refcount[r] = getrefcount(r)+x
+end
+
+function getrefcount(r::DRef)
+    if haskey(refcount, r)
+        return refcount[r]
+    else
+        0
+    end
+end
+
 Base.@deprecate_binding AbstractPart Union{Chunk, Thunk}
 Base.@deprecate_binding Part Chunk
 Base.@deprecate parts(args...) chunks(args...)