RFC: Context variables

[tkf]

This PR proposes context variables API that can be used to propagate context-dependent information across task boundaries. It is conceptually similar to task_local_storage but with a main difference that it "copies" all key-value pairs to child Tasks. The context variables are conceptually similar to dynamically scoped variables.

Motivations

There are several places context variables can be useful or required.

ENV

It was proposed to use task_local_storage to fix thread-safety of withenv #34726 (comment) by maintaining task-local copy (or overlay) of ENV. However, it would mean that the environment variables cannot cross task boundaries. Context variable can fix this shortcoming.

@testset

@testset uses task_local_storage to track current active test set. However, following example does not work (prints No tests) because the information of the current test does not propagate to the child task:

@testset begin
    @sync @async @test true
end

Logging

Task has the logstate field that propagates to child tasks; i.e., it works as a hard-coded context variable. Once context variable handling is sufficiently matured, it may be possible to eliminate the special handling from Task and use a context variable for Task. Furthermore, context variable allows users to develop logging-like interfaces.

Custom worker pool abstraction

In #35757 and Propagation of available/assigned worker-IDs in hierarchical computations? - Domains / Julia at Scale - JuliaLang, it was brought up that propagating "computation resources" (thread/process pools, etc.) across tasks and processes is required for implementing custom worker pool interfaces.

Misc.

Request-related data, such as security tokens and request data in web applications, language context for gettext, etc.

--- Rationale section in PEP 0567

It would be useful for implementing a better nestable progress information handling JuliaLogging/ProgressLogging.jl#13 (comment).

Other languages

Implicit context

• Python has contextvars. See PEP 567 and PEP 550 for discussion.
• .NET has ExecutionContext.
• Kotlin has CoroutineContext.

Explicit context

• Go has context.

Proposed design

I propose an API with the following basic usages.

• Declare context variable: @contextvar x
• Get the value: x[]
• Set the value: x[] = value

For a tutorial and the full reference API of the proposed design, see https://tkf.github.io/ContextVariablesX.jl/dev/

Internally, @contextvar creates an instance of ContextVar which is defined as

struct ContextVar{T}
    key::UUID
    default::T
    ...
end

Then x[] and x[] = value invoke current_task().ctxvars[x.key] and an "immutable version" of current_task().ctxvars[x.key] = value, respectively (roughly speaking).

This has a couple of nice properties:

1. x[] can be inferred
2. x is forced to be namespaced (i.e., it has to exist in some module name space.)
3. x can be backed up by an efficient concrete key type (e.g., UUID)
4. x allows small-size optimization when the value type can be inlined into the context storage (in principle)

See also #35757 which already contains some discussion on this API.

[tkf]

It is called contextvars.copy_context in Python and ExecutionContext.CreateCopy in .NET. But I was not sure if "copy" is the right word when the underlying data is (treated as) immutable.

[tkf]

I couldn't find a proper way to "import" stdlib to Base. Is this an OK approach?

[twavv]

Why not just rand(UInt64) or something involving objectid?

[twavv]

It looks to me like most of the uses of key could just be replaced with objectid(::ContextVar)

[tkf]

Actually, the main reason I wanted to use UUID was uuid5 because I can then use it to get deterministic and namespaced key for "global const" context variables (while uuid4 is used only for local and "global non-const" context variables):

julia/base/contextvariables.jl

Line 181 in ea0e4a7

return uuid5(pkgid.uuid, join(fullpath, '.'))

(where join(fullpath, '.') is something like "PackageName.SubModuleName.varname")

