From 0fab5cb74b52cc8be190b5e402e8d80cd4f19816 Mon Sep 17 00:00:00 2001
From: Chris Foster <chris42f@gmail.com>
Date: Tue, 15 Sep 2020 11:17:54 +1000
Subject: [PATCH] Add global lock around the caching mechanism

---
 src/RuntimeGeneratedFunctions.jl | 63 +++++++++++++++++++-------------
 test/runtests.jl                 | 19 ++++++++++
 2 files changed, 56 insertions(+), 26 deletions(-)

diff --git a/src/RuntimeGeneratedFunctions.jl b/src/RuntimeGeneratedFunctions.jl
index 61a7b1f..3da68eb 100644
--- a/src/RuntimeGeneratedFunctions.jl
+++ b/src/RuntimeGeneratedFunctions.jl
@@ -84,42 +84,53 @@ end
 # little non-robust to weird special cases like Main.eval being
 # Base.MainInclude.eval.)
 
+# It appears we can't use a ReentrantLock here, as contention seems to lead to
+# deadlock. Perhaps because it triggers a task switch while compiling the
+# @generated function.
+_cache_lock = Threads.SpinLock()
 _cachename = Symbol("#_RuntimeGeneratedFunctions_cache")
 _tagname = Symbol("#_RuntimeGeneratedFunctions_ModTag")
 
 function _cache_body(moduletag, id, body)
-    cache = getfield(parentmodule(moduletag), _cachename)
-    # Caching is tricky when `id` is the same for different AST instances:
-    #
-    # Tricky case #1: If a function body with the same `id` was cached
-    # previously, we need to use that older instance of the body AST as the
-    # canonical one rather than `body`. This ensures the lifetime of the
-    # body in the cache will always cover the lifetime of the parent
-    # `RuntimeGeneratedFunction`s when they share the same `id`.
-    #
-    # Tricky case #2: Unless we hold a separate reference to
-    # `cache[id].value`, the GC can collect it (causing it to become
-    # `nothing`). So root it in a local variable first.
-    #
-    cached_body = haskey(cache, id) ? cache[id].value : nothing
-    cached_body = cached_body !== nothing ? cached_body : body
-    # Use a WeakRef to allow `body` to be garbage collected. (After GC, the
-    # cache will still contain an empty entry with key `id`.)
-    cache[id] = WeakRef(cached_body)
-    return cached_body
+    lock(_cache_lock) do
+        cache = getfield(parentmodule(moduletag), _cachename)
+        # Caching is tricky when `id` is the same for different AST instances:
+        #
+        # Tricky case #1: If a function body with the same `id` was cached
+        # previously, we need to use that older instance of the body AST as the
+        # canonical one rather than `body`. This ensures the lifetime of the
+        # body in the cache will always cover the lifetime of the parent
+        # `RuntimeGeneratedFunction`s when they share the same `id`.
+        #
+        # Tricky case #2: Unless we hold a separate reference to
+        # `cache[id].value`, the GC can collect it (causing it to become
+        # `nothing`). So root it in a local variable first.
+        #
+        cached_body = haskey(cache, id) ? cache[id].value : nothing
+        cached_body = cached_body !== nothing ? cached_body : body
+        # Use a WeakRef to allow `body` to be garbage collected. (After GC, the
+        # cache will still contain an empty entry with key `id`.)
+        cache[id] = WeakRef(cached_body)
+        return cached_body
+    end
 end
 
 function _lookup_body(moduletag, id)
-    getfield(parentmodule(moduletag), _cachename)[id].value
+    lock(_cache_lock) do
+        cache = getfield(parentmodule(moduletag), _cachename)
+        cache[id].value
+    end
 end
 
 function _ensure_cache_exists!(mod)
-    if !isdefined(mod, _cachename)
-        mod.eval(quote
-            const $_cachename = Dict()
-            struct $_tagname
-            end
-        end)
+    lock(_cache_lock) do
+        if !isdefined(mod, _cachename)
+            mod.eval(quote
+                const $_cachename = Dict()
+                struct $_tagname
+                end
+            end)
+        end
     end
 end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 508fabe..f461b4d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -88,3 +88,22 @@ let
 end
 GC.gc()
 @test f_gc(1,-1) == 100001
+
+# Test that threaded use works
+tasks = []
+for k=1:4
+    let k=k
+        t = Threads.@spawn begin
+            r = Bool[]
+            for i=1:100
+                f = @RuntimeGeneratedFunction(Base.remove_linenums!(:((x,y)->x+y+$i*$k)))
+                x = 1; y = 2;
+                push!(r, f(x,y) == x + y + i*k)
+            end
+            r
+        end
+        push!(tasks, t)
+    end
+end
+@test all(all.(fetch.(tasks)))
+