fix: lock various methods

Author: MilesCranmer

src/Convert/pyconvert.jl

@@ -73,7 +73,7 @@ function pyconvert_add_rule(
         get!(Vector{PyConvertRule}, PYCONVERT_RULES, pytypename),
         PyConvertRule(type, func, priority),
     )
-    empty!.(values(PYCONVERT_RULES_CACHE))
+    Base.@lock PYCONVERT_RULES_CACHE_LOCK empty!.(values(PYCONVERT_RULES_CACHE))
     return
 end
 
@@ -262,9 +262,10 @@ function _pyconvert_get_rules(pytype::Py)
 end
 
 const PYCONVERT_PREFERRED_TYPE = Dict{Py,Type}()
+const PYCONVERT_PREFERRED_TYPE_LOCK = Threads.SpinLock()
 
 pyconvert_preferred_type(pytype::Py) =
-    get!(PYCONVERT_PREFERRED_TYPE, pytype) do
+    Base.@lock PYCONVERT_PREFERRED_TYPE_LOCK get!(PYCONVERT_PREFERRED_TYPE, pytype) do
         if pyissubclass(pytype, pybuiltins.int)
             Union{Int,BigInt}
         else
@@ -308,9 +309,10 @@ end
 pyconvert_fix(::Type{T}, func) where {T} = x -> func(T, x)
 
 const PYCONVERT_RULES_CACHE = Dict{Type,Dict{C.PyPtr,Vector{Function}}}()
+const PYCONVERT_RULES_CACHE_LOCK = Threads.SpinLock()
 
 @generated pyconvert_rules_cache(::Type{T}) where {T} =
-    get!(Dict{C.PyPtr,Vector{Function}}, PYCONVERT_RULES_CACHE, T)
+    Base.@lock PYCONVERT_RULES_CACHE_LOCK get!(Dict{C.PyPtr,Vector{Function}}, PYCONVERT_RULES_CACHE, T)
 
 function pyconvert_rule_fast(::Type{T}, x::Py) where {T}
     if T isa Union
@@ -352,7 +354,7 @@ function pytryconvert(::Type{T}, x_) where {T}
     # TODO: we should hold weak references and clear the cache if types get deleted
     tptr = C.Py_Type(x)
     trules = pyconvert_rules_cache(T)
-    rules = get!(trules, tptr) do
+    rules = Base.@lock PYCONVERT_RULES_CACHE_LOCK get!(trules, tptr) do
         t = pynew(incref(tptr))
         ans = pyconvert_get_rules(T, t)::Vector{Function}
         pydel!(t)

src/Core/Py.jl

@@ -57,6 +57,7 @@ decref(x::Py) = Base.GC.@preserve x (decref(getptr(x)); x)
 Base.unsafe_convert(::Type{C.PyPtr}, x::Py) = getptr(x)
 
 const PYNULL_CACHE = Py[]
+const PYNULL_CACHE_LOCK = Threads.SpinLock()
 
 """
     pynew([ptr])
@@ -69,12 +70,13 @@ points at, i.e. the new `Py` object owns a reference.
 Note that NULL Python objects are not safe in the sense that most API functions will probably
 crash your Julia session if you pass a NULL argument.
 """
-pynew() =
+pynew() = Base.@lock PYNULL_CACHE_LOCK begin
     if isempty(PYNULL_CACHE)
         Py(Val(:new), C.PyNULL)
     else
         pop!(PYNULL_CACHE)
     end
+end
 
 const PyNULL = pynew()
 
@@ -119,7 +121,7 @@ function pydel!(x::Py)
         C.Py_DecRef(ptr)
         setptr!(x, C.PyNULL)
     end
-    push!(PYNULL_CACHE, x)
+    Base.@lock PYNULL_CACHE_LOCK push!(PYNULL_CACHE, x)
     return
 end
 

src/JlWrap/C.jl

@@ -19,6 +19,8 @@ const PyJuliaBase_Type = Ref(C.PyNULL)
 const PYJLVALUES = []
 # unused indices in PYJLVALUES
 const PYJLFREEVALUES = Int[]
+# Thread safety for PYJLVALUES and PYJLFREEVALUES
+const PYJLVALUES_LOCK = Threads.SpinLock()
 
 function _pyjl_new(t::C.PyPtr, ::C.PyPtr, ::C.PyPtr)
     o = ccall(UnsafePtr{C.PyTypeObject}(t).alloc[!], C.PyPtr, (C.PyPtr, C.Py_ssize_t), t, 0)
@@ -31,20 +33,25 @@ end
 function _pyjl_dealloc(o::C.PyPtr)
     idx = UnsafePtr{PyJuliaValueObject}(o).value[]
     if idx != 0
-        PYJLVALUES[idx] = nothing
-        push!(PYJLFREEVALUES, idx)
+        Base.@lock PYJLVALUES_LOCK begin
+            PYJLVALUES[idx] = nothing
+            push!(PYJLFREEVALUES, idx)
+        end
     end
     UnsafePtr{PyJuliaValueObject}(o).weaklist[!] == C.PyNULL || C.PyObject_ClearWeakRefs(o)
     ccall(UnsafePtr{C.PyTypeObject}(C.Py_Type(o)).free[!], Cvoid, (C.PyPtr,), o)
     nothing
 end
 
 const PYJLMETHODS = Vector{Any}()
