From ea0e4a7f8da164ba02b2bab612ab73fae97dfbd9 Mon Sep 17 00:00:00 2001 From: Takafumi Arakaki Date: Thu, 7 May 2020 22:49:36 -0700 Subject: [PATCH] Implement context variables --- base/Base.jl | 3 +- base/contextvariables.jl | 415 ++++++++++++++++++++++++++++++++++++++ base/exports.jl | 7 + src/jltypes.c | 8 +- src/julia.h | 1 + src/task.c | 2 + stdlib/UUIDs/src/UUIDs.jl | 2 + test/choosetests.jl | 3 +- test/contextvariables.jl | 60 ++++++ 9 files changed, 496 insertions(+), 5 deletions(-) create mode 100644 base/contextvariables.jl create mode 100644 test/contextvariables.jl diff --git a/base/Base.jl b/base/Base.jl index c5dd34271493d..c8bc03d7e72ea 100644 --- a/base/Base.jl +++ b/base/Base.jl @@ -224,6 +224,8 @@ include("lock.jl") include("channels.jl") include("task.jl") include("weakkeydict.jl") +include("uuid.jl") +include("contextvariables.jl") # Logging include("logging.jl") @@ -332,7 +334,6 @@ include("initdefs.jl") include("threadcall.jl") # code loading -include("uuid.jl") include("loading.jl") # misc useful functions & macros diff --git a/base/contextvariables.jl b/base/contextvariables.jl new file mode 100644 index 0000000000000..25674eb717274 --- /dev/null +++ b/base/contextvariables.jl @@ -0,0 +1,415 @@ +# This file is a part of Julia. License is MIT: https://julialang.org/license + +function _uuid4 end # to be defined by UUIDs.jl + +function _ContextVar end + +# `ContextVar` object itself does not hold any data (except the +# default value). It is actually just a key into the task-local +# context storage. +""" + ContextVar{T} + +Context variable type. This is the type of the object `var` created by +[`@contextvar var`](@ref @contextvar). This acts as a reference to the +value stored in a task-local. The macro `@contextvar` is the only public +API to construct this object. + +!!! warning + + It is unspecified if this type is concrete or not. It may be + changed to an abstract type and/or include more type parameters in + the future. +""" +struct ContextVar{T} + name::Symbol + _module::Module + key::UUID + has_default::Bool + default::T + + global _ContextVar + _ContextVar(name, _module, key, ::Type{T}, default) where {T} = + new{T}(name, _module, key, true, default) + _ContextVar(name, _module, key, ::Type{T}) where {T} = new{T}(name, _module, key, false) +end + +_ContextVar(name, _module, key, ::Nothing, default) = + _ContextVar(name, _module, key, typeof(default), default) +_ContextVar(name, _module, key, ::Nothing) = _ContextVar(name, _module, key, Any) + +eltype(::Type{ContextVar{T}}) where {T} = T + +function show(io::IO, var::ContextVar) + print(io, ContextVar) + if eltype(var) !== Any && !(var.has_default && typeof(var.default) === eltype(var)) + print(io, '{', eltype(var), '}') + end + print(io, '(', repr(var.name)) + if var.has_default + print(io, ", ") + show(io, var.default) + end + print(io, ')') +end + +function show(io::IO, ::MIME"text/plain", var::ContextVar) + print(io, var._module, '.', var.name, " :: ContextVar") + if get(io, :compact, false) === false + print(io, " [", var.key, ']') + if get(var) === nothing + print(io, " (not assigned)") + else + print(io, " => ") + show(IOContext(io, :compact => true), MIME"text/plain"(), var[]) + end + end +end + +# The primitives that can be monkey-patched to play with new context +# variable storage types: +""" + merge_ctxvars(ctx::Union{Nothing,T}, kvs) -> ctx′:: Union{Nothing,T} + +!!! warning + + This is not a public API. This documentation is for making it easier to + experiment with different implementations of the context variable storage + backend, by monkey-patching it at run-time. + +The first argument `ctx` is either `nothing` or a dict-like object of type `T` where +its `keytype` is `UUID` and `valtype` is `Any`. The second argument `kvs` is an +iterable of `Pair{UUID,<:Union{Some,Nothing}}` values. Iterable `kvs` must have +length. + +If `ctx` is `nothing` and `kvs` is non-empty, `merge_ctxvars` creates a new +instance of `T`. If `ctx` is not `nothing`, it returns a shallow-copy `ctx′` of +`ctx` where `k => v` is inserted to `ctx′` for each `k => Some(v)` in `kvs` +and `k` is deleted from `ctx′` for each `k => nothing` in `kvs`. +""" +function merge_ctxvars(ctx, kvs) + # Assumption: eltype(kvs) <: Pair{UUID,<:Union{Some,Nothing}} + if isempty(kvs) + return ctx + else + # Copy-or-create-on-write: + vars = ctx === nothing ? Dict{UUID,Any}() : copy(ctx) + for (k, v) in kvs + if v === nothing + delete!(vars, k) + else + vars[k] = something(v) + end + end + isempty(vars) && return nothing # should we? + return vars + end +end + +get_task_ctxvars(t = current_task()) = t.ctxvars +set_task_ctxvars(t, ctx) = t.ctxvars = ctx +set_task_ctxvars(ctx) = set_task_ctxvars(current_task(), ctx) + +struct _NoValue end + +""" + get(var::ContextVar{T}) -> Union{Some{T},Nothing} + +Return `Some(value)` if `value` is assigned to `var`. Return `nothing` if +unassigned. +""" +function get(var::ContextVar{T}) where {T} + ctx = get_task_ctxvars() + if ctx === nothing + var.has_default && return Some(var.default) + return nothing + end + if var.has_default + return Some(get(ctx, var.key, var.default)::T) + else + y = get(ctx, var.key, _NoValue()) + y isa _NoValue || return Some(ctx[var.key]::T) + end + return nothing +end + +""" + getindex(var::ContextVar{T}) -> value::T + +Return the `value` assigned to `var`. Throw a `KeyError` if unassigned. +""" +function getindex(var::ContextVar{T}) where {T} + maybe = get(var) + maybe === nothing && throw(KeyError(var)) + return something(maybe)::T +end + +setindex!(var::ContextVar, value) = set!(var, Some(value)) +delete!(var::ContextVar) = set!(var, nothing) + +""" + set!(var::ContextVar, Some(value)) + set!(var::ContextVar, nothing) + +Set the value of context variable `var` to `value` or delete it. +""" +function set!(var::ContextVar{T}, value::Union{Some,Nothing}) where {T} + value = convert(Union{Some{T},Nothing}, value) + set_task_ctxvars(merge_ctxvars(get_task_ctxvars(), (var.key => value,))) + return var +end + +""" + genkey(__module__::Module, varname::Symbol) -> UUID + +Generate a stable UUID for a context variable `__module__.\$varname`. +""" +function genkey(__module__::Module, varname::Symbol) + pkgid = PkgId(__module__) + if pkgid.uuid === nothing + throw(ArgumentError( + "Module `$__module__` is not a part of a package. " * + "`@contextvar` can only be used inside a package.", + )) + end + fullpath = push!(collect(fullname(__module__)), varname) + if any(x -> contains(string(x), "."), fullpath) + throw(ArgumentError( + "Modules and variable names must not contain a dot:\n" * join(fullpath, "\n"), + )) + end + return uuid5(pkgid.uuid, join(fullpath, '.')) +end + +""" + @contextvar [local|global] var[::T] [= default] + +Declare a context variable named `var`. + +!!! warning + + Context variables declared with `global` does not work with `Distributed`. + +# Examples + +Top-level context variables needs to be declared in a package: + +``` +module MyPackage +@contextvar cvar1 +@contextvar cvar2 = 1 +@contextvar cvar3::Int +end +``` + +Context variables can be declared in local scope by using `local` prefix: + +```jldoctest +julia> function demo() + @contextvar local x = 1 + function f() + x[] += 1 + return x[] + end + return f() + end; + +julia> demo() +2 +``` + +To use `@contextvar` in a non-package namespace like REPL, prefix the variable +with `global`: + +```jldoctest global_vars +julia> @contextvar global X; + +julia> X[] = 1; + +julia> X[] +1 +``` +""" +macro contextvar(ex0) + ex = ex0 + qualifier = :const + if Meta.isexpr(ex, :local) + length(ex.args) != 1 && throw(ArgumentError("Malformed input:\n$ex0")) + ex, = ex.args + qualifier = :local + elseif Meta.isexpr(ex, :global) + length(ex.args) != 1 && throw(ArgumentError("Malformed input:\n$ex0")) + ex, = ex.args + qualifier = :global + end + if Meta.isexpr(ex, :(=)) + length(ex.args) != 2 && throw(ArgumentError("Unsupported syntax:\n$ex0")) + ex, default = ex.args + args = Any[esc(default)] + else + args = [] + end + if Meta.isexpr(ex, :(::)) + length(ex.args) != 2 && throw(ArgumentError("Malformed input:\n$ex0")) + ex, vartype = ex.args + pushfirst!(args, esc(vartype)) + else + pushfirst!(args, nothing) + end + if !(ex isa Symbol) + if ex === ex0 + throw(ArgumentError("Unsupported syntax:\n$ex0")) + else + throw(ArgumentError(""" + Not a variable name: + $ex + Input: + $ex0 + """)) + end + end + varname = QuoteNode(ex) + if qualifier === :const + key = genkey(__module__, ex) + else + # Creating a UUID at macro expansion time because: + # * It would be a memory leak because context variable storage can be + # filled with UUIDs created at run-time. + # * Creating it at run-time is doable with function-based interface like + # `ContextVar(:name, default)`. + key = _uuid4() + end + return Expr( + qualifier, + :($(esc(ex)) = _ContextVar($varname, $__module__, $key, $(args...))), + ) +end + +""" + with_context(f, var1 => value1, var2 => value2, ...) + with_context(f, pairs) + +Run `f` in a context with given values set to the context variables. Variables +specified in this form are rolled back to the original value when `with_context` +returns. It act like a dynamically scoped `let`. If `nothing` is passed as +a value, corresponding context variable is cleared; i.e., it is unassigned or +takes the default value. Use `Some(value)` to set `value` if `value` can be +`nothing`. + + with_context(f, nothing) + +Run `f` in a new empty context. All variables are rewind to the original values +when `with_context` returns. + +Note that + +```julia +var2[] = value2 +with_context(var1 => value1) do + @show var2[] # shows value2 + var3[] = value3 +end +@show var3[] # shows value3 +``` + +and + +```julia +var2[] = value2 +with_context(nothing) do + var1[] = value1 + @show var2[] # shows default (or throws) + var3[] = value3 +end +@show var3[] # does not show value3 +``` + +are not equivalent. +""" +function with_context(f, kvs::Pair{<:ContextVar}...) + orig = map(((k, _),) -> k => get(k), kvs) + set_context(kvs) + try + return f() + finally + set_context(orig) + end +end + +function with_context(f, ::Nothing) + ctx = get_task_ctxvars() + try + set_task_ctxvars(nothing) + return f() + finally + set_task_ctxvars(ctx) + end +end + +# Not using `set_context!` so that `snapshot` as an input makes sense. +""" + set_context(var1 => value1, var2 => value2, ...) + set_context(pairs) + set_context(snapshot::ContextSnapshot) + +Equivalent to `var1[] = value1`, `var2[] = value2`, and so on. The second form +expect an iterable of pairs of context variable and values. This function is +more efficient than setting individual context variables sequentially. Like +[`with_context`](@ref), `nothing` value means to clear the context variable +and `Some` is always unwrapped. + +A "snapshot" of the context returned from [`snapshot_context`](@ref) can +be specified as an input (the third form). + +# Examples +```jldoctest +julia> @contextvar global x + @contextvar global y; + +julia> set_context(x => 1, y => 2) + +julia> x[] +1 + +julia> y[] +2 + +julia> set_context([x => 10, y => 20]) + +julia> x[] +10 + +julia> y[] +20 +``` +""" +set_context(kvs::Pair{<:ContextVar}...) = set_context(kvs) +function set_context(kvs) + ctx = merge_ctxvars(get_task_ctxvars(), (k.key => v for (k, v) in kvs)) + set_task_ctxvars(ctx) + return +end + +struct ContextSnapshot{T} + vars::T +end + +# TODO: Do we need to implement `Dict{ContextVar}(::ContextSnapshot)`? +# This requires storing UUID-to-ContextVar mapping somewhere. + +""" + snapshot_context() -> snapshot::ContextSnapshot + +Get a snapshot of a context that can be passed to [`reset_context`](@ref) to +rewind all changes in the context variables. +""" +snapshot_context() = ContextSnapshot(get_task_ctxvars()) + +""" + reset_context(snapshot::ContextSnapshot) + +Rest the entire context of the current task to the state at which `snapshot` +is obtained via [`snapshot_context`](@ref). +""" +reset_context(snapshot::ContextSnapshot) = set_task_ctxvars(snapshot.vars) +set_context(snapshot::ContextSnapshot) = set_context(snapshot.vars) diff --git a/base/exports.jl b/base/exports.jl index 316025db9ce6c..c2b3fab2d21b6 100644 --- a/base/exports.jl +++ b/base/exports.jl @@ -38,6 +38,7 @@ export ComplexF64, ComplexF32, ComplexF16, + ContextVar, DenseMatrix, DenseVecOrMat, DenseVector, @@ -525,6 +526,7 @@ export mergewith, pairs, reduce, + set!, setdiff!, setdiff, setindex!, @@ -672,6 +674,11 @@ export timedwait, asyncmap, asyncmap!, + @contextvar, + reset_context, + set_context, + snapshot_context, + with_context, # channels take!, diff --git a/src/jltypes.c b/src/jltypes.c index ac2a60df9145c..5d9aa807fd5bd 100644 --- a/src/jltypes.c +++ b/src/jltypes.c @@ -2273,10 +2273,11 @@ void jl_init_types(void) JL_GC_DISABLED NULL, jl_any_type, jl_emptysvec, - jl_perm_symsvec(11, + jl_perm_symsvec(12, "next", "queue", "storage", + "ctxvars", "state", "donenotify", "result", @@ -2285,7 +2286,8 @@ void jl_init_types(void) JL_GC_DISABLED "logstate", "code", "sticky"), - jl_svec(11, + jl_svec(12, + jl_any_type, jl_any_type, jl_any_type, jl_any_type, @@ -2297,7 +2299,7 @@ void jl_init_types(void) JL_GC_DISABLED jl_any_type, jl_any_type, jl_bool_type), - 0, 1, 9); + 0, 1, 10); jl_value_t *listt = jl_new_struct(jl_uniontype_type, jl_task_type, jl_nothing_type); jl_svecset(jl_task_type->types, 0, listt); diff --git a/src/julia.h b/src/julia.h index a90aa6dc8b265..a3c256b7bf1a3 100644 --- a/src/julia.h +++ b/src/julia.h @@ -1728,6 +1728,7 @@ typedef struct _jl_task_t { jl_value_t *next; // invasive linked list for scheduler jl_value_t *queue; // invasive linked list for scheduler jl_value_t *tls; + jl_value_t *ctxvars; jl_sym_t *state; jl_value_t *donenotify; jl_value_t *result; diff --git a/src/task.c b/src/task.c index 9d88306dd4eee..abd01abd8b1d7 100644 --- a/src/task.c +++ b/src/task.c @@ -570,6 +570,7 @@ JL_DLLEXPORT jl_task_t *jl_new_task(jl_function_t *start, jl_value_t *completion t->next = jl_nothing; t->queue = jl_nothing; t->tls = jl_nothing; + t->ctxvars = ptls->current_task->ctxvars; // inherit parent context t->state = runnable_sym; t->start = start; t->result = jl_nothing; @@ -1093,6 +1094,7 @@ void jl_init_root_task(void *stack_lo, void *stack_hi) ptls->current_task->started = 1; ptls->current_task->next = jl_nothing; ptls->current_task->queue = jl_nothing; + ptls->current_task->ctxvars = jl_nothing; ptls->current_task->state = runnable_sym; ptls->current_task->start = NULL; ptls->current_task->result = jl_nothing; diff --git a/stdlib/UUIDs/src/UUIDs.jl b/stdlib/UUIDs/src/UUIDs.jl index ee2064605c7c8..36ab8811c4c5d 100644 --- a/stdlib/UUIDs/src/UUIDs.jl +++ b/stdlib/UUIDs/src/UUIDs.jl @@ -94,6 +94,8 @@ function uuid4(rng::AbstractRNG=Random.default_rng()) UUID(u) end +Base._uuid4(rng::AbstractRNG=Random.default_rng()) = uuid4(rng) + """ uuid5(ns::UUID, name::String) -> UUID diff --git a/test/choosetests.jl b/test/choosetests.jl index 8e8b4ae87ed6d..7abec10cd1b1d 100644 --- a/test/choosetests.jl +++ b/test/choosetests.jl @@ -54,7 +54,8 @@ function choosetests(choices = []) "checked", "bitset", "floatfuncs", "precompile", "boundscheck", "error", "ambiguous", "cartesian", "osutils", "channels", "iostream", "secretbuffer", "specificity", - "reinterpretarray", "syntax", "logging", "missing", "asyncmap", "atexit" + "reinterpretarray", "syntax", "logging", "missing", "asyncmap", "atexit", + "contextvariables", ] tests = [] diff --git a/test/contextvariables.jl b/test/contextvariables.jl new file mode 100644 index 0000000000000..9e345bfb705c5 --- /dev/null +++ b/test/contextvariables.jl @@ -0,0 +1,60 @@ +module TestContextVariables + +using Test + +@contextvar global cvar1 = 42 +@contextvar global cvar2::Int +@contextvar global cvar3 + +@testset "typed, w/ default" begin + ok = Ref(0) + @sync @async begin + with_context() do + @test cvar1[] == 42 + cvar1[] = 0 + @test cvar1[] == 0 + ok[] += 1 + @async begin + @test cvar1[] == 0 + ok[] += 1 + end + with_context(cvar1 => 1) do + @test cvar1[] == 1 + ok[] += 1 + end + @test cvar1[] == 0 + ok[] += 1 + end + end + @test ok[] == 4 +end + +@testset "typed, w/o default" begin + with_context() do + @test_throws InexactError cvar2[] = 0.5 + @test_throws KeyError cvar2[] + cvar2[] = 1.0 + @test cvar2[] === 1 + end +end + +@testset "untyped, w/o default" begin + with_context() do + cvar3[] = 1 + @test cvar3[] === 1 + cvar3[] = 'a' + @test cvar3[] === 'a' + end +end + +@testset "show" begin + @test endswith(sprint(show, cvar1), "ContextVar(:cvar1, 42)") + @test endswith(sprint(show, cvar2), "ContextVar{$Int}(:cvar2)") + @test endswith(sprint(show, cvar3), "ContextVar(:cvar3)") + @test endswith( + sprint(show, @contextvar local x::Union{Missing,Int64} = 1), + "ContextVar{Union{Missing, Int64}}(:x, 1)", + ) +end + +end # module