switch to Lockable

Author: MilesCranmer

pysrc/juliacall/juliapkg.json

@@ -3,7 +3,9 @@
     "packages": {
         "PythonCall": {
             "uuid": "6099a3de-0909-46bc-b1f4-468b9a2dfc0d",
-            "version": "=0.9.25"
+            "version": "=0.9.25",
+            "path": "/Users/mcranmer/PermaDocuments/PythonCall.jl/",
+            "dev": true
         },
         "OpenSSL_jll": {
             "uuid": "458c3c95-2e84-50aa-8efc-19380b2a3a95",

src/Convert/Convert.jl

@@ -9,6 +9,7 @@ using ..Core
 using ..Core:
     C,
     Utils,
+    Lockable,
     @autopy,
     getptr,
     incref,

src/Convert/pyconvert.jl

@@ -12,8 +12,8 @@ struct PyConvertRule
     priority::PyConvertPriority
 end
 
-const PYCONVERT_RULES = Dict{String,Vector{PyConvertRule}}()
-const PYCONVERT_EXTRATYPES = Py[]
+const PYCONVERT_RULES = Lockable(Dict{String,Vector{PyConvertRule}}())
+const PYCONVERT_EXTRATYPES = Lockable(Py[])
 
 """
     pyconvert_add_rule(tname::String, T::Type, func::Function, priority::PyConvertPriority=PYCONVERT_PRIORITY_NORMAL)
@@ -69,11 +69,11 @@ function pyconvert_add_rule(
     priority::PyConvertPriority = PYCONVERT_PRIORITY_NORMAL,
 )
     @nospecialize type func
-    push!(
-        get!(Vector{PyConvertRule}, PYCONVERT_RULES, pytypename),
+    Base.@lock PYCONVERT_RULES push!(
+        get!(Vector{PyConvertRule}, PYCONVERT_RULES[], pytypename),
         PyConvertRule(type, func, priority),
     )
-    Base.@lock PYCONVERT_RULES_CACHE_LOCK empty!.(values(PYCONVERT_RULES_CACHE))
+    Base.@lock PYCONVERT_RULES_CACHE empty!.(values(PYCONVERT_RULES_CACHE[]))
     return
 end
 
@@ -163,7 +163,7 @@ function _pyconvert_get_rules(pytype::Py)
     omro = collect(pytype.__mro__)
     basetypes = Py[pytype]
     basemros = Vector{Py}[omro]
-    for xtype in PYCONVERT_EXTRATYPES
+    Base.@lock PYCONVERT_EXTRATYPES for xtype in PYCONVERT_EXTRATYPES[]
         # find the topmost supertype of
         xbase = PyNULL
         for base in omro
@@ -248,9 +248,9 @@ function _pyconvert_get_rules(pytype::Py)
     mro = String[x for xs in xmro for x in xs]
 
     # get corresponding rules
-    rules = PyConvertRule[
+    rules = Base.@lock PYCONVERT_RULES PyConvertRule[
         rule for tname in mro for
-        rule in get!(Vector{PyConvertRule}, PYCONVERT_RULES, tname)
+        rule in get!(Vector{PyConvertRule}, PYCONVERT_RULES[], tname)
     ]
 
     # order the rules by priority, then by original order
@@ -261,11 +261,10 @@ function _pyconvert_get_rules(pytype::Py)
     return rules
 end
 
-const PYCONVERT_PREFERRED_TYPE = Dict{Py,Type}()
-const PYCONVERT_PREFERRED_TYPE_LOCK = Threads.SpinLock()
+const PYCONVERT_PREFERRED_TYPE = Lockable(Dict{Py,Type}())
 
 pyconvert_preferred_type(pytype::Py) =
-    Base.@lock PYCONVERT_PREFERRED_TYPE_LOCK get!(PYCONVERT_PREFERRED_TYPE, pytype) do
+    Base.@lock PYCONVERT_PREFERRED_TYPE get!(PYCONVERT_PREFERRED_TYPE[], pytype) do
         if pyissubclass(pytype, pybuiltins.int)
             Union{Int,BigInt}
         else
@@ -308,11 +307,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()
+const PYCONVERT_RULES_CACHE = Lockable(Dict{Type,Dict{C.PyPtr,Vector{Function}}}())
 
 @generated pyconvert_rules_cache(::Type{T}) where {T} =
-    Base.@lock PYCONVERT_RULES_CACHE_LOCK get!(Dict{C.PyPtr,Vector{Function}}, PYCONVERT_RULES_CACHE, T)
+    Base.@lock PYCONVERT_RULES_CACHE get!(Dict{C.PyPtr,Vector{Function}}, PYCONVERT_RULES_CACHE[], T)
 
 function pyconvert_rule_fast(::Type{T}, x::Py) where {T}
     if T isa Union
@@ -353,12 +351,13 @@ function pytryconvert(::Type{T}, x_) where {T}
     # get rules from the cache
     # 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 = Base.@lock PYCONVERT_RULES_CACHE_LOCK get!(trules, tptr) do
-        t = pynew(incref(tptr))
-        ans = pyconvert_get_rules(T, t)::Vector{Function}
-        pydel!(t)
-        ans
+    rules = Base.@lock PYCONVERT_RULES_CACHE let trules = pyconvert_rules_cache(T)
+        get!(trules, tptr) do
+            t = pynew(incref(tptr))
+            ans = pyconvert_get_rules(T, t)::Vector{Function}
+            pydel!(t)
+            ans
+        end
     end
 
     # apply the rules
@@ -420,15 +419,17 @@ pyconvertarg(::Type{T}, x, name) where {T} = @autopy x @pyconvert T x_ begin
 end
 
 function init_pyconvert()
-    push!(PYCONVERT_EXTRATYPES, pyimport("io" => "IOBase"))
-    push!(
-        PYCONVERT_EXTRATYPES,
-        pyimport("numbers" => ("Number", "Complex", "Real", "Rational", "Integral"))...,
-    )
-    push!(
-        PYCONVERT_EXTRATYPES,
-        pyimport("collections.abc" => ("Iterable", "Sequence", "Set", "Mapping"))...,
-    )
+    Base.@lock PYCONVERT_EXTRATYPES begin
+        push!(PYCONVERT_EXTRATYPES[], pyimport("io" => "IOBase"))
+        push!(
+            PYCONVERT_EXTRATYPES[],
+            pyimport("numbers" => ("Number", "Complex", "Real", "Rational", "Integral"))...,
+        )
+        push!(
+            PYCONVERT_EXTRATYPES[],
+            pyimport("collections.abc" => ("Iterable", "Sequence", "Set", "Mapping"))...,
+        )
+    end
 
     priority = PYCONVERT_PRIORITY_CANONICAL
     pyconvert_add_rule("builtins:NoneType", Nothing, pyconvert_rule_none, priority)

src/Core/Core.jl

@@ -11,7 +11,7 @@ const ROOT_DIR = dirname(dirname(@__DIR__))
 using ..PythonCall: PythonCall  # needed for docstring cross-refs
 using ..C: C
 using ..GC: GC
-using ..Utils: Utils
+using ..Utils: Utils, Lockable
 using Base: @propagate_inbounds, @kwdef
 using Dates:
     Date,

src/Core/Py.jl

@@ -56,8 +56,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()
+const PYNULL_CACHE = Lockable(Py[])
 
 """
     pynew([ptr])
@@ -70,11 +69,11 @@ 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() = Base.@lock PYNULL_CACHE_LOCK begin
-    if isempty(PYNULL_CACHE)
+pynew() = Base.@lock PYNULL_CACHE begin
+    if isempty(PYNULL_CACHE[])
         Py(Val(:new), C.PyNULL)
     else
-        pop!(PYNULL_CACHE)
+        pop!(PYNULL_CACHE[])
     end
 end
 
@@ -121,7 +120,7 @@ function pydel!(x::Py)
         C.Py_DecRef(ptr)
         setptr!(x, C.PyNULL)
     end
-    Base.@lock PYNULL_CACHE_LOCK push!(PYNULL_CACHE, x)
+    Base.@lock PYNULL_CACHE push!(PYNULL_CACHE[], x)
     return
 end
 

src/JlWrap/C.jl

@@ -1,7 +1,7 @@
 module Cjl
 
 using ...C: C
-using ...Utils: Utils
+using ...Utils: Utils, Lockable
 using Base: @kwdef
 using UnsafePointers: UnsafePtr
 using Serialization: serialize, deserialize
@@ -16,13 +16,7 @@ const PyJuliaBase_Type = Ref(C.PyNULL)
 
 # we store the actual julia values here
 # the `value` field of `PyJuliaValueObject` indexes into here
-const PYJLVALUES = IdDict{Int,Any}()
-# unused indices in PYJLVALUES
-const PYJLFREEVALUES = Int[]
-# Thread safety for PYJLVALUES and PYJLFREEVALUES
-const PYJLVALUES_LOCK = Threads.SpinLock()
-# Track next available index
-const PYJLVALUES_NEXT_IDX = Ref(1)
+const PYJLVALUES = Lockable((; values=IdDict{Int,Any}(), free=Int[], next=Ref(1)))
 
 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)
@@ -35,24 +29,23 @@ end
 function _pyjl_dealloc(o::C.PyPtr)
     idx = UnsafePtr{PyJuliaValueObject}(o).value[]
     if idx != 0
-        Base.@lock PYJLVALUES_LOCK begin
-            delete!(PYJLVALUES, idx)
-            push!(PYJLFREEVALUES, idx)
+        Base.@lock PYJLVALUES begin
+            delete!(PYJLVALUES[].values, idx)
+            push!(PYJLVALUES[].free, 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()
+const PYJLMETHODS = Lockable([])
 
 function PyJulia_MethodNum(f)
     @nospecialize f
-    Base.@lock PYJLMETHODS_LOCK begin
-        push!(PYJLMETHODS, f)
-        return length(PYJLMETHODS)
+    Base.@lock PYJLMETHODS begin
+        push!(PYJLMETHODS[], f)
+        return length(PYJLMETHODS[])
     end
 end
 
@@ -67,13 +60,12 @@ 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 = Base.@lock PYJLMETHODS_LOCK PYJLMETHODS[num]
+    f = Base.@lock PYJLMETHODS 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()
+const PYJLBUFCACHE = Lockable(Dict{Ptr{Cvoid},Any}())
 
 @kwdef struct PyBufferInfo{N}
     # data
@@ -187,9 +179,7 @@ function _pyjl_get_buffer_impl(
 
     # internal
     cptr = Base.pointer_from_objref(c)
-    Base.@lock PYJLBUFCACHE_LOCK begin
-        PYJLBUFCACHE[cptr] = c
-    end
+    Base.@lock PYJLBUFCACHE PYJLBUFCACHE[][cptr] = c
     b.internal[] = cptr
 
     # obj
@@ -207,7 +197,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 = Base.@lock PYJLMETHODS_LOCK PYJLMETHODS[num]
+        f = Base.@lock PYJLMETHODS PYJLMETHODS[][num]
         x = PyJuliaValue_GetValue(o)
         return _pyjl_get_buffer_impl(o, buf, flags, x, f)::Cint
     catch exc
@@ -221,9 +211,7 @@ 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})
-    Base.@lock PYJLBUFCACHE_LOCK begin
-        delete!(PYJLBUFCACHE, UnsafePtr(buf).internal[!])
-    end
+    Base.@lock PYJLBUFCACHE delete!(PYJLBUFCACHE[], UnsafePtr(buf).internal[!])
     nothing
 end
 
@@ -355,28 +343,26 @@ PyJuliaValue_IsNull(o) = Base.GC.@preserve o UnsafePtr{PyJuliaValueObject}(C.asp
 
 PyJuliaValue_GetValue(o) = Base.GC.@preserve o begin
     idx = UnsafePtr{PyJuliaValueObject}(C.asptr(o)).value[]
-    Base.@lock PYJLVALUES_LOCK begin
-        PYJLVALUES[idx]
-    end
+    Base.@lock PYJLVALUES PYJLVALUES[].values[idx]
 end
 
 PyJuliaValue_SetValue(_o, @nospecialize(v)) = Base.GC.@preserve _o begin
     o = C.asptr(_o)
     idx = UnsafePtr{PyJuliaValueObject}(o).value[]
     if idx == 0
-        Base.@lock PYJLVALUES_LOCK begin
-            if isempty(PYJLFREEVALUES)
-                idx = PYJLVALUES_NEXT_IDX[]
-                PYJLVALUES_NEXT_IDX[] += 1
+        Base.@lock PYJLVALUES begin
+            if isempty(PYJLVALUES[].free)
+                idx = PYJLVALUES[].next[]
+                PYJLVALUES[].next[] += 1
             else
-                idx = pop!(PYJLFREEVALUES)
+                idx = pop!(PYJLVALUES[].free)
             end
-            PYJLVALUES[idx] = v
+            PYJLVALUES[].values[idx] = v
         end
         UnsafePtr{PyJuliaValueObject}(o).value[] = idx
     else
-        Base.@lock PYJLVALUES_LOCK begin
-            PYJLVALUES[idx] = v
+        Base.@lock PYJLVALUES begin
+            PYJLVALUES[].values[idx] = v
         end
     end
     nothing

src/Utils/Utils.jl

@@ -308,4 +308,30 @@ function Base.iterate(x::StaticString{UInt32,N}, i::Int = 1) where {N}
     end
 end
 
+@static if !isdefined(Base, :Lockable)
+    """
+    Compat for `Base.Lockable` (introduced in Julia 1.11)
+    """
+    struct Lockable{T, L<:AbstractLock}
+        value::T
+        lock::L
+    end
+    
+    Lockable(value) = Lockable(value, ReentrantLock())
+
+    function Base.lock(f, l::Lockable)
+        lock(l.lock) do
+            f(l.value)
+        end
+    end
+
+    Base.lock(l::Lockable) = lock(l.lock)
+    Base.trylock(l::Lockable) = trylock(l.lock)
+    Base.unlock(l::Lockable) = unlock(l.lock)
+    Base.islocked(l::Lockable) = islocked(l.lock)
+    Base.getindex(l::Lockable) = (@assert islocked(l); l.value)
+else
+    const Lockable = Base.Lockable
+end
+
 end