Add ScopedVariables

ScopedVariables are containers whose observed value depends the current
dynamic scope. This implementation is inspired by https://openjdk.org/jeps/446

A scope is introduced with the `scoped` function that takes a lambda to
execute within the new scope. The value of a `ScopedVariable` is
constant within that scope and can only be set upon introduction
of a new scope.

Scopes are propagated across tasks boundaries.

In contrast to #35833 the storage of the per-scope data is assoicated
with the ScopedVariables object and does not require copies upon scope
entry. This also means that libraries can use scoped variables without
paying for scoped variables introduces in other libraries.

Finding the current value of a ScopedVariable, involves walking the
scope chain upwards and checking if the scoped variable has a value
for the current or one of its parent scopes. This means the cost of
a lookup scales with the depth of the dynamic scoping. This could be
amortized by using a task-local cache.

Author: vchuravy

base/Base.jl

@@ -330,10 +330,6 @@ using .Libc: getpid, gethostname, time, memcpy, memset, memmove, memcmp
 const libblas_name = "libblastrampoline" * (Sys.iswindows() ? "-5" : "")
 const liblapack_name = libblas_name
 
-# Logging
-include("logging.jl")
-using .CoreLogging
-
 # Concurrency (part 2)
 # Note that `atomics.jl` here should be deprecated
 Core.eval(Threads, :(include("atomics.jl")))
@@ -343,6 +339,14 @@ include("task.jl")
 include("threads_overloads.jl")
 include("weakkeydict.jl")
 
+# ScopedVariables
+include("scopedvariables.jl")
+using .ScopedVariables
+
+# Logging
+include("logging.jl")
+using .CoreLogging
+
 include("env.jl")
 
 # functions defined in Random

base/boot.jl

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

base/exports.jl

@@ -648,6 +648,10 @@ export
     sprint,
     summary,
 
+# ScopedVariable
+    scoped,
+    ScopedVariable,
+
 # logging
     @debug,
     @info,

base/logging.jl

@@ -492,8 +492,10 @@ end
 
 LogState(logger) = LogState(LogLevel(_invoked_min_enabled_level(logger)), logger)
 
+const CURRENT_LOGSTATE = ScopedVariable{Union{Nothing, LogState}}(nothing)
+
 function current_logstate()
-    logstate = current_task().logstate
+    logstate = CURRENT_LOGSTATE[]
     return (logstate !== nothing ? logstate : _global_logstate)::LogState
 end
 
@@ -506,17 +508,7 @@ end
     return nothing
 end
 
-function with_logstate(f::Function, logstate)
-    @nospecialize
-    t = current_task()
-    old = t.logstate
-    try
-        t.logstate = logstate
-        f()
-    finally
-        t.logstate = old
-    end
-end
+with_logstate(f::Function, logstate) = scoped(f, CURRENT_LOGSTATE => logstate)
 
 #-------------------------------------------------------------------------------
 # Control of the current logger and early log filtering

base/scopedvariables.jl