Making the key for "global const" context variable is important for Distributed (otherwise different processes can't agree on the key). So no, we can't use objectid.

[c42f]

I couldn't find a proper way to "import" stdlib to Base. Is this an OK approach?

Yeah it's a tricky problem: you'd like a uuid4 to use a non-terrible RNG. But we can't pull all of Random into Base.

Another workaround for this would be to use root_module to look up the UUIDs module without loading it. There's precedent for this in the way LibGit2 is looked up in one place:

julia/base/methodshow.jl

Line 316 in 97e3fe8

LibGit2 = root_module(libgit2_id)

In either case, it's not great to have Base missing some functionality when compiled without a stdlib present.

[tkf]

I see. Thanks for the explanation. How about vendoring libc-based uuid4 in Base?

function _uuid4()
    u = reinterpret(UInt128, [Libc.rand() for _ in 1:(sizeof(UInt128) ÷ sizeof(Cint))])[1]
    u &= 0xffffffffffff0fff3fffffffffffffff
    u |= 0x00000000000040008000000000000000
    UUID(u)
end

[c42f]

What quality of random numbers do we need for this? I guess a large period and good seeding are desired, as the UUIDs may end up in a distributed system? Actually this would suggest that the global RNG from Random isn't a great choice for UUIDs anyway because the user may control the seeding for reasons unrelated to UUIDs.

[tkf]

Actually this would suggest that the global RNG from Random isn't a great choice for UUIDs anyway because the user may control the seeding for reasons unrelated to UUIDs.

Oooh, that's a good point. Maybe we should use Random.RandomDevice() instead of Random.default_rng() in UUIDs? See also #32954 (comment)

But I think _uuid4 above may be OK as the intial implementation since:

• It is used only for @contextvar local var and @contextvar global var (the latter is available mainly for exploration in REPL and testing). But the "main entrypoint" is @contextvar var which uses uuid5.
• The UUIDs are generated during macro expansion time. So, messing with it requires you to call Libc.srand(my_seed) at the top-level of your module. It's totally possible but somewhat unlikely.

In the long run, I think we'd need something like Base.CoreRandom.RandomDevice.

[twavv]

I don't see why you've "scoped" ContextVars to modules the way you have. Why not just scope them using normal objects (i.e., why do the vars themselves need to know about what module they were declared in?)? This is the way it works in Python at least.

[tkf]

As I said #35833 (comment), that's a requirement for this to work with Distributed across different machines. Consider

module MyPackage
const KEY = uuid4()
end

MyPackage.KEY changes every time I re-compile it. So, a key generated similar way (including objectid at compile-time and run-time) can't be used to communicate with Distributed workers on different machines.

[c42f]

Ah very nice chaining these UUIDs. I was just coming here to suggest it should be done this way. But you've already done it :-)

[tkf]

Good to know that we independently arrived at the same solution :) It's an indication that this solution makes sense.

[c42f]

Well perhaps I shouldn't be so definitive as to say "should" ... but I think it makes sense :-)

[c42f]

Speaking of UUIDs, where does uuid5 come from? I guess it will need to be moved it into base/uuid.jl from the UUIDs stdlib?

[tkf]

There is a fake uuid5 definition for loading:

julia/base/loading.jl

Lines 103 to 116 in 8f512f3

	# fake uuid5 function (for self-assigned UUIDs)
	# TODO: delete and use real uuid5 once it's in stdlib
	function uuid5(namespace::UUID, key::String)
	u::UInt128 = 0
	h = hash(namespace)
	for _ = 1:sizeof(u)÷sizeof(h)
	u <<= sizeof(h) << 3
	u |= (h = hash(key, h))
	end
	u &= 0xffffffffffff0fff3fffffffffffffff
	u |= 0x00000000000050008000000000000000
	return UUID(u)
	end

[c42f]

Ah I see. Not your fault at all, but ick!

It seems some of the hashing stuff should move back into Base (at least some minimal parts of the implementation of SHA1 and RandomDevice, not the whole stdlib). Good hashing and randomness is core functionality which isn't really optional in some of these cases.

[tkf]

Here is a PR to move RandomDevice to Base: #35894

[tkf]

Example usecase: circular reference detection

As a short somewhat-real usecase, here is how you can use the context variable to detect circular references.

Consider this innocent-looking show definition:

struct WAT
    value
end

Base.show(io::IO, x::WAT) = print(io, "WAT(", repr(x.value), ")")

However, this causes a stack overflow with

x = Any[WAT(Any[1])]
x[1].value[1] = x
@show x  # throws StackOverflowError

This is because current circular reference detection relies on IOContext. So, using repr instead of show "cuts" IOContext propagation (this is why you should be using repr(x; context=io) but this is very error prone) and it fools Base.show_circular.

We can use a context variable instead of IOContext to propagate the history of objects that are already shown:

