Discussion: Make ThreadContextElement or its analog available on all platforms

[eymar]

Currently, ThreadContextElement interface is available only for jvm and it has 2 functions: 1st invoked before a coroutine is resumed, 2nd invoked after that coroutine is suspended.

These functions should never throw an exception:

Thrown exceptions will leave coroutine which context is updated in an undefined state and may crash an application.

public interface ThreadContextElement<S> : CoroutineContext.Element {
    /**
     * This function is invoked before the coroutine in the specified [context] is resumed in the current thread
     * when the context of the coroutine this element.
     */
    public fun updateThreadContext(context: CoroutineContext): S

    /**
     * This function is invoked after the coroutine in the specified [context] is suspended in the current thread
     * if [updateThreadContext] was previously invoked on resume of this coroutine.
     */
    public fun restoreThreadContext(context: CoroutineContext, oldState: S)
}

The need to have it for all targets (jvm, k/js, k/native) appeared in Compose Runtime library. Currently it's used in JVM only, but something similar needed to provide correct implementations of new Compose API for k/js and k/native .

Use Case
Here is a test from compose to showcase an example of desired behaviour:

/**
To have better intuition about this example, it might be useful to  look at Snapshot in Compose. 
Snapshot resembles a version control system with branches and read, write, merge operations. 
This system manages the states of different objects (Compose's state objects). 
Mutations made with one snapshot/branch do not interfere with other snapshots, 
but a Mutable Snapshot can be applied/commited to its parent/root.
Here is a good article: https://dev.to/zachklipp/introduction-to-the-compose-snapshot-system-19cn

*/

/* Example dependencies */

fun Snapshot.takeSnapshot(): Snapshot // "branches" a new Snapshot from a current Snapshot (from Global by default)

// Snapshot.current is a public static property. Its `getter` is accessing a ThreadLocal (Compose has an internal expect/actual ThreadLocal)
val current: Snapshot
    get() = threadSnapshot.get() ?: currentGlobalSnapshot.get() //threadSnapshot is ThreadLocal
    
// Snapshot.asContextElement() in an extension that converts a Snapshot to CoroutineContext.Element:
fun Snapshot.asContextElement(): SnapshotContextElement = SnapshotContextElementImpl(this) //  implements ThreadContextElement and sets/restores a values accessed via Snapshot.current

/* Example */
val snapshotOne = Snapshot.takeSnapshot() // branch from the Global snapshot
val snapshotTwo = Snapshot.takeSnapshot() // branch from the Global snapshot
runBlocking {
    val stopA = Job()
    val jobA = launch(snapshotOne.asContextElement()) { // associate  snapshotOne with this coroutine
        assertSame(snapshotOne, Snapshot.current, "expected snapshotOne, A")
        stopA.join() // suspend, so Snapshot.current will be restored
        assertSame(snapshotOne, Snapshot.current, "expected snapshotOne, B")
    }
    launch(snapshotTwo.asContextElement()) { // associate  snapshotTwo with this coroutine
        assertSame(snapshotTwo, Snapshot.current, "expected snapshotTwo, A")
        stopA.complete()
        jobA.join() // suspend, so Snapshot.current will be restored
        assertSame(snapshotTwo, Snapshot.current, "expected snapshotTwo, B")
    }
}

A coroutine can read/write a snapshot by accessing it via Snapshot.current (even if there is no SnapshotContextElement, Snapshot.current will return a snapshot from either a ThreadLocal or a global snapshot).
SnapshotContextElement lets us have independent snapshots in different coroutines within one thread, so a coroutine can have an associated Snapshot.

ThreadContextElement can be further helpful to perform Snapshot.apply (commit changes) in restoreThreadContext, when a coroutine suspends.

Not existing feature: (asked by Adam Powell)
Let's say we want to commit a Snapshot when coroutine with SnapshotContextElement suspends: so we call Snapshot.apply from ThreadContextElement.restoreThreadContext. The thing is apply can return SnapshotApplyResult.Failure.
In that case, it's desirable to resume a suspended coroutine right away with an exception. It would need to be an arbitrary exception and not a CancellationException, as it should eventually result in failing the job tree.

