Move logstate into task-local storage

Author: jpsamaroo

base/boot.jl

@@ -163,7 +163,6 @@
 #    result::Any
 #    exception::Any
 #    backtrace::Any
-#    logstate::Any
 #    code::Any
 #end
 

base/logging.jl

@@ -492,7 +492,7 @@ end
 LogState(logger) = LogState(LogLevel(_invoked_min_enabled_level(logger)), logger)
 
 function current_logstate()
-    logstate = current_task().logstate
+    logstate = get(task_local_storage(), :logstate, nothing)
     return (logstate !== nothing ? logstate : _global_logstate)::LogState
 end
 
@@ -505,15 +505,28 @@ end
     return nothing
 end
 
+function logstate_task_hook(code::UInt8, to)
+    if code == 0
+        to::Task
+        @assert to.storage === nothing
+        to.storage = IdDict{Any,Any}()
+        to.storage[:logstate] = get(task_local_storage(), :logstate, nothing)
+        attach_task_hook!(to, logstate_task_hook_ptr)
+    end
+    return
+end
+const logstate_task_hook_ptr = @cfunction(logstate_task_hook, Nothing, (UInt8, Any))
+
 function with_logstate(f::Function, logstate)
     @nospecialize
-    t = current_task()
-    old = t.logstate
+    tls = task_local_storage()
+    old = get(tls, :logstate, nothing)
+    tls[:logstate] = logstate
+    attach_task_hook!(to, logstate_task_hook_ptr)
     try
-        t.logstate = logstate
         f()
     finally
-        t.logstate = old
+        tls[:logstate] = old
     end
 end
 

src/jltypes.c

@@ -2769,13 +2769,12 @@ void jl_init_types(void) JL_GC_DISABLED
                         NULL,
                         jl_any_type,
                         jl_emptysvec,
-                        jl_perm_symsvec(17,
+                        jl_perm_symsvec(16,
                                         "next",
                                         "queue",
                                         "storage",
                                         "donenotify",
                                         "result",
-                                        "logstate",
                                         "code",
                                         "rngState0",
                                         "rngState1",
@@ -2787,8 +2786,7 @@ void jl_init_types(void) JL_GC_DISABLED
                                         "sticky",
                                         "_isexception",
                                         "priority"),
-                        jl_svec(17,
-                                jl_any_type,
+                        jl_svec(16,
                                 jl_any_type,
                                 jl_any_type,
                                 jl_any_type,

src/julia.h

@@ -1920,10 +1920,7 @@ typedef struct _jl_task_t {
     jl_value_t *tls;
     jl_value_t *donenotify;
     jl_value_t *result;
-    jl_value_t *logstate;
     jl_function_t *start;
-    // 4 byte padding on 32-bit systems
-    // uint32_t padding0;
     uint64_t rngState[JL_RNG_SIZE];
     jl_value_t *hooks;
     _Atomic(uint8_t) _state;

src/task.c

@@ -1113,8 +1113,6 @@ JL_DLLEXPORT jl_task_t *jl_new_task(jl_function_t *start, jl_value_t *completion
     t->result = jl_nothing;
     t->donenotify = completion_future;
     jl_atomic_store_relaxed(&t->_isexception, 0);
-    // Inherit logger state from parent task
-    t->logstate = ct->logstate;
     // Fork task-local random state from parent
     jl_rng_split(t->rngState, ct->rngState);
     // there is no active exception handler available on this stack yet
@@ -1714,7 +1712,6 @@ jl_task_t *jl_init_root_task(jl_ptls_t ptls, void *stack_lo, void *stack_hi)
     ct->result = jl_nothing;
     ct->donenotify = jl_nothing;
     jl_atomic_store_relaxed(&ct->_isexception, 0);
-    ct->logstate = jl_nothing;
     ct->eh = NULL;
     ct->gcstack = NULL;
     ct->excstack = NULL;