function nocircular(f, io, @nospecialize(x))
    @contextvar local shown::Cons
    history = get(shown)
    d = 1
    for y in something(history, ())
        if x === y
            print(io, "#= circular reference @-$d =#")
            return
        end
        d += 1
    end
    with_context(f, shown => Some(Cons(x, something(history, Some(nothing)))))
end

struct Cons
    car
    cdr
end

Base.iterate(list::Cons) = iterate(list, list)
Base.iterate(::Cons, list::Cons) = (list.car, list.cdr)
Base.iterate(::Cons, ::Nothing) = nothing

We can use this to implement our show method of objects that may contain circular references:

mutable struct Mutable
    value
end

function Base.show(io::IO, x::Mutable)
    nocircular(io, x) do
        print(io, "Mutable(")
        show(io, x.value)
        print(io, ")")
    end
end

With this approach, circular references can be detected without using IOContext:

x = Mutable(WAT(Mutable(1)))
x.value.value.value = x
@show x  # => Mutable(WAT(Mutable(#= circular reference @-2 =#)))

Of course, it works even if you switch the task inside show:

struct WAAAAAT
    value
end

Base.show(io::IO, x::WAAAAAT) = print(io, "WAAAAAT(", fetch(@async repr(x.value)), ")")

x = Mutable(WAAAAAT(Mutable(1)))
x.value.value.value = x
@show x  # => Mutable(WAAAAAT(Mutable(#= circular reference @-2 =#)))

(Edit: Actually, nocircular being io-independent is not really nice property. show(io, x) could be called for a completely different purpose in a given call chain.)

[c42f]

I just found a rather nice post on the use and abuse of Go's equivalent Context.Value:

https://medium.com/@cep21/how-to-correctly-use-context-context-in-go-1-7-8f2c0fafdf39

They argue quite persuasively against Context.Value, so it's a thought-provoking read. In particular, they argue against using it for request-related data as mentioned in the PEP-567 rationale, preferring a closure. I think they're right about that!

I'm not arguing against ContextVar btw — I believe we should eventually have something to replace logstate — but it does make me wonder whether we should be cautious about making it too convenient ;-)

[tkf]

A bit tangential, but I find the section "Does Context.Value even belong?" interesting. I haven't realized that the cancellation context (= nursery) and the context variable is so tightly coupled in Go. Though given the focus on server-side application in the context package documentation (and how Go is used in general), maybe it's not so puzzling.

Anyway, I agree dynamically scoped variables in any form can be used to write bad programs. It's more or less just a "less bad" form of global variables, in the end. For example, you can put a nursery in the context variable to easily create dynamically-scoped nursery. We already have global but people don't use it often. So, I'd hope it'd be used with caution.

[c42f]

A bit tangential, but I find the section "Does Context.Value even belong?" interesting.

Yes I thought this was interesting, personally I feel it's a weird mixture of abstractions :-) But perhaps it's the best which can be done with the desire to explicitly pass no more than one context, and with the limitations of Go interfaces.

A couple more half-formed thoughts about the Julia interface:

• I think the difference between const, local and global versions may be perplexing for users, especially because in all cases they are global to the task (IIUC the local contextvar doesn't behave like a local variable). Could we get away with only the const version? It seems a bit less convenient but might be sufficient in practice?
• I like the look of with_context, it's pretty clear that the variable is set in dynamic scope. Reading the variables is less nice somehow, as the syntax doesn't suggest where these variables come from, somewhat in contrast to task_local_storage(key). I guess it's no different from normal global variables though.

I also think it worth considering what API we're able to offer if we want to replace logstate with this system and have it be just as fast or faster. In that respect I like that you've documented @contextvar as the only official way to create these, it gives a lot of implementation flexibility. I think the compiler could model dynamic scope enough to allow the load of a contextvar to be hoisted out of any loop. (I think it would only be necessary that the same contextvar is not set in the intervening code, and that the contextvar use is inlined into the loop body. We might need to make the UUID part of the type for that to work neatly, but that seems it could be done in future as necessary.)

[tkf]

Could we get away with only the const version?

Yeah, I think that's what you need most of the time. The global version is just too handy to use in the REPL. Otherwise, you have to create a full blown package just for playing with the context variables.

The local version is least important but it's needed (or it's the easiest to do it) if we want to support context variables in scripts that works with Distributed, without creating a package. This is because, since local ContextVar object (i.e., UUID) is captured by the closure, it'd be sent to the peer processes.

