Adapted StructuredTaskScope for Context and cancellation

[Wesley Hill]

Hi, and thanks for your work on StructuredTaskScope. I recently adapted it
for use in our services to handle (a) propagation of the OpenTelemetry
Context, and (b) the richer cancellation model we use. Overall it went
fairly well, and I'll give a few details and some minor feedback in case
that's of interest.

Background: We have an implementation of io.opentelemetry.context.Context
for the usual purposes of passing around tracing, logging, authentication,
and other context. We also added a cancellation signal in that Context
implementation, very much along the lines of
grpc-java's Context/CancellableContext classes (which also makes it similar
to golang's Context). Cancelling a parent Context cancels all the child
contexts, and a Context can cancel itself when it reaches a Deadline, and
callbacks can be registered to run upon cancellation. We do this
because thread interruption alone is not a sufficient approach for
cancellation when we need to cancel things like HTTP or gRPC calls, JDBC
Statements, and other I/O which isn't interruptible. (I understand that
it's not a goal of the current structured concurrency JEPs to improve
cancellation beyond thread interruption).

Details:

StructuredTaskScope is a sealed interface with a hidden implementation, so
this required that I make a wrapper class which I called
CancellableTaskScope. I gave it exactly the same API though, so that it
could be a drop-in replacement for StructuredTaskScope. It uses the same
Joiners and Subtask types, and internally it delegates to a
StructuredTaskScope.

In order to propagate the current Context to all forked tasks, I could not
rely on the ScopedValue propagation of StructuredTaskScope, since
ScopedValue's model is not compatible with
io.opentelemetry.context.Context's approach of a makeCurrent() method which
returns a Scope to close, meaning that Context still has to be stored in a
ThreadLocal. So, I wrap the ThreadFactory for the task scope in a decorator
which first makes the scope's Context current on all new threads.

This is where the main challenge arose: The
StructuredTaskScope.Configuration interface has only withers, no accessors.
So my CancellableTaskScope could not work with that same Configuration
interface, as there was no way to get the configured ThreadFactory out of
it so that it could be wrapped before delegating to StructuredTaskScope. I
worked around this by duplicating the Configuration type in my class.
That's not terrible, and I ended up adding an additional configuration
parameter anyway for specifying an explicit parent Context.

Similarly, I take the timeout configured in the Configuration to set the
Deadline on our Context, and I don't pass that timeout through to the
StructuredTaskScope since our Context can manage it.

Propagating the Context to the forked tasks meant that our cancellation
signal was also propagated to them. The fork and join methods in
CancellableTaskScope are written to respect our Context cancellation. Next,
I needed to cancel the scope's Context when the Joiner decided to cancel
the StructuredTaskScope. This meant another decorator, this time wrapping
the Joiner to spy on the boolean returned by onFork/onComplete.

There is one minor mismatch, in that our Context can be cancelled
asynchronously at any time (such as a gRPC client cancelling a call), but
the StructuredTaskScope can only be cancelled at specific times: by the
owner thread calling close(), or by the Joiner when a task happens to be
forked or completed. I don't think this is an issue in practice though,
since the owner thread should respect an asynchronous cancellation signal
and quickly proceed to calling close().

So StructuredTaskScope turned out to be reasonably well extended by
composition even though a lot of it is sealed. I had some thoughts on the
exceptions StructuredTaskScope throws, which I will leave to another post.

Regards,
Wesley Hill
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <https://mail.openjdk.org/pipermail/loom-dev/attachments/20251024/dceb6d8c/attachment.htm>