JuliaConcurrent · tkf · May 13, 2021 · May 10, 2021 · May 10, 2021 · May 10, 2021
diff --git a/.github/workflows/test-broken.yml b/.github/workflows/test-broken.yml
diff --git a/src/atomicsutils.jl b/src/atomicsutils.jl
@@ -35,6 +35,8 @@ end
     else
         fptr = fieldpointer(obj, field)
         vint = UInt(pointer_from_objref(value))
+        julia_write_barrier(value)
+        julia_write_barrier(obj, value)
         GC.@preserve obj value begin
             UnsafeAtomics.store!(fptr, vint, order)
         end

diff --git a/src/dict.jl b/src/dict.jl
@@ -274,11 +274,13 @@ end
 
 allocate_slot(::AbstractVector{<:AbstractPair}) = nothing
 
-@inline cas_slot!(slotref, new_slot, key) = cas_slot!(slotref, new_slot, key, NoValue())
+@inline cas_slot!(slotref, new_slot, root, key) =
+    cas_slot!(slotref, new_slot, root, key, NoValue())
 
 @inline function cas_slot!(
     slotref::InlinedSlotRef{Slot},
     ::Nothing,
+    root,
     key::KeyUnion{Key},
     value,
 ) where {Key,Value,Slot<:AbstractPair{Key,Value}}
@@ -300,8 +302,10 @@ allocate_slot(::AbstractVector{<:AbstractPair}) = nothing
     if Slot <: InlinedPair
         newkeyint = uint_from(Inlined{KeyUnion{Key}}(key))
     elseif key isa Moved{Key}
-        ref = forceheap(Ref(key))
+        ref = Ref(key)
         newkeyint = UInt(pointer_from_objref(ref))
+        julia_write_barrier(ref)
+        julia_write_barrier(root, ref)
     else
         newkeyint = UInt(_pointer_from_objref(key))
     end
@@ -349,6 +353,7 @@ allocate_slot(::AbstractVector{Slot}) where {Slot<:Ref} = forceheap(Slot())
 @inline function cas_slot!(
     slotref::RefSlotRef{Slot},
     new_slot::Slot,
+    root,
     key,
     value,
 ) where {P,Slot<:Ref{P}}
@@ -359,6 +364,8 @@ allocate_slot(::AbstractVector{Slot}) where {Slot<:Ref} = forceheap(Slot())
     GC.@preserve new_slot begin
         fu = UnsafeAtomics.cas!(Ptr{typeof(nu)}(ptr), ou, nu)
     end
+    julia_write_barrier(new_slot)
+    julia_write_barrier(root, new_slot)
     return fu == ou
 end
 
@@ -462,7 +469,7 @@ function ConcurrentCollections.modify!(
                 if reply isa Keep
                     return reply
                 elseif reply isa Union{Nothing,Delete}
-                    if cas_slot!(slotref, new_slot, Deleted())
+                    if cas_slot!(slotref, new_slot, slots, Deleted())
                         ndeleted = Threads.atomic_add!(dict.ndeleted, 1) + 1
                         approx_len = dict.nadded[] - ndeleted
                         if approx_len < length(slots) ÷ 2
@@ -471,7 +478,7 @@ function ConcurrentCollections.modify!(
                         return reply
                     end
                 else
-                    if cas_slot!(slotref, new_slot, nsk, something(reply))
+                    if cas_slot!(slotref, new_slot, slots, nsk, something(reply))
                         if sk isa Empty
                             Threads.atomic_add!(dict.nadded, 1)
                         end
@@ -661,7 +668,7 @@ function migrate_impl!(
             continue
         elseif sk isa Empty
             # Mark that this slot is not usable anymore
-            if !cas_slot!(slotref, allocate_slot(slots), MovedEmpty())
+            if !cas_slot!(slotref, allocate_slot(slots), slots, MovedEmpty())
                 @goto reload
             end
             stop_on_empty == Val(true) && return Stopped(i, nadded)
@@ -675,7 +682,7 @@ function migrate_impl!(
         if sk isa Moved
             key = sk.key
         else
-            if !cas_slot!(slotref, allocate_slot(slots), Moved(sk), sv)
+            if !cas_slot!(slotref, allocate_slot(slots), slots, Moved(sk), sv)
                 @goto reload
             end
             key = sk

diff --git a/src/utils.jl b/src/utils.jl
@@ -176,6 +176,43 @@ const _FALSE_ = Ref(false)
     return x
 end
 
+@generated function julia_write_barrier(args::Vararg{Any,N}) where {N}
+    pointer_exprs = map(1:N) do i
+        :(_pointer_from_objref(args[$i]))
+    end
+    jlp = "{} addrspace(10)*"
+    llvm_args = string.("%", 0:N-1)
+    word = "i$(Base.Sys.WORD_SIZE)"
+    entry_sig = join(word .* llvm_args, ", ")
+    ptrs = string.("%ptr", 0:N-1)
+    wb_sig = join("$jlp " .* ptrs, ", ")
+    inttoptr = join(
+        (ptrs .* "_tmp = inttoptr $word " .* llvm_args .* " to {}*\n") .*
+        (ptrs .* " = addrspacecast {}* " .* ptrs .* "_tmp to $jlp"),
+        "\n",
+    )
+    IR = (
+        """
+        define void @entry($entry_sig) #0 {
+        top:
+            $inttoptr
+            call void ($jlp, ...) @julia.write_barrier($wb_sig)
+            ret void
+        }
+
+        declare void @julia.write_barrier($jlp, ...) #1
+
+        attributes #0 = { alwaysinline }
+        attributes #1 = { inaccessiblememonly norecurse nounwind }
+        """,
+        "entry",
+    )
+    quote
+        $(Expr(:meta, :inline))
+        Base.llvmcall($IR, Cvoid, NTuple{N,Ptr{Cvoid}}, $(pointer_exprs...))
+    end
+end
+
 @generated allocate_singleton_ref(::Type{T}) where {T} = Ref{Any}(T.instance)
 
 @inline function pointer_from_singleton(::T) where {T}

diff --git a/test/test_dict.jl b/test/test_dict.jl
@@ -4,8 +4,6 @@ using ConcurrentCollections
 using ConcurrentCollections.Implementations: clusters
 using Test
 
-const IS_CI = lowercase(get(ENV, "CI", "false")) == "true"
-
 @testset for npairs in [2, 100]
     @testset for Key in [Int8, Int32, Int64], Value in [Int]
         d = ConcurrentDict{Key,Value}()
@@ -52,11 +50,7 @@ const IS_CI = lowercase(get(ENV, "CI", "false")) == "true"
         end
 
         @testset "clusters" begin
-            if VERSION >= v"1.7-" && IS_CI
-                @info "skipping `clusters` (known bug)"
-            else
-                @test length(clusters(d)::Vector{UnitRange{Int}}) > 0
-            end
+            @test length(clusters(d)::Vector{UnitRange{Int}}) > 0
         end
     end
 
@@ -84,11 +78,7 @@ const IS_CI = lowercase(get(ENV, "CI", "false")) == "true"
             @test "001" ∉ sort!(collect(keys(d)))
         end
         @testset "clusters" begin
-            if VERSION >= v"1.7-" && IS_CI
-                @info "skipping `clusters` (known bug)"
-            else
-                @test length(clusters(d)::Vector{UnitRange{Int}}) > 0
-            end
+            @test length(clusters(d)::Vector{UnitRange{Int}}) > 0
         end
     end
 end