I think we can special-case Main to support the const version in the REPL. But I imagine that would be pretty hairy (involving finalize?).

Another aspect is that, even if we only support the const version, people are free to experiment with the internal functions and create the local/global version themselves. So maybe it's a wise choice to play conservative and only provide the const version in the first version of the API.

Reading the variables is less nice somehow, as the syntax doesn't suggest where these variables come from, somewhat in contrast to task_local_storage(key).

Do you mean the x[] syntax? I just thought it's kind of cute way to do it. We can always make it uglier to make it explicit (e.g., get_context_variable(x)).

I guess it's no different from normal global variables though.

Maybe we should recommend using uppercase for context variables, too?

I think it would only be necessary that the same contextvar is not set in the intervening code, and that the contextvar use is inlined into the loop body.

Yeah, I'd hope so. Context variable lookup is just a dictionary look up in the end. So, if the compiler can assume it's effect-free, I suppose it can hoist it out?

We might need a better dictionary implantation, though. For now, I'm just using Dict in a copy-on-write manner. I wonder if it's better to use immutable data structures like HAMT backed up with Base.Experimental.Const(::Vector). I made this part "pluggable" so that people can easily experiment with it; i.e., you can monkey-patch a single function merge_ctxvars to use an alternative storage type. It's not a public stable API, of course.

[c42f]

I got a chance to chat to @JeffBezanson about this. Jeff had some great questions which I was ill-prepared for (and I know @tkf could have answered better) but it was productive nonetheless. Here's a summary:

• The local/const/global thing came up again as seeming somewhat complex and confusing, and I couldn't describe the distinction off the top of my head. (I'm pretty sure I did understand this a month ago, but I found I'd largely forgotten what it was all about!)
• Jeff mentioned that UUIDs seemed kinda heavy weight as keys. (Agreed, though I'm not sure whether there are alternatives in the distributed setting?)
• Having ContextVar storage as a Dict didn't seem great. I mentioned that I thought this was just a simple reference implementation, with the internal merge_ctxvars designed to allow experimentation and gradual improvement.
• Obviously ContextVars and task local storage really do the same thing, so Jeff commented that having two of them is kind of bad, and that taking up an extra pointer in the Task isn't great. My thought/comment about this is that we just need to be sure that ContextVars are a satisfying API covering all the cases that task local storage do (and more!), and that we might tolerate the extra field for a few releases until we can delete the current task local storage API. If we can delete the logstate field we'd also make up for it ;-)

I think the most important point discussed was ability to implement this efficiently in the future: exposing set! kills the compiler's ability to generate good code for context vars — it cannot reason locally about their value if a call to a child function can mutate the state seen by the parent. This seems like a general problem for any proposed implementation of dynamic scope, so I'll write some further thoughts about it in #35757.

[tkf]

The local/const/global thing came up again as seeming somewhat complex and confusing

I think the best way forward might be to just expose the const version (i.e., the default one) as the API at least for the first version. Adding others is pretty easy because there is no change required for the backend. For example, it can be provided as a package (provided that the package author is brave enough to rely on the implementation detail).

Anyway, here is how they work in the reference implementation:

• const (@contextvar x): Usable at top-level scope of packages. It cannot be used outside packages because the UUID has to be salted with the package UUID. Usable with Distributed.jl.
• local (@contextvar local x): To be used inside functions. Usable with Distributed.jl because the closure will capture ContextVar object which includes the UUID.
• global (@contextvar global x): To be used in REPL or scripts. Not usable with Distributed.jl because the UUID is not sent to remote.

It'd be possible to merge const and local if a macro can know if they are in the top-level scope or in the function-level scope. Merging const and global may be possible since the macro can detect if the package UUID is available or not. If the package UUID is not available, it can generate UUID4 (and maybe even broadcast the definition via Distributed.jl).

UUIDs seemed kinda heavy weight as keys

