make Dict safe to use from finalizers

we may want to start indicating which APIs are safe to call from finalizers
which generally will require disabling finalizers in most of the
functions of that API

to facilitate that, expose control from userspace of the finalizer-disable
mechanism already used by codegen (modeled similarly to gc-disable)

to make both easier to use, automatically restore them when an error is thrown

fix #14445

Author: vtjnash

base/base.jl

@@ -80,7 +80,8 @@ end
 finalize(o::ANY) = ccall(:jl_finalize, Void, (Any,), o)
 
 gc(full::Bool=true) = ccall(:jl_gc_collect, Void, (Cint,), full)
-gc_enable(on::Bool) = ccall(:jl_gc_enable, Cint, (Cint,), on)!=0
+gc_enable(on::Bool) = ccall(:jl_gc_enable, Cint, (Cint,), on) != 0
+gc_finalizers_enable(on::Bool) = ccall(:jl_finalizers_enable, Cint, (Cint,), on) != 0
 
 bytestring(str::String) = str
 

base/coreimg.jl

@@ -64,6 +64,7 @@ end
 include("reduce.jl")
 
 ## core structures
+gc_finalizers_enable(on::Bool) = true # Core.Inference doesn't use finalizers
 include("intset.jl")
 include("dict.jl")
 include("iterator.jl")

base/dict.jl

@@ -537,6 +537,8 @@ function rehash!{K,V}(h::Dict{K,V}, newsz = length(h.keys))
     count = 0
     maxprobe = h.maxprobe
 
+    # make sure the following is atomic wrt finalizers
+    prev_finalizers = gc_finalizers_enable(false)
     for i = 1:sz
         if olds[i] == 0x1
             k = oldk[i]
@@ -551,11 +553,6 @@ function rehash!{K,V}(h::Dict{K,V}, newsz = length(h.keys))
             keys[index] = k
             vals[index] = v
             count += 1
-
-            if h.age != age0
-                # if `h` is changed by a finalizer, retry
-                return rehash!(h, newsz)
-            end
         end
     end
 
@@ -565,7 +562,8 @@ function rehash!{K,V}(h::Dict{K,V}, newsz = length(h.keys))
     h.count = count
     h.ndel = 0
     h.maxprobe = maxprobe
-    assert(h.age == age0)
+    @assert h.age == age0
+    gc_finalizers_enable(prev_finalizers)
 
     return h
 end
@@ -584,6 +582,7 @@ function sizehint!(d::Dict, newsz)
 end
 
 function empty!{K,V}(h::Dict{K,V})
+    prev_finalizers = gc_finalizers_enable(false)
     fill!(h.slots, 0x0)
     sz = length(h.slots)
     empty!(h.keys)
@@ -594,6 +593,7 @@ function empty!{K,V}(h::Dict{K,V})
     h.count = 0
     h.age += 1
     h.idxfloor = 1
+    gc_finalizers_enable(prev_finalizers)
     return h
 end
 
@@ -698,6 +698,7 @@ function setindex!{K,V}(h::Dict{K,V}, v0, key0)
     end
     v = convert(V, v0)
 
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex2(h, key)
 
     if index > 0
@@ -707,37 +708,31 @@ function setindex!{K,V}(h::Dict{K,V}, v0, key0)
     else
         _setindex!(h, v, key, -index)
     end
+    gc_finalizers_enable(prev_finalizers)
 
     return h
 end
 
-function get!{K,V}(h::Dict{K,V}, key0, default)
-    key = convert(K,key0)
-    if !isequal(key,key0)
-        throw(ArgumentError("$key0 is not a valid key for type $K"))
-    end
-
-    index = ht_keyindex2(h, key)
-
-    index > 0 && return h.vals[index]
-
-    v = convert(V,  default)
-    _setindex!(h, v, key, -index)
-    return v
-end
-
+get!{K,V}(h::Dict{K,V}, key0, default) = get!(()->default, h, key0)
 function get!{K,V}(default::Callable, h::Dict{K,V}, key0)
     key = convert(K,key0)
     if !isequal(key,key0)
         throw(ArgumentError("$key0 is not a valid key for type $K"))
     end
 
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex2(h, key)
 
