task_local_storage (#15)

JuliaGeometry · Jul 19, 2024 · ebd85bf · ebd85bf · DanielVandH · Jul 19, 2024
1 parent 6a01673
commit ebd85bf
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 42 deletions.
diff --git a/src/AdaptivePredicates.jl b/src/AdaptivePredicates.jl
@@ -10,7 +10,7 @@ export orient2, orient3, incircle, insphere
 export orient2p, orient3p, incirclep, inspherep
 
 @static if VERSION ≥ v"1.11.0-DEV.469"
-    eval(Meta.parse("public orient2, orient3, incircle, insphere, orient2p, orient3p, incirclep, inspherep, free!"))
+    eval(Meta.parse("public orient2, orient3, incircle, insphere, orient2p, orient3p, incirclep, inspherep"))
 end
 
 end # module
diff --git a/src/caches.jl b/src/caches.jl
@@ -9,47 +9,26 @@ struct CacheKey{T}
     id::UInt8
 end
 
-const CacheLock = Base.Threads.SpinLock()
+struct APMarker{T} end # Avoid collisions with other packages using the task_local_storage
+
+const APCache{T} = Dict{CacheKey{T},Vec{T}}
 
-const TASK_LOCAL_F64CACHE = Dict{Task,Dict{CacheKey{Float64},Vec{Float64}}}()
-const TASK_LOCAL_F32CACHE = Dict{Task,Dict{CacheKey{Float32},Vec{Float32}}}()
-@inline TASK_LOCAL_CACHE(::Type{Float64}) = TASK_LOCAL_F64CACHE
-@inline TASK_LOCAL_CACHE(::Type{Float32}) = TASK_LOCAL_F32CACHE
 @inline function task_local_cache(::Type{T}) where {T}
-    tls = TASK_LOCAL_CACHE(T)
-    t = current_task()
-    if haskey(tls, t)
-        return tls[t]
-    else
-        d = Dict{CacheKey{T},Vec{T}}()
-        lock(CacheLock) do
-            tls[t] = d
-        end
-        return tls[t]
-    end
+    tls = get!(task_local_storage(), APMarker{T}()) do
+        APCache{T}()
+    end::APCache{T}
+    return tls::APCache{T}
 end
 
-"""
-    free!()
-
-Empties the caches used for computing the predicates.
-"""
-free!() = (empty!(TASK_LOCAL_F64CACHE); empty!(TASK_LOCAL_F32CACHE))
-
-@inline cache_eltype(::Dict{CacheKey{Float64},Vec{Float64}}) = Float64
-@inline cache_eltype(::Dict{CacheKey{Float32},Vec{Float32}}) = Float32
 
-@inline function get_cache!(tls, size, id)
-    T = cache_eltype(tls)
+@inline function get_cache!(tls::APCache{T}, size, id) where {T}
     cache::Vec{T} = get!(tls, CacheKey{T}(size, id)) do
         Vec{T}(zeros(T, Int(size))) # Memory{T}(undef, Int(size)) has weird concurrency issues sometimes?
     end
     return cache::Vec{T}
 end
 
-abstract type AbstractCache{T} end
-
-struct Orient2Cache{T} <: AbstractCache{T}
+struct Orient2Cache{T}
     h4::NTuple{4,T}
     h8::NTuple{8,T}
     h12::NTuple{12,T}
@@ -63,7 +42,7 @@ end
     return Orient2Cache{T}(h4, h8, h12, h16)
 end
 
-struct Orient3Cache{T} <: AbstractCache{T}
+struct Orient3Cache{T}
     h4::NTuple{4,T}
     h8::NTuple{8,T}
     h12::NTuple{12,T}
@@ -83,6 +62,9 @@ struct Orient3Cache{T} <: AbstractCache{T}
 end
 @inline function Orient3Cache{T}() where {T}
     tls = task_local_cache(T)
+    return Orient3Cache{T}(tls)
+end
+@inline function Orient3Cache{T}(tls::APCache{T}) where {T}
     h4 = ntuple(_ -> zero(T), Val(4))
     h8 = ntuple(_ -> zero(T), Val(8))
     h12 = ntuple(_ -> zero(T), Val(12))
@@ -105,7 +87,7 @@ end
     )
 end
 
-struct IncircleCache{T} <: AbstractCache{T}
+struct IncircleCache{T}
     h4::NTuple{4,T}
     h8::NTuple{8,T}
     h12::NTuple{12,T}
@@ -137,6 +119,9 @@ struct IncircleCache{T} <: AbstractCache{T}
 end
 @inline function IncircleCache{T}() where {T}
     tls = task_local_cache(T)
+    return IncircleCache{T}(tls)
+end
+@inline function IncircleCache{T}(tls::APCache{T}) where {T}
     h4 = ntuple(_ -> zero(T), Val(4))
     h8 = ntuple(_ -> zero(T), Val(8))
     h12 = ntuple(_ -> zero(T), Val(12))
@@ -174,7 +159,7 @@ end
     )
 end
 
-struct InsphereCache{T} <: AbstractCache{T}
+struct InsphereCache{T}
     h4::NTuple{4,T}
     h8::NTuple{8,T}
     h12::NTuple{12,T}
@@ -238,6 +223,9 @@ struct InsphereCache{T} <: AbstractCache{T}
 end
 @inline function InsphereCache{T}() where {T}
     tls = task_local_cache(T)
+    return InsphereCache{T}(tls)
+end
+@inline function InsphereCache{T}(tls::APCache{T}) where {T}
     h4 = ntuple(_ -> zero(T), Val(4))
     h8 = ntuple(_ -> zero(T), Val(8))
     h12 = ntuple(_ -> zero(T), Val(12))

diff --git a/test/Project.toml b/test/Project.toml
@@ -1,5 +1,9 @@
 [deps]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ExactPredicates = "429591f6-91af-11e9-00e2-59fbe8cec110"
 Supposition = "5a0628fe-1738-4658-9b6d-0b7605a9755b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[compat]
+Aqua = "0.8.7"
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -3,6 +3,12 @@ using Test
 using Supposition
 using BenchmarkTools
 import ExactPredicates: ExactPredicates
+using Aqua
+
+@testset "Aqua" begin
+    Aqua.test_all(AdaptivePredicates)
+end
+
 const AP = AdaptivePredicates
 cd("original") do
     include("compile.jl")
@@ -194,12 +200,4 @@ setup_insphere(T) = ntuple(_ -> (_rand(T), _rand(T), _rand(T)), 5)
     @test iszero(@ballocated incircle(args...) setup = (args = setup_incircle(Float32)))
     @test iszero(@ballocated insphere(args...) setup = (args = setup_insphere(Float64)))
     @test iszero(@ballocated insphere(args...) setup = (args = setup_insphere(Float32)))
-end
-
-@testset "free!" begin
-    F64C = AP.TASK_LOCAL_F64CACHE
-    F32C = AP.TASK_LOCAL_F32CACHE
-    @test !isempty(F64C) && !isempty(F32C)
-    AP.free!()
-    @test isempty(F64C) && isempty(F32C)
 end