Conveying Context with AsyncLocal

Sometimes you want custom, ambient- or cross-cutting-context to flow across multiple components or multiple task continuations. Maybe it’s some ad-hoc TenantId, CorrelationId, vector-clock info, IsTest or IsShadow categorization, A/B test flight ids, remaining latency-budget (e.g. for cache or retry aggressiveness), quality-of-service level (e.g. for throttling perf-SLA’d requests vs. low-priority, re-runnable batch-reqs). Whatever.

There could be a lot of distance between where the context arrives and where its used, and you likely want to avoid so-called Tramp Data (passing data to one routine just so that routine can pass it to another routine), as that’s bad conceptual coupling and quite messy given more context ideas invariably emerge later. What you really want is a way to add the data similarly to how Thread.CurrentCulture or .CurrentPrincipal make UI and security cross-cutting data ubiquitously available. Those are ‘global’, yet per-request at the same time.

AsyncLocal

Happily, the AsyncLocal class helps us out a lot here.

Now, at first glance in the doc page, you’re going to see a static instance and property-setter, and immediately huff that there’s no-way that could be thread-safe. But it is. The property-setter is really acting as a static-method, updating the execution-context that .Net itself uses for flowing state into child tasks no matter which thread(s) they’re scheduled onto. CurrentPrincipal and CurrentCulture use the technique.

using System.Threading;

class Example
{
    static AsyncLocal<string> _asyncLocalString = new AsyncLocal<string>();

    static async Task AsyncMethodA()
    {
        _asyncLocalString.Value = req.TenantId; // 👈 Thread-safe? 

under the covers this executes:

public T Value
{
 ...
 set =>
     ExecutionContext.SetLocalValue(this, value, _valueChangedHandler != null);

Granted, adding hidden-inputs to functions is contrary to purity, but there are certain cases where marshalling some data through the backdoor (ideally to turn it into a passed parameter at a good spot) is the practical choice. Up to you and your team (and your priest or counselor).

Bare-Bones AmbientContext Utility

Now, we tried to motivate this exploration by saying we have a whole bundle of context-data that we’re likely to add to over time. Dedicated static-members of strongly-typed AsyncLocal<T>, akin to CultureInfo’s, is the prudent way to go for most apps, but should you be giving customers a library they can hang context on, a dict of AsyncLocals seems to work.

// Utility to set/get ambient context that flows across async tasks
public static class AmbientContext
{
    public static readonly ConcurrentDictionary<string, AsyncLocal<object?>> 
        _contexts = new(StringComparer.Ordinal);
    
    public static void Set<T>(string key, [MaybeNull] T val)
    {
        AsyncLocal<object?> keyctx = _contexts.AddOrUpdate( 
                key,
                k      => new AsyncLocal<object?>(),
                (k,al) => al);
        keyctx.Value = (object?) val;
    }
    
    [return: MaybeNull]
    public static T Get<T>(string key)
    {
       return _contexts.TryGetValue( key, out AsyncLocal<object?>? keyctx)
                ? (T) (keyctx!.Value ?? default(T)!)
                : default(T);
    }
}

Perhaps your DefaultController could call AmbientContext.Set( TenantIdKey, claims… ) and several layers deeper the data-access-layer calls AmbientContext.Get<string>( TenantIdKey ) to lookup this customer’s storage.

(Having to supply the type-parameter to Get<T> is a ptential source of bugs, so, again, dedicated fields are desired in most cases. A hybrid with fields and a dict might be advisable.)

Scoped State

If a method spawns multiple child tasks, those might try to add or override AmbientContexts (e.g. if an ML layer runs two models, live and IsShadow, you could have the ML engine skip billing for the latter context). Changes made by a child percolate to grand-child tasks, but don’t bleed-over into sibling child tasks, nor back into the parent task.

(Yes, in the naive AmbientContext class above, the ConcurrentDictionary still holds the memory, but the ExceutionContext does not. Pruning keys from this implementation of AmbientContext is inherently dangerous.)

But what if I want to mutate the context? Well, it’s changes to AsyncLocal.Value that have scope isolation. Getting the same Value to a mutable reference object would let you change it.

RunStats reqStats = new RunStats { LatencyBudget = 500L, ... };
AmbientContext.Set( "RunStats", reqStats );
await DoStuff();
if (reqStats.LatencyBudget < 0L) { throw new TimeoutEx("slow-poke"); }

...

Task<X> DoStuff()
{
    RunStats stats = AmbientContext.Get<RunStats>("RunStats");
    long msecStart = Environment.TickCount64;
    ...
    stats.LatencyBudget -= (Environment.TickCount64 - msecStart);  // interlocked ideally

Unit-Testing

We want our unit-tests to be able to run in parallel. Luckily, this is no different from our web-server running http-requests in parallel: task-schedulers will draw from thread-pools, flowing ExecutionContext either way. Remember to have each UT assign to all AsyncLocal.Value properties they intend to inspect (or if mutable modify) to preserve the isolation.

Context Across RPCs / Persistence

No doubt you’d want an ability to flow context even over the wire to other micro- or external-services (perhaps on a per-endpoint basis, e.g. CorrelationId to everyone, sometimes as http-header, sometimes as a gRPC field, but QoS level only to some endpoints). We won’t delve deeply into this, but one approach would web-apis might be to use ASP.Net’s middleware and client-factory middleware support.

using Microsoft.AspNetCore.Builder;
IApplicationBuilder builder = ... 
builder.UseMiddleware<AmbientContextMiddleware>();

...

public class AmbientContextMiddleware
{
    private readonly RequestDelegate _next;

    public MyMiddleware(RequestDelegate next) => 
        _next = next;

    public async Task Invoke(HttpContext context)
    {
        string corrId = 
            context.Request
                   .Headers["x-foocorp-correlationId"]
                   .FirstOrDefault( Guid.NewGuid().ToString() );
        
        AmbientContext.Set( CorrelationIdKey, corrId );
        
        await _next.Invoke(context);

        // TODO: emit corrid response header, esp if _we_ created it
    }
}

public class OutboundHttpAmbientContext 
    : System.Net.Http.DelegatingHandler
{
    protected override Task<HttpResponseMessage> SendAsync(
        HttpRequestMessage request,
        CancellationToken  cancellationToken)
    {
        string corrId = AmbientContext.Get<string>(CorrelationIdKey);
        if (corrId is not null)
        {
            request.Headers["x-foocorp-correlationId"]
                   .Append( corrId );
        }

        return base.SendAsync( request, cancellationToken );
}

Something similar would be desired for persistence to, say, your Data Science training pipeline to capture the TenantId, flight-id(s), etc. from context for them to condition their models on.

Conclusion

We’ve explored how AsyncLocal can help you convey custom execution-context that can keep your code clean from tramp-data. I hope you’ve found it an interesting technique. As always, leave a comment about your approach or experience with this one.