-    index > 0 && return h.vals[index]
+    if index > 0
+        val = h.vals[index]
+        gc_finalizers_enable(prev_finalizers)
+        return val
+    end
 
     age0 = h.age
-    v = convert(V,  default())
+    gc_finalizers_enable(prev_finalizers)
+    v = convert(V, default())
+    prev_finalizers = gc_finalizers_enable(false)
     if h.age != age0
         index = ht_keyindex2(h, key)
     end
@@ -748,6 +743,7 @@ function get!{K,V}(default::Callable, h::Dict{K,V}, key0)
     else
         _setindex!(h, v, key, -index)
     end
+    gc_finalizers_enable(prev_finalizers)
     return v
 end
 
@@ -761,26 +757,47 @@ end
 
 
 function getindex{K,V}(h::Dict{K,V}, key)
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex(h, key)
-    return (index<0) ? throw(KeyError(key)) : h.vals[index]::V
+    index < 0 && throw(KeyError(key))
+    val = h.vals[index]::V
+    gc_finalizers_enable(prev_finalizers)
+    return val
 end
 
 function get{K,V}(h::Dict{K,V}, key, default)
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex(h, key)
-    return (index<0) ? default : h.vals[index]::V
+    val =  (index < 0) ? default : h.vals[index]::V
+    gc_finalizers_enable(prev_finalizers)
+    return val
 end
 
 function get{K,V}(default::Callable, h::Dict{K,V}, key)
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex(h, key)
-    return (index<0) ? default() : h.vals[index]::V
+    if index < 0
+        gc_finalizers_enable(prev_finalizers)
+        return default()
+    end
+    val = h.vals[index]::V
+    gc_finalizers_enable(prev_finalizers)
+    return val
 end
 
 haskey(h::Dict, key) = (ht_keyindex(h, key) >= 0)
 in{T<:Dict}(key, v::KeyIterator{T}) = (ht_keyindex(v.dict, key) >= 0)
 
 function getkey{K,V}(h::Dict{K,V}, key, default)
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex(h, key)
-    return (index<0) ? default : h.keys[index]::K
+    if index < 0
+        gc_finalizers_enable(prev_finalizers)
+        return default
+    end
+    val = h.keys[index]::V
+    gc_finalizers_enable(prev_finalizers)
+    return val
 end
 
 function _pop!(h::Dict, index)
@@ -790,13 +807,24 @@ function _pop!(h::Dict, index)
 end
 
 function pop!(h::Dict, key)
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex(h, key)
-    index > 0 ? _pop!(h, index) : throw(KeyError(key))
+    index < 0 && throw(KeyError(key))
+    val = _pop!(h, index)
+    gc_finalizers_enable(prev_finalizers)
+    return val
 end
 
 function pop!(h::Dict, key, default)
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex(h, key)
-    index > 0 ? _pop!(h, index) : default
+    if index < 0
+        gc_finalizers_enable(prev_finalizers)
+        return default
+    end
+    val = _pop!(h, index)
+    gc_finalizers_enable(prev_finalizers)
+    return val
 end
 
 function _delete!(h::Dict, index)
@@ -806,13 +834,17 @@ function _delete!(h::Dict, index)
     h.ndel += 1
     h.count -= 1
     h.age += 1
-    h
+    return h
 end
 
 function delete!(h::Dict, key)
+    prev_finalizers = gc_finalizers_enable(false)
     index = ht_keyindex(h, key)
-    if index > 0; _delete!(h, index); end
-    h
+    if index > 0
+        _delete!(h, index)
+    end
+    gc_finalizers_enable(prev_finalizers)
+    return h
 end
 
 function skip_deleted(h::Dict, i)
