bpart: Properly track methods with invalidated source after require_world (#58830)

There are three categories of methods we need to worry about during
staticdata validation:
1. New methods added to existing generic functions
2. New methods added to new generic functions
3. Existing methods that now have new CodeInstances

In each of these cases, we need to check whether any of the implicit
binding edges from the method's source was invalidated. Currently, we
handle this for 1 and 2 by explicitly scanning the method on load.
However, we were not tracking it for case 3. Fix that by using an extra
bit in did_scan_method that gets set when we see an existing method
getting invalidated, so we know that we need to drop the corresponding
CodeInstances during load.

Fixes #58346

Author: Keno

base/invalidation.jl

@@ -69,6 +69,9 @@ function invalidate_method_for_globalref!(gr::GlobalRef, method::Method, invalid
         src = _uncompressed_ir(method)
         invalidate_all = should_invalidate_code_for_globalref(gr, src)
     end
+    if invalidate_all && !Base.generating_output()
+        @atomic method.did_scan_source |= 0x4
+    end
     invalidated_any = false
     for mi in specializations(method)
         isdefined(mi, :cache) || continue
@@ -182,7 +185,7 @@ function binding_was_invalidated(b::Core.Binding)
     b.partitions.min_world > unsafe_load(cglobal(:jl_require_world, UInt))
 end
 
-function scan_new_method!(methods_with_invalidated_source::IdSet{Method}, method::Method, image_backedges_only::Bool)
+function scan_new_method!(method::Method, image_backedges_only::Bool)
     isdefined(method, :source) || return
     if image_backedges_only && !has_image_globalref(method)
         return
@@ -195,21 +198,20 @@ function scan_new_method!(methods_with_invalidated_source::IdSet{Method}, method
             # TODO: We could turn this into an addition if condition. For now, use it as a reasonably cheap
             # additional consistency check
             @assert !image_backedges_only
-            push!(methods_with_invalidated_source, method)
+            @atomic method.did_scan_source |= 0x4
         end
         maybe_add_binding_backedge!(b, method)
     end
+    @atomic method.did_scan_source |= 0x1
 end
 
-function scan_new_methods(extext_methods::Vector{Any}, internal_methods::Vector{Any}, image_backedges_only::Bool)
-    methods_with_invalidated_source = IdSet{Method}()
+function scan_new_methods!(extext_methods::Vector{Any}, internal_methods::Vector{Any}, image_backedges_only::Bool)
     for method in internal_methods
         if isa(method, Method)
-           scan_new_method!(methods_with_invalidated_source, method, image_backedges_only)
+           scan_new_method!(method, image_backedges_only)
         end
     end
     for tme::Core.TypeMapEntry in extext_methods
-        scan_new_method!(methods_with_invalidated_source, tme.func::Method, image_backedges_only)
+        scan_new_method!(tme.func::Method, image_backedges_only)
     end
-    return methods_with_invalidated_source
 end

base/runtime_internals.jl

@@ -1414,6 +1414,8 @@ Returns the world the [current_task()](@ref) is executing within.
 """
 tls_world_age() = ccall(:jl_get_tls_world_age, UInt, ())
 
+get_require_world() = unsafe_load(cglobal(:jl_require_world, UInt))
+
 """
     propertynames(x, private=false)
 

base/staticdata.jl

@@ -17,20 +17,20 @@ function insert_backedges(edges::Vector{Any}, ext_ci_list::Union{Nothing,Vector{
     # determine which CodeInstance objects are still valid in our image
     # to enable any applicable new codes
     backedges_only = unsafe_load(cglobal(:jl_first_image_replacement_world, UInt)) == typemax(UInt)
-    methods_with_invalidated_source = Base.scan_new_methods(extext_methods, internal_methods, backedges_only)
+    Base.scan_new_methods!(extext_methods, internal_methods, backedges_only)
     stack = CodeInstance[]
     visiting = IdDict{CodeInstance,Int}()
-    _insert_backedges(edges, stack, visiting, methods_with_invalidated_source)
+    _insert_backedges(edges, stack, visiting)
     if ext_ci_list !== nothing
-        _insert_backedges(ext_ci_list, stack, visiting, methods_with_invalidated_source, #=external=#true)
+        _insert_backedges(ext_ci_list, stack, visiting, #=external=#true)
     end
 end
 
-function _insert_backedges(edges::Vector{Any}, stack::Vector{CodeInstance}, visiting::IdDict{CodeInstance,Int}, mwis::IdSet{Method}, external::Bool=false)
+function _insert_backedges(edges::Vector{Any}, stack::Vector{CodeInstance}, visiting::IdDict{CodeInstance,Int}, external::Bool=false)
     for i = 1:length(edges)
         codeinst = edges[i]::CodeInstance
         validation_world = get_world_counter()
-        verify_method_graph(codeinst, stack, visiting, mwis, validation_world)
+        verify_method_graph(codeinst, stack, visiting, validation_world)
         # After validation, under the world_counter_lock, set max_world to typemax(UInt) for all dependencies
         # (recursively). From that point onward the ordinary backedge mechanism is responsible for maintaining
         # validity.
@@ -54,16 +54,14 @@ function _insert_backedges(edges::Vector{Any}, stack::Vector{CodeInstance}, visi
     end
 end
 
-function verify_method_graph(codeinst::CodeInstance, stack::Vector{CodeInstance}, visiting::IdDict{CodeInstance,Int}, mwis::IdSet{Method}, validation_world::UInt)
+function verify_method_graph(codeinst::CodeInstance, stack::Vector{CodeInstance}, visiting::IdDict{CodeInstance,Int}, validation_world::UInt)
     @assert isempty(stack); @assert isempty(visiting);
-    child_cycle, minworld, maxworld = verify_method(codeinst, stack, visiting, mwis, validation_world)
+    child_cycle, minworld, maxworld = verify_method(codeinst, stack, visiting, validation_world)
     @assert child_cycle == 0
     @assert isempty(stack); @assert isempty(visiting);
     nothing
 end
 
-get_require_world() = unsafe_load(cglobal(:jl_require_world, UInt))
-
 function gen_staged_sig(def::Method, mi::MethodInstance)
     isdefined(def, :generator) || return nothing
     isdispatchtuple(mi.specTypes) || return nothing
@@ -113,7 +111,7 @@ end
 # - Visit the entire call graph, starting from edges[idx] to determine if that method is valid
 # - Implements Tarjan's SCC (strongly connected components) algorithm, simplified to remove the count variable
 #   and slightly modified with an early termination option once the computation reaches its minimum
-function verify_method(codeinst::CodeInstance, stack::Vector{CodeInstance}, visiting::IdDict{CodeInstance,Int}, mwis::IdSet{Method}, validation_world::UInt)
+function verify_method(codeinst::CodeInstance, stack::Vector{CodeInstance}, visiting::IdDict{CodeInstance,Int}, validation_world::UInt)
     world = codeinst.min_world
     let max_valid2 = codeinst.max_world
         if max_valid2 ≠ WORLD_AGE_REVALIDATION_SENTINEL
@@ -127,13 +125,13 @@ function verify_method(codeinst::CodeInstance, stack::Vector{CodeInstance}, visi
     end
 
     # Implicitly referenced bindings in the current module do not get explicit edges.
-    # If they were invalidated, they'll be in `mwis`. If they weren't, they imply a minworld
+    # If they were invalidated, they'll have the flag set in did_scan_source. If they weren't, they imply a minworld
     # of `get_require_world`. In principle, this is only required for methods that do reference
     # an implicit globalref. However, we already don't perform this validation for methods that
     # don't have any (implicit or explicit) edges at all. The remaining corner case (some explicit,
     # but no implicit edges) is rare and there would be little benefit to lower the minworld for it
     # in any case, so we just always use `get_require_world` here.
-    local minworld::UInt, maxworld::UInt = get_require_world(), validation_world
+    local minworld::UInt, maxworld::UInt = Base.get_require_world(), validation_world
     if haskey(visiting, codeinst)
         return visiting[codeinst], minworld, maxworld
     end
@@ -143,7 +141,11 @@ function verify_method(codeinst::CodeInstance, stack::Vector{CodeInstance}, visi
     # TODO JL_TIMING(VERIFY_IMAGE, VERIFY_Methods)
     callees = codeinst.edges
     # Check for invalidation of the implicit edges from GlobalRef in the Method source
-    if def in mwis
+    if (def.did_scan_source & 0x1) == 0x0
+        backedges_only = unsafe_load(cglobal(:jl_first_image_replacement_world, UInt)) == typemax(UInt)
+        Base.scan_new_method!(def, backedges_only)
+    end
+    if (def.did_scan_source & 0x4) != 0x0
         maxworld = 0
         invalidations = _jl_debug_method_invalidation[]
         if invalidations !== nothing
@@ -153,7 +155,7 @@ function verify_method(codeinst::CodeInstance, stack::Vector{CodeInstance}, visi
     # verify current edges
     if isempty(callees)
         # quick return: no edges to verify (though we probably shouldn't have gotten here from WORLD_AGE_REVALIDATION_SENTINEL)
-    elseif maxworld == get_require_world()
+    elseif maxworld == Base.get_require_world()
         # if no new worlds were allocated since serializing the base module, then no new validation is worth doing right now either
     else
         j = 1
@@ -231,7 +233,7 @@ function verify_method(codeinst::CodeInstance, stack::Vector{CodeInstance}, visi
             end
             callee = edge
             local min_valid2::UInt, max_valid2::UInt
-            child_cycle, min_valid2, max_valid2 = verify_method(callee, stack, visiting, mwis, validation_world)
+            child_cycle, min_valid2, max_valid2 = verify_method(callee, stack, visiting, validation_world)
             if minworld < min_valid2
                 minworld = min_valid2
             end

src/jltypes.c

@@ -3607,7 +3607,7 @@ void jl_init_types(void) JL_GC_DISABLED
                         0, 1, 10);
     //const static uint32_t method_constfields[1] = { 0b0 }; // (1<<0)|(1<<1)|(1<<2)|(1<<3)|(1<<6)|(1<<9)|(1<<10)|(1<<17)|(1<<21)|(1<<22)|(1<<23)|(1<<24)|(1<<25)|(1<<26)|(1<<27)|(1<<28)|(1<<29)|(1<<30);
     //jl_method_type->name->constfields = method_constfields;
-    const static uint32_t method_atomicfields[1] = { 0x00000030 }; // (1<<4)|(1<<5)
+    const static uint32_t method_atomicfields[1] = { 0x10000030 }; // (1<<4)|(1<<5)||(1<<28)
     jl_method_type->name->atomicfields = method_atomicfields;
 
     jl_method_instance_type =

src/julia.h

@@ -368,6 +368,7 @@ typedef struct _jl_method_t {
     uint8_t nospecializeinfer;
     // bit flags, 0x01 = scanned
     // 0x02 = added to module scanned list (either from scanning or inference edge)
+    // 0x04 = Source was invalidated since jl_require_world
     _Atomic(uint8_t) did_scan_source;
 
     // uint8 settings

test/core.jl

@@ -36,7 +36,7 @@ end
 for (T, c) in (
         (Core.CodeInfo, []),
         (Core.CodeInstance, [:next, :min_world, :max_world, :inferred, :edges, :debuginfo, :ipo_purity_bits, :invoke, :specptr, :specsigflags, :precompile, :time_compile]),
-        (Core.Method, [:primary_world, :dispatch_status]),
+        (Core.Method, [:primary_world, :did_scan_source, :dispatch_status]),
         (Core.MethodInstance, [:cache, :flags]),
         (Core.MethodTable, [:defs]),
         (Core.MethodCache, [:leafcache, :cache, :var""]),

test/precompile.jl

@@ -1083,9 +1083,8 @@ precompile_test_harness("code caching") do dir
         @test hasvalid(mi, world)       # was compiled with the new method
         m = only(methods(MA.fib))
         mi = m.specializations::Core.MethodInstance
-        @test isdefined(mi, :cache)     # it was precompiled by StaleB
-        @test_broken !hasvalid(mi, world)      # invalidated by redefining `gib` before loading StaleB
-        @test_broken MA.fib() === 2.0
+        @test !hasvalid(mi, world)      # invalidated by redefining `gib` before loading StaleB
+        @test MA.fib() === 2.0
 
         # Reporting test (ensure SnoopCompile works)
         @test all(i -> isassigned(invalidations, i), eachindex(invalidations))