Question: Can we somehow have an access to a continuation of a suspended coroutine within a CoroutineContext.Element (right after it was suspended)?

Question to conclude the above:

• Can we have a common way to hook into the coroutine lifecycle? To have a kind of callback when a coroutine resumes and suspends? (like ThreadContextElement in jvm)
• If we had such hooks, can we optionally manage that coroutine within those hooks? Like resumeWith(error) immediately after suspension in restoreThreadContext ?

[adamp]

Happy to provide additional context on the request if it helps.

[chuckjaz]

A bit of background around Snapshots.

Snapshots and Compose

Snapshots solve two problems for Compose, observation and consistency. It does this through an in-memory transaction system implemented using an MVCC algorithm.

Consistency

When a snapshot is taken the values of state objects are isolated from changes made in any other snapshot. They retain the values they had at the start of the snapshot regardless of any changes made outside the snapshot. This allows Compose to take a snapshot at the beginning of composition and produce frame that is consistent with a instantaneous snapshot of state at the beginning of composition. For example, if a composition produces a column of numbers and a sum of the column at the bottom, the sum is guaranteed to be the sum of the values displayed even if the numbers shown are actively changing by some background thread. This is because the changes are not visible to composition as composition takes a snapshot before it starts.

Observation

As changes are made to state objects, the Snapshot system will provide the equivalent of a flow of the objects changed (using a more primitive callback notification). Observers can register to be notified whenever a snapshot commits with changes. When a snapshot is created a callback can be supplied that notified whenever a state object is read. These two features combine to form a publish/subscribe system that allows composition to automatically subscribe to be notified when state objects it reads are modified, even indirectly. This is how Compose knows when and what to re-execute to update the composition.

Global snapshot

Transactions are typically hard to use as all modifications must be in a transaction and and committing that transaction might fail as it might contain changes that collide with changes that have already been committed by another transaction.

Snapshots simplify this by allowing the user to trade consistency for availability by using the global snapshot. The global snapshot is the default snapshot a thread is in when the thread not explicitly in another snapshot. The global snapshot is guaranteed to commit and is committed before any new, top level snapshot is taken. This means that any changes made in the global snapshot are guaranteed to be accepted and visible to all new snapshots (existing snapshots are still isolated from the global snapshot). This allows most code to ignore snapshots, they just treat state objects as if they were normal mutable memory. If they need consistency, can then either use a synchronization mechanism of their choice (such as locks) or create a snapshot and make the changes in that, which might fail and require the work to be repeated.

The users of snapshots always have a choice, make consistent changes that can fail or make potentially inconsistent changes that will always succeed.

Snapshots and coroutines

Currently snapshots are of limited use in coroutines as they rely on a thread local variable to preserve the isolation for a thread. As coroutines hop threads the thread local variable is not carried with the coroutine. The .asContextElement() allows the bridging of the thread local to the coroutine allowing them to be used as well as they are currently with threads. Without something like .asContextElement() then coroutines are limited to using the global snapshot.

However, a context element is helpful but not sufficient if used with structured concurrency. For snapshots to work well with structured concurrency, as the coroutine splits and joins the snapshot should also split and join. If the coroutine is cancelled then the snapshot it is using should also be cancelled. Snapshots support nested snapshots. That is, in any snapshot a child snapshot can be taken and when it is applied (a.k.a. committed) it is applied to the parent snapshot. When used with structured concurrency, whenever a child context is created a corresponding child snapshot should be taken. When the child context joins the parent then the snapshot should be applied to the parent. If, when applying the child snapshot, it fails because of a conflicting change, the corresponding context should be cancelled with an exception.

The context element helps bring coroutines on level footing with threads when using snapshots. This is the motivation for asking for the equivalent of of a thread context element. Allowing snapshot state to track the structure of concurrency allows them to be truly useful. This is the motivation for additional hooks into coroutines.

[dkhalanskyjb]

Thanks for the detailed explanation!

Can we somehow have an access to a continuation of a suspended coroutine within a CoroutineContext.Element (right after it was suspended)?

