fix deadlock

Author: sbrantq

src/Types.jl

@@ -215,6 +215,7 @@ function ConcretePJRTArray(
 end
 
 Base.wait(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = foreach(wait, x.data)
+Base.isready(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = all(isready, x.data)
 XLA.client(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = XLA.client(x.data)
 function XLA.device(x::Union{ConcretePJRTArray,ConcretePJRTNumber})
     x.sharding isa Sharding.NoShardInfo && return XLA.device(only(x.data))
@@ -412,6 +413,7 @@ function ConcreteIFRTArray(
 end
 
 Base.wait(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber}) = wait(x.data)
+Base.isready(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber}) = isready(x.data)
 XLA.client(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber}) = XLA.client(x.data)
 function XLA.device(x::Union{ConcreteIFRTArray,ConcreteIFRTNumber})
     return XLA.device(x.data)

src/probprog/Modeling.jl

@@ -199,11 +199,13 @@ function simulate(rng::AbstractRNG, f::Function, args::Vararg{Any,Nargs}) where
 
     compiled_fn = @compile optimize = :probprog simulate_internal(rng, f, args...)
 
-    old_gc_state = GC.enable(false)
-    try
+    seed_buffer = only(rng.seed.data).buffer
+    GC.@preserve seed_buffer begin
         trace, _, _ = compiled_fn(rng, f, args...)
-    finally
-        GC.enable(old_gc_state)
+
+        while !isready(trace)
+            yield()
+        end
     end
 
     trace = unsafe_pointer_to_objref(Ptr{Any}(Array(trace)[1]))
@@ -278,11 +280,13 @@ function generate(
 
     compiled_fn = @compile optimize = :probprog wrapper_fn(rng, constraint_ptr, args...)
 
-    old_gc_state = GC.enable(false)
-    try
+    seed_buffer = only(rng.seed.data).buffer
+    GC.@preserve seed_buffer constraint begin
         trace, _, _ = compiled_fn(rng, constraint_ptr, args...)
-    finally
-        GC.enable(old_gc_state)
+
+        while !isready(trace)
+            yield()
+        end
     end
 
     trace = unsafe_pointer_to_objref(Ptr{Any}(Array(trace)[1]))