+const PYJLMETHODS_LOCK = Threads.SpinLock()
 
 function PyJulia_MethodNum(f)
     @nospecialize f
-    push!(PYJLMETHODS, f)
-    return length(PYJLMETHODS)
+    Base.@lock PYJLMETHODS_LOCK begin
+        push!(PYJLMETHODS, f)
+        return length(PYJLMETHODS)
+    end
 end
 
 function _pyjl_isnull(o::C.PyPtr, ::C.PyPtr)
@@ -58,12 +65,13 @@ function _pyjl_callmethod(o::C.PyPtr, args::C.PyPtr)
     @assert nargs > 0
     num = C.PyLong_AsLongLong(C.PyTuple_GetItem(args, 0))
     num == -1 && return C.PyNULL
-    f = PYJLMETHODS[num]
+    f = Base.@lock PYJLMETHODS_LOCK PYJLMETHODS[num]
     # this form gets defined in jlwrap/base.jl
     return _pyjl_callmethod(f, o, args, nargs)::C.PyPtr
 end
 
 const PYJLBUFCACHE = Dict{Ptr{Cvoid},Any}()
+const PYJLBUFCACHE_LOCK = Threads.SpinLock()
 
 @kwdef struct PyBufferInfo{N}
     # data
@@ -177,7 +185,9 @@ function _pyjl_get_buffer_impl(
 
     # internal
     cptr = Base.pointer_from_objref(c)
-    PYJLBUFCACHE[cptr] = c
+    Base.@lock PYJLBUFCACHE_LOCK begin
+        PYJLBUFCACHE[cptr] = c
+    end
     b.internal[] = cptr
 
     # obj
@@ -195,7 +205,7 @@ function _pyjl_get_buffer(o::C.PyPtr, buf::Ptr{C.Py_buffer}, flags::Cint)
     C.Py_DecRef(num_)
     num == -1 && return Cint(-1)
     try
-        f = PYJLMETHODS[num]
+        f = Base.@lock PYJLMETHODS_LOCK PYJLMETHODS[num]
         x = PyJuliaValue_GetValue(o)
         return _pyjl_get_buffer_impl(o, buf, flags, x, f)::Cint
     catch exc
@@ -209,7 +219,9 @@ function _pyjl_get_buffer(o::C.PyPtr, buf::Ptr{C.Py_buffer}, flags::Cint)
 end
 
 function _pyjl_release_buffer(xo::C.PyPtr, buf::Ptr{C.Py_buffer})
-    delete!(PYJLBUFCACHE, UnsafePtr(buf).internal[!])
+    Base.@lock PYJLBUFCACHE_LOCK begin
+        delete!(PYJLBUFCACHE, UnsafePtr(buf).internal[!])
+    end
     nothing
 end
 
@@ -339,22 +351,31 @@ end
 
 PyJuliaValue_IsNull(o) = Base.GC.@preserve o UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[] == 0
 
-PyJuliaValue_GetValue(o) = Base.GC.@preserve o PYJLVALUES[UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[]]
+PyJuliaValue_GetValue(o) = Base.GC.@preserve o begin
+    idx = UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[]
+    Base.@lock PYJLVALUES_LOCK begin
+        PYJLVALUES[idx]
+    end
+end
 
 PyJuliaValue_SetValue(_o, @nospecialize(v)) = Base.GC.@preserve _o begin
     o = C.asptr(_o)
     idx = UnsafePtr{PyJuliaValueObject}(o).value[]
     if idx == 0
-        if isempty(PYJLFREEVALUES)
-            push!(PYJLVALUES, v)
-            idx = length(PYJLVALUES)
-        else
-            idx = pop!(PYJLFREEVALUES)
-            PYJLVALUES[idx] = v
+        Base.@lock PYJLVALUES_LOCK begin
+            if isempty(PYJLFREEVALUES)
+                push!(PYJLVALUES, v)
+                idx = length(PYJLVALUES)
+            else
+                idx = pop!(PYJLFREEVALUES)
+                PYJLVALUES[idx] = v
+            end
         end
         UnsafePtr{PyJuliaValueObject}(o).value[] = idx
     else
-        PYJLVALUES[idx] = v
+        Base.@lock PYJLVALUES_LOCK begin
+            PYJLVALUES[idx] = v
+        end
     end
     nothing
 end

src/JlWrap/base.jl

@@ -84,6 +84,7 @@ function Cjl._pyjl_callmethod(f, self_::C.PyPtr, args_::C.PyPtr, nargs::C.Py_ssi
                 pybuiltins.NotImplementedError,
                 "__jl_callmethod not implemented for this many arguments",
             )
+            return C.PyNULL
         end
         return getptr(incref(ans))
     catch exc