Fix concurrenty issue with global MC cache

HISTORY.md

@@ -1,3 +1,7 @@
+# 0.9.7
+
+Fix a concurrency bug, where Libtask would sometimes crash with a "Multiple concurrent writes to Dict detected!" error when TapedTasks were being executed concurrently.
+
 # 0.9.6
 
 Add support for Julia v1.12.

Project.toml

@@ -3,7 +3,7 @@ uuid = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 license = "MIT"
 desc = "Tape based task copying in Turing"
 repo = "https://github.com/TuringLang/Libtask.jl.git"
-version = "0.9.6"
+version = "0.9.7"
 
 [deps]
 MistyClosures = "dbe65cb8-6be2-42dd-bbc5-4196aaced4f4"

src/copyable_task.jl

@@ -113,7 +113,24 @@ struct CacheKey
     key::Any
 end
 
-const mc_cache = Dict{CacheKey,MistyClosure}()
+# A cache for holding MistyClosures, keyed by CacheKey.
+# The reason to have a type for this, rather than just a global const
+# Dict{CacheKey,MistyClosure}, is thread-safety. The global cache is written to when new
+# TapedTasks are compiled, and if multiple threads are doing this concurrently races might
+# occur. This type uses a ReentrantLock to ensure that any `setindex!` operations are
+# atomic, solving the problem.
+struct GlobalMCCache
+    cache::Dict{CacheKey,MistyClosure}
+    lock::ReentrantLock
+
+    GlobalMCCache() = new(Dict{CacheKey,MistyClosure}(), ReentrantLock())
+end
+
+Base.haskey(c::GlobalMCCache, key) = haskey(c.cache, key)
+Base.getindex(c::GlobalMCCache, key) = getindex(c.cache, key)
+Base.setindex!(c::GlobalMCCache, val, key) = @lock c.lock setindex!(c.cache, val, key)
+
+const mc_cache = GlobalMCCache()
 
 """
     TapedTask(taped_globals::Any, f, args...; kwargs...)