I guess we can swap it with something better once we find out a better way to do it? For example, maybe we can have a global "key broker" that maps an UUID to a small key UInt32 or UInt64 and store the small key in the Task-local storage. When a ContextVar object is instantiated (also when it's loaded from precompilation cache), it asks the key broker for the short key and store it in its field. It'd be tricky to make it Distributed.jl-compatible, though.

Having ContextVar storage as a Dict didn't seem great. I mentioned that I thought this was just a simple reference implementation, with the internal merge_ctxvars designed to allow experimentation and gradual improvement.

Yeah, I just want to get API correct and make it easy to experiment with better backend storage. HAMT is an obvious candidate.

[c42f]

HAMT is an obvious candidate.

Ah yes, I see the wikipedia page mentions persistent variants of HAMTs. This seems like the thing we want 👍

[oschulz]

@contextvar x [...] x is forced to be namespaced

Would that allow code to access context via @contextvar SomeModule.x, to avoid clashes with local names?

[tkf]

How @contextvar SomeModule.x behaves? Just to be clear, by namespacing I mean:

module A
    module B
        @contextvar x

        get_x() = x[]

        function shadowing()
            x = 1    # shadows context variable `x`
            sum = 1  # shadows `Base.sum`
            return get_x()
        end
    end

    get_B_x() = B.x[]
end

# These work and equivalent:
A.B.x[]
A.B.get_x()
A.get_B_x()
A.B.shadowing()

# These do not work:
x[]
A.x[]

@contextvar x behaves just like an ordinal global (const) name. You'd already make sure that local variables don't shadow functions, constants, and globals. I don't think we need something extra here.

[oschulz]

Just to be clear, by namespacing I mean [...]

Oh, of course - very nice!

How would you handle the case where modified context is to be passed to tasks/workers? I am in a certain context, but I want to modify part of that context in a different way for each worker (without changing by current context). E.g. to partition resources listed in the context among tasks/workers, etc. How should be express that, API-wise?

[c42f]

Is there anything blocking the specific changes in this PR?

I think API changes from ContextVariablesX still need porting in here:

This PR still includes set! but I assume it'd be removed in an update?

Ah, yes, I forgot that I didn't backport the change. Currently, the source of truth is ContextVariablesX.jl since it's much easier to tweak code in a package.

[Tokazama]

Is this something that we want to also support context in the sense that Cassette.Context does, or is this solely intended for concurrency?

[oschulz]

or is this solely intended for concurrency?

I don't think it's only useful for concurrency, it's about propagating argument-like "settings" that can't be well handled via function arguments. Things like logger, progress monitor, choice of computational resources, etc. This is of course very important for thread- and process-level parallelism, but I think it would benefit serial chains of function calls as well.

[MasonProtter]

Just to add another potential motivation for this sort of thing, the way that propagate_inbounds works is basically a hacky, single use implementation of a context variable.

It'd be pretty cool if we could have context variables be lightweight enough that things like flagging that it's okay to turn off error checks could use this mechanism.

[c42f]

I don't think propagate_inbounds is very similar — it only propagates through one additional level of inlining. This "locality" is an important feature, especially in the reliable implementation of higher order functions. For example, here's what can go wrong with too much @propagate_inbounds: andyferris/Dictionaries.jl#64

[Tokazama]

If we had an index set and could say that "in the context of this collection it is inbounds" then limiting to one level of propagation wouldn't be necessary (and ensure not mutations too). Unless we need to further process indices we just create a view, so it's not really a critical example to get working. It's more of a proof of concept for propagating other relationships

[NHDaly]

I want to ask about another related usage of Context in Golang:

• Cancellation and timeout

https://pkg.go.dev/context#WithCancel

The Go Context object can also contain an indication of whether a request has been cancelled by the client, and/or whether a timeout/duration has been exceeded.

This requires the server code to thread the Context through all of the functions on the server, across module boundaries, and check for cancellation at very places.

Could we also cover that logic with a ContextVariable as proposed here / in https://github.com/tkf/ContextVariablesX.jl?
In particular: there a few extra pieces besides just inheriting named variables across task boundaries:

• WithCancel returns a copy of parent with a new Done channel. The returned context's Done channel is closed when the returned cancel function is called or when the parent context's Done channel is closed, whichever happens first.

  Canceling this context releases resources associated with it, so code should call cancel as soon as the operations running in this Context complete.

I'm thinking maybe this could be done by having a function like with_cancel(ctx) do that listens to the parent context's cancellation variable and propagates that state to its own?

The main place I see where this might not mesh well with the current design is:

• The current design doesn't have a local context variable to copy from, nor assign to, so i don't see how you could keep multiple separate cancellation context handles?

It seems like supporting multiple separate cancellation contexts is a bit at-odds with the feature this package seems to provide, which is avoiding having to thread the variable through all your code by instead making it a global variable.
Which I must admit is very appealing.

Maybe the approach is to have these cancellation contexts also be named context variables? Like

@contextvar request_cancelled = CancellationContext()

And then every new incoming server request could set its own request_cancelled variable?

Has anyone else thought about this already? We're working on implementing user-requested Transaction Cancellation in our async web server at RAI now, and it occurred to me that these go-style Context objects seem like the right design.

[c42f]

I think you can use an implementation of dynamic scope which is inherited across @async/@spawn (such as this) to build a scope-based cancellation system.

However I don't think there's a need to bake anything like that into the context variables system itself.

Safe and convenient cancellation is a very deep topic and it's not clear to me that any language has fully solved this problem. But I still think the way that cancellation is handled in structured concurrency systems is the best option seen so far. On that note, you should definitely check out https://github.com/JuliaConcurrent/Julio.jl if you haven't seen it already :-)