I'm not sure this is what should happen if there's a merge conflict. What if we were going to resume the coroutine with an exception anyway? We'd have two exceptions to choose from in that case. The one with which we wanted to resume the coroutine anyway is probably expected and properly handled, whereas a merge conflict looks like a more severe issue.

I agree with you that it should at least propagate throughout the coroutine hierarchy, but let's expand on this: perhaps we could use the uncaught exception handler for this instead? If merge conflicts are considered purely a programming mistake (so, not a normal mode of operation), crashing an app by default could be a proper solution. If someone does need explicit handling of such exceptions, it's possible to override such behavior by registering an uncaught exception handler.

What do you think?

[eymar]

@dkhalanskyjb Thank you for reply! Sorry for a long absence.

We'd have two exceptions to choose from in that case. The one with which we wanted to resume the coroutine anyway is probably expected and properly handled, whereas a merge conflict looks like a more severe issue.

As an option, maybe overriding the Result.Failure can be restricted. If there was an exception before, we probably shouldn't even attempt to apply the snapshot.

If merge conflicts are considered purely a programming mistake (so, not a normal mode of operation), crashing an app by default could be a proper solution.

I'm not sure that a merge conflict when applying snapshots (SnapshotApplyResult.Failure) is considered a purely programming mistake.

Let's hear from Adam and Chuck on this. cc: @adamp @chuckjaz

---

P.S. Back then, I made a simple prototype of ThreadContextElement for k/js and k/native #3325, but of course it doesn't solve all extra questions.

[chuckjaz]

What if we were going to resume the coroutine with an exception anyway? We'd have two exceptions to choose from in that case.

Dispose (a.k.a abort) the snapshot. A snapshot does not need to be applied. If the coroutine is going to cancel with an exception then its snapshot should be disposed(). If the developer wants the different behavior they can, for example, commit the snapshot in a catch or finally and they can then decide how to handle conficts. The default behavior should be to discard the snapshot.

I'm not sure that a merge conflict when applying snapshots (SnapshotApplyResult.Failure) is considered a purely programming mistake.

It depends on the context of how the snapshot is being used.

For the states created in composition, they should only ever be modified during composition and, since is never multiple coroutines operating on the same composition, merge conflicts are hard to generate, and rely on give a reference to state created in composition to some outside observer, and are be considered fatal errors.

However, if the state is part of the application model then snapshot failures would be more common depending on how the mutators are controlled (unique mutators of state cannot conflict).

The default should be that conflicts should throw an exception (e.g. SnapshotApplyConflictException generated by check() of the apply result) when they are merged as the work of the lanched coroutine failed to produce a valid result. One can imagine a retryingLaunch { } which launches the lambda in a coroutine and, if applying the snapshot fails, takes a new snapshot launches the coroutine again. With similar such primitives, the user can control how snapshot failure is handled at the point of launching the nested coroutine.

What I would like to be able to do is have a context element observe the splits and joins of the coroutine scopes for which it is an element. That is, if the coroutine splits a new element should be able to be created for the new context and then, when it later rejoins the parent, the element should be able to merged back to the parent element. It would also be nice to know when the resources used by the element can be discarded. If I had such an API, snapshots would be much more useful. Without it (and with thread elements) it is still useful but requires a lot of hand-holding. This same mechanism would be generally useful for any pooled resource (a pool of database connections, for example) as resource aquired by a coroutine context can be returned as a natural part of the structure of the coroutine. Snapshot are, if you squint hard enough, a pooled resource and would benefit from any support for pooled resources.

[hfhbd]

Another use-case:
I want to implement a Kotlin Native database driver. Each database connection is stateful and needs to be executed on its own thread, this is the limitation/design of the native API. To create a thread pool, there already is newFixedThreadPoolContext(). But there is no option to connect these dispatcher workers to a database connection on Kotlin Native.

On JVM, you could use ThreadContextElement, like Exposed does, but there is no Thread/WorkerContextElement on Kotlin Native to (re-) store a database connection in the current context based on the underlaying executing worker.