@@ -0,0 +1,128 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+module ScopedVariables
+
+export ScopedVariable, scoped
+
+mutable struct Scope
+    const parent::Union{Nothing, Scope}
+end
+
+current_scope() = current_task().scope::Union{Nothing, Scope}
+
+"""
+    ScopedVariable(x)
+
+Create a container that propagates values across scopes.
+Use [`scoped`](@ref) to create and enter a new scope.
+
+Values can only be set when entering a new scope,
+and the value referred to will be constant during the
+execution of a scope.
+
+Dynamic scopes are propagated across tasks.
+
+# Examples
+```jldoctest
+julia> const svar = ScopedVariable(1);
+
+julia> svar[]
+1
+
+julia> scoped(svar => 2) do
+           svar[]
+       end
+2
+```
+
+!!! compat "Julia 1.11"
+    This method requires at least Julia 1.11. In Julia 1.7+ this
+    is available from the package ScopedVariables.jl.
+"""
+mutable struct ScopedVariable{T}
+    const values::WeakKeyDict{Scope, T}
+    const initial_value::T
+    ScopedVariable{T}(initial_value) where {T} = new{T}(WeakKeyDict{Scope, T}(), initial_value)
+end
+ScopedVariable(initial_value::T) where {T} = ScopedVariable{T}(initial_value)
+
+Base.eltype(::Type{ScopedVariable{T}}) where {T} = T
+
+function Base.getindex(var::ScopedVariable{T})::T where T
+    scope = current_scope()
+    if scope === nothing
+        return var.initial_value
+    end
+    @lock var.values begin
+        while scope !== nothing
+            if haskey(var.values.ht, scope)
+                return var.values.ht[scope]
+            end
+            scope = scope.parent
+        end
+    end
+    return var.initial_value
+end
+
+function Base.show(io::IO, var::ScopedVariable)
+    print(io, ScopedVariable)
+    print(io, '{', eltype(var), '}')
+    print(io, '(')
+    show(io, var[])
+    print(io, ')')
+end
+
+function __set_var!(scope::Scope, var::ScopedVariable{T}, val::T) where T
+    # PRIVATE API! Wrong usage will break invariants of ScopedVariable.
+    if scope === nothing
+        error("ScopedVariable: Currently not in scope.")
+    end
+    @lock var.values begin
+        if haskey(var.values.ht, scope)
+            error("ScopedVariable: Variable is already set for this scope.")
+        end
+        var.values[scope] = val
+    end
+end
+
+"""
+    scoped(f, var::ScopedVariable{T} => val::T)
+
+Execute `f` in a new scope with `var` set to `val`.
+"""
+function scoped(f, pair::Pair{<:ScopedVariable{T}, T}) where T
+    @nospecialize
+    ct = Base.current_task()
+    current_scope = ct.scope::Union{Nothing, Scope}
+    try
+        scope = Scope(current_scope)
+        __set_var!(scope, pair...)
+        ct.scope = scope
+        return f()
+    finally
+        ct.scope = current_scope
+    end
+end
+
+"""
+    scoped(f, vars...::ScopedVariable{T} => val::T)
+
+Execute `f` in a new scope with each scoped variable set to the provided `val`.
+"""
+function scoped(f, pairs::Pair{<:ScopedVariable}...)
+    @nospecialize
+    ct = Base.current_task()
+    current_scope = ct.scope::Union{Nothing, Scope}
+    try
+        scope = Scope(current_scope)
+        for (var, val) in pairs
+            __set_var!(scope, var, val)
+        end
+        ct.scope = scope
+        return f()
+    finally
+        ct.scope = current_scope
+    end
+end
+
+end # module ScopedVariables

src/jltypes.c

@@ -3231,7 +3231,7 @@ void jl_init_types(void) JL_GC_DISABLED
                                         "storage",
                                         "donenotify",
                                         "result",
-                                        "logstate",
+                                        "scope",
                                         "code",
                                         "rngState0",
                                         "rngState1",

src/julia.h

@@ -2012,7 +2012,7 @@ typedef struct _jl_task_t {
     jl_value_t *tls;
     jl_value_t *donenotify;
     jl_value_t *result;
-    jl_value_t *logstate;
+    jl_value_t *scope;
     jl_function_t *start;
     // 4 byte padding on 32-bit systems
     // uint32_t padding0;

src/task.c

@@ -1068,8 +1068,8 @@ 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;
+    // Inherit scope from parent task
+    t->scope = ct->scope;
     // 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
@@ -1670,7 +1670,7 @@ 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->scope = jl_nothing;
     ct->eh = NULL;
     ct->gcstack = NULL;
     ct->excstack = NULL;

test/choosetests.jl

@@ -29,6 +29,7 @@ const TESTNAMES = [
         "channels", "iostream", "secretbuffer", "specificity",
         "reinterpretarray", "syntax", "corelogging", "missing", "asyncmap",
         "smallarrayshrink", "opaque_closure", "filesystem", "download",
+        "scopedvariables",
 ]
 
 const INTERNET_REQUIRED_LIST = [

test/scopedvariables.jl

@@ -0,0 +1,48 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+const svar1 = ScopedVariable(1)
+
+@testset "errors" begin
+    var = ScopedVariable(1)
+    @test_throws MethodError var[] = 2
+    scoped() do
+        @test_throws MethodError var[] = 2
+    end
+end
+
+const svar = ScopedVariable(1)
+@testset "inheritance" begin
+    @test svar[] == 1
+    scoped() do
+        @test svar[] == 1
+        scoped() do
+            @test svar[] == 1
+        end
+        scoped(svar => 2) do
+            @test svar[] == 2
+        end
+        @test svar[] == 1
+    end
+    @test svar[] == 1
+end
+
+const svar_float = ScopedVariable(1.0)
+
+@testset "multiple scoped variables" begin
+    scoped(svar => 2, svar_float => 2.0) do
+        @test svar[] == 2
+        @test svar_float[] == 2.0
+    end
+end
+
+import Base.Threads: @spawn
+@testset "tasks" begin
+    @test fetch(@spawn begin
+        svar[]
+    end) == 1
+    scoped(svar => 2) do
+        @test fetch(@spawn begin
+            svar[]
+        end) == 2
+    end
+end