[vchuravy]

One comment from triage was that the @contextvar is slightly to magical. Would it be feasible do have something akin to Ref e.g. ContextVar so that the v[] is "expected"?

[tkf]

I totally agree that anything that can be implemented as a macro must be implemented without a macro. However, I couldn't find any other ways to provide some desired properties without macro. (This was precisely what was discussed in #35833 (comment))

Distributed computing is the main reason why the primary construction API is provided through macro. With macro-based construction, it supports the stability of the context variable key across precompilations. This key is used for identifying the value assigned to the context variable within a context (say, a task-local Dict{UUID,Any}). I believe it is desirable to maintain the identity of context variable MyPackage.MyModule.MY_VARIABLE across different processes of different machines in order to use this mechanism for distributed computing. To support this, the context variable key is generated using uuid5(PkgId(MyPackage).uuid, "MyPackage.MyModule.MY_VARIABLE"). This way, context variables precompiled in different machines have exactly the same UUID, thanks to the UUID-based packaging infrastructure. This makes copying entire context variables across machines as "easy" ¹ as copying the key-value pairs since the keys have the same UUIDs across machines. If we used (say) uuid4 from a constructor invocation ContextVariable() the key associated with this context variable changes for each precompilation. This still works for single-process application. However, a context variable cannot be guaranteed to be identified across multiple processes anymore.

So, it would be nice if triage can discuss:

(1) Do we want to support sending values assigned to context variables across distributed processes?

(2) If so, is it reasonable to use what is implemented in ContextVariablesX?

Even if the answer to the second question is "maybe?", I think it would be better to use the current macro-based API even we ended up using different implementation strategies because it is constrained enough to allow the implementation used by ContextVariablesX.

Footnotes

1. Of course, copying everything unconditionally is a bad idea unless only immutable values -- not mutable objects -- are stored in context variables. Some additional mechanism must be built upon this PR to support automatic cross-process propagation. However, stable context variable key is still a prerequisite for deciding cross-process propagation based on context variable identity (and not merely based on values and/or types). ↩

[oschulz]

I believe it is desirable to maintain the identity of context variable MyPackage.MyModule.MY_VARIABLE across different processes of different machines in order to use this mechanism for distributed computing.

Yes please!

[oschulz]

To elaborate: I think distributed use cases would profit from context variables very much. And if the machinery doesn't forward context variables through remote calls and the like it will be very hard for the user to do so. And that will mean that code that relies on context variables would malfunction when used distributed - so lot's of special handling would need to be added manually for the distributed case, making code much less generic.

[simonbyrne]

I broadly like the overall design here, but I have a minor suggestion: provide some mechanism for specifying how context variables should be passed when spawning subtasks on local and remote processes.

This was motivated by #48121: the idea is that when you spawn a task on a remote process, you could wrap the logger in a RemoteLogger (so that any log messages are passed back to the main process).

[vchuravy]

Another implementation is ScopedVariables https://openjdk.org/jeps/446

[vchuravy]

Replaced by #50958