Thread safe jl_gc_enable

Author: yuyichao

src/ast.c

@@ -201,7 +201,7 @@ static jl_value_t *resolve_globals(jl_value_t *expr, jl_lambda_info_t *lam);
 
 static jl_value_t *scm_to_julia(value_t e, int expronly)
 {
-    int en = jl_gc_enable(0);
+    int en = jl_gc_enable(0); // Might GC
     jl_value_t *v;
     JL_TRY {
         v = scm_to_julia_(e, expronly);

src/dump.c

@@ -1974,7 +1974,7 @@ JL_DLLEXPORT jl_value_t *jl_ast_rettype(jl_lambda_info_t *li, jl_value_t *ast)
     ios_mem(&src, 0);
     ios_setbuf(&src, (char*)bytes->data, jl_array_len(bytes), 0);
     src.size = jl_array_len(bytes);
-    int en = jl_gc_enable(0);
+    int en = jl_gc_enable(0); // Might GC
     jl_value_t *rt = jl_deserialize_value(&src, NULL);
     jl_gc_enable(en);
     tree_literal_values = NULL;
@@ -1994,7 +1994,7 @@ JL_DLLEXPORT jl_value_t *jl_compress_ast(jl_lambda_info_t *li, jl_value_t *ast)
     ios_mem(&dest, 0);
     jl_array_t *last_tlv = tree_literal_values;
     jl_module_t *last_tem = tree_enclosing_module;
-    int en = jl_gc_enable(0);
+    int en = jl_gc_enable(0); // Might GC
 
     if (li->module->constant_table == NULL) {
         li->module->constant_table = jl_alloc_cell_1d(0);
@@ -2038,7 +2038,7 @@ JL_DLLEXPORT jl_value_t *jl_uncompress_ast(jl_lambda_info_t *li, jl_value_t *dat
     ios_mem(&src, 0);
     ios_setbuf(&src, (char*)bytes->data, jl_array_len(bytes), 0);
     src.size = jl_array_len(bytes);
-    int en = jl_gc_enable(0);
+    int en = jl_gc_enable(0); // Might GC
     (void)jl_deserialize_value(&src, NULL); // skip ret type
     jl_value_t *v = jl_deserialize_value(&src, NULL);
     jl_gc_enable(en);

src/gc.c

@@ -2144,15 +2144,29 @@ static void post_mark(arraylist_t *list, int dryrun)
 }
 
 // collector entry point and control
+static volatile uint64_t jl_gc_disable_counter = 0;
 
-static int is_gc_enabled = 1;
 JL_DLLEXPORT int jl_gc_enable(int on)
 {
-    int prev = is_gc_enabled;
-    is_gc_enabled = (on!=0);
+    jl_tls_states_t *ptls = jl_get_ptls_states();
+    int prev = !ptls->disable_gc;
+    ptls->disable_gc = (on == 0);
+    if (on && !prev) {
+        // disable -> enable
+        JL_ATOMIC_FETCH_AND_ADD(jl_gc_disable_counter, -1);
+    }
+    else if (prev && !on) {
+        // enable -> disable
+        JL_ATOMIC_FETCH_AND_ADD(jl_gc_disable_counter, 1);
+        // check if the GC is running and wait for it to finish
+        jl_gc_safepoint();
+    }
     return prev;
 }
-JL_DLLEXPORT int jl_gc_is_enabled(void) { return is_gc_enabled; }
+JL_DLLEXPORT int jl_gc_is_enabled(void)
+{
+    return !jl_get_ptls_states()->disable_gc;
+}
 
 JL_DLLEXPORT int64_t jl_gc_total_bytes(void) { return total_allocd_bytes + allocd_bytes + collect_interval; }
 JL_DLLEXPORT uint64_t jl_gc_total_hrtime(void) { return total_gc_time; }
@@ -2455,7 +2469,7 @@ static void _jl_gc_collect(int full, char *stack_hi)
 
 JL_DLLEXPORT void jl_gc_collect(int full)
 {
-    if (!is_gc_enabled || jl_in_gc)
+    if (jl_gc_disable_counter || jl_in_gc)
         return;
     char *stack_hi = (char*)gc_get_stack_ptr();
     gc_debug_print();
@@ -2482,7 +2496,8 @@ JL_DLLEXPORT void jl_gc_collect(int full)
     jl_gc_signal_begin();
 
     jl_in_gc = 1;
-    _jl_gc_collect(full, stack_hi);
+    if (!jl_gc_disable_counter)
+        _jl_gc_collect(full, stack_hi);
     jl_in_gc = 0;
 
     // Need to reset the page protection before resetting the flag since

src/julia.h

@@ -1459,6 +1459,7 @@ typedef struct _jl_tls_states_t {
     //              execute at the same time with the GC.
     volatile int8_t gc_state;
     volatile int8_t in_gc;
+    int8_t disable_gc;
     struct _jl_thread_heap_t *heap;
     jl_task_t *volatile current_task;
     jl_task_t *root_task;