@@ -840,32 +872,28 @@ next{T<:Dict}(v::ValueIterator{T}, i) = (v.dict.vals[i], skip_deleted(v.dict,i+1
 # weak key dictionaries
 
 type WeakKeyDict{K,V} <: Associative{K,V}
-    ht::Dict{Any,V}
-    deleter::Function
-
-    WeakKeyDict() = new(Dict{Any,V}(), identity)
+    ht::Dict{WeakRef,V}
+    WeakKeyDict() = new(Dict{WeakRef,V}())
 end
 WeakKeyDict() = WeakKeyDict{Any,Any}()
 
-function weak_key_delete!(t::Dict, k)
+function weak_key_delete!(t::Dict, k::WeakRef)
     # when a weak key is finalized, remove from dictionary if it is still there
     wk = getkey(t, k, secret_table_token)
-    if !is(wk,secret_table_token) && is(wk.value, k)
+    if is(wk, k)
         delete!(t, k)
     end
 end
 
 function setindex!{K}(wkh::WeakKeyDict{K}, v, key)
     t = wkh.ht
     k = convert(K, key)
-    if is(wkh.deleter, identity)
-        wkh.deleter = x->weak_key_delete!(t, x)
-    end
-    t[WeakRef(k)] = v
+    wr = WeakRef(k)
+    t[wr] = v
     # TODO: it might be better to avoid the finalizer, allow
     # wiped WeakRefs to remain in the table, and delete them as
     # they are discovered by getindex and setindex!.
-    finalizer(k, wkh.deleter)
+    finalizer(k, x->weak_key_delete!(t, wr))
     return t
 end
 

src/builtins.c

@@ -195,12 +195,16 @@ JL_CALLABLE(jl_f_throw)
 JL_DLLEXPORT void jl_enter_handler(jl_handler_t *eh)
 {
     JL_SIGATOMIC_BEGIN();
-    eh->prev = jl_current_task->eh;
-    eh->gcstack = jl_pgcstack;
+    jl_tls_states_t *ptls = jl_get_ptls_states();
+    eh->prev = ptls->current_task->eh;
+    eh->gcstack = ptls->pgcstack;
 #ifdef JULIA_ENABLE_THREADING
-    eh->gc_state = jl_gc_state();
-    eh->locks_len = jl_current_task->locks.len;
+    eh->gc_state = ptls->gc_state;
+    eh->locks_len = ptls->current_task->locks.len;
 #endif
+    eh->gc_enabled = !ptls->disable_gc;
+    eh->finalizers_enabled = !ptls->disable_finalizers;
+    eh->prev = jl_current_task->eh;
     jl_current_task->eh = eh;
     // TODO: this should really go after setjmp(). see comment in
     // ctx_switch in task.c.

src/codegen.cpp

@@ -813,15 +813,16 @@ static void to_function(jl_lambda_info_t *li)
     if (!nested_compile && dump_compiles_stream != NULL)
         last_time = jl_hrtime();
     nested_compile = true;
-    jl_gc_inhibit_finalizers(nested_compile);
     std::unique_ptr<Module> m;
     Function *f = NULL, *specf = NULL;
     // actually do the work of emitting the function
     JL_TRY {
+        int last_inhibited = jl_finalizers_enable(0);
         m = emit_function(li, &li->functionObjectsDecls);
         f = (Function*)li->functionObjectsDecls.functionObject;
         specf = (Function*)li->functionObjectsDecls.specFunctionObject;
         //n_emit++;
+        jl_finalizers_enable(last_inhibited);
     }
     JL_CATCH {
         // something failed! this is very bad, since other WIP may be pointing to this function
@@ -871,7 +872,6 @@ static void to_function(jl_lambda_info_t *li)
     }
     li->inCompile = 0;
     nested_compile = last_n_c;
-    jl_gc_inhibit_finalizers(nested_compile);
     JL_UNLOCK(&codegen_lock);
     JL_SIGATOMIC_END();
     if (dump_compiles_stream != NULL) {
@@ -3866,15 +3866,15 @@ static Function *jl_cfunction_object(jl_function_t *ff, jl_value_t *declrt, jl_t
     DebugLoc olddl = builder.getCurrentDebugLocation();
     bool last_n_c = nested_compile;
     nested_compile = true;
-    jl_gc_inhibit_finalizers(nested_compile);
+    int last_inhibited = jl_finalizers_enable(0);
     Function *f = gen_cfun_wrapper(ff, crt, (jl_tupletype_t*)argt, sf, declrt, (jl_tupletype_t*)sigt);
     // Restore the previous compile context
     if (old != NULL) {
         builder.SetInsertPoint(old);
         builder.SetCurrentDebugLocation(olddl);
     }
     nested_compile = last_n_c;
-    jl_gc_inhibit_finalizers(nested_compile);
+    jl_finalizers_enable(last_inhibited);
     JL_SIGATOMIC_END();
     JL_GC_POP();
     return f;

src/gc.c

@@ -443,7 +443,8 @@ static void jl_gc_signal_end(void)
 
 #endif
 
-static int jl_gc_finalizers_inhibited; // don't run finalizers during codegen #11956
+static volatile uint64_t jl_finalizers_disabled_counter; // don't run finalizers during codegen (for example, #11956)
+                                                         // counter keeps track of per-thread state
 
 // malloc wrappers, aligned allocation
 
@@ -539,16 +540,28 @@ static void run_finalizers(void)
     arraylist_free(&copied_list);
 }
 
-void jl_gc_inhibit_finalizers(int state)
+JL_DLLEXPORT int jl_finalizers_enable(int on)
 {
-    // NOTE: currently only called with the codegen lock held, but might need
-    // more synchronization in the future
-    if (jl_gc_finalizers_inhibited && !state && !jl_in_finalizer) {
-        jl_in_finalizer = 1;
-        run_finalizers();
-        jl_in_finalizer = 0;
+    jl_tls_states_t *ptls = jl_get_ptls_states();
+    int prev = !ptls->disable_finalizers;
+    ptls->disable_finalizers = (on == 0);
+    if (on && !prev) {
+        // disable -> enable
+        // TODO: the atomics are in the wrong place / missing some thread locks
+        jl_atomic_fetch_add(&jl_finalizers_disabled_counter, -1);
+        if (jl_finalizers_disabled_counter == 0) {
+            int8_t was_in_finalizer = jl_in_finalizer;
+            jl_in_finalizer = 1;
+            run_finalizers();
+            jl_in_finalizer = was_in_finalizer;
+        }
     }
-    jl_gc_finalizers_inhibited = state;
+    else if (prev && !on) {
+        // enable -> disable
+        jl_atomic_fetch_add(&jl_finalizers_disabled_counter, 1);
+        // TODO: check if finalizers are running and wait for them to finish
+    }
+    return prev;
 }
 
 static void schedule_all_finalizers(arraylist_t *flist)
@@ -2056,7 +2069,7 @@ static void post_mark(arraylist_t *list, int dryrun)
 }
 
 // collector entry point and control
-static volatile uint64_t jl_gc_disable_counter = 0;
+static volatile uint64_t jl_gc_disable_counter = 0; // counter keeps track of per-thread state
 
 JL_DLLEXPORT int jl_gc_enable(int on)
 {
@@ -2357,7 +2370,7 @@ JL_DLLEXPORT void jl_gc_collect(int full)
     JL_SIGATOMIC_END();
     jl_gc_state_set(old_state, JL_GC_STATE_WAITING);
 
-    if (!jl_gc_finalizers_inhibited) {
+    if (!jl_finalizers_disabled_counter) {
         int8_t was_in_finalizer = jl_in_finalizer;
         jl_in_finalizer = 1;
         run_finalizers();