Eduasync part 15: implementing COMEFROM with a horrible hack

skeet — Fri, 11 Nov 2011 22:15:11 GMT

Ages ago when I wrote my previous Eduasync post, I said we'd look at a pipeline model of coroutines. I've decided to skip that, as I do want to cover the topic of this post, and I've got some more "normal" async ideas to write about too. If you want to look at the pipeline coroutines code, it's project 20 in the source repository. Have fun, and don't blame me if you get confused reading it - so do I.

The code I am going to write about is horrible too. It's almost as tricky to understand, and it does far nastier things. Things that the C# 5 specification explicitly says you shouldn't do.

If it makes you feel any better when your head hurts reading this code, spare a thought for me - I haven't looked at it in over six months, and I don't have a blog post explaining how it's meant to work. I just have almost entirely uncommented code which is designed to be hard to understand (in terms of the main program flow).

On no account should any code like this ever be used for anything remotely serious.

With that health warning out of the way, let's have a look at it...

COMEFROM at the caller level

The idea is to implement the COMEFROM control structure, which is sort of the opposite of GOTO (or in my implementation, more of a GOSUB). There are two operations, effectively:

ComeFrom(label): Register interest in a particular label.
Label(label): If anyone has registered interested in the given label, keep going from their registration point (which will be within a method), then continue from where we left off afterwards.

In some senses it's a little like the observer pattern, with labels taking the place of events. However, it looks entirely different and is much harder to get your head round, because instead of having a nicely-encapsulated action which is subscribed to an event, we just have a ComeFrom call which lets us jump back into a method somewhat arbitrarily.

I have two implementations, in project 22 and project 23 in source control. Project 22 is almost sane; a little funky, but not too bad. Project 23 is where the fun really happens. In addition to the operations listed above, there's an Execute operation which is sort of an implementation detail - it allows an async method containing ComeFrom calls to be executed without returning earlier than we might want.

Let's look at some code and the output, and try to work out what's going on.

internal class Program
{
    private static void Main(string[] args)
    {
        Coordinator coordinator = new Coordinator(SimpleEntryPoint);
        coordinator.Start();
    }

    private static async void SimpleEntryPoint(Coordinator coordinator)
    {
        await coordinator.Execute(SimpleOtherMethod);

        Console.WriteLine("First call to Label(x)");
        await coordinator.Label("x");

        Console.WriteLine("Second call to Label(x)");
        await coordinator.Label("x");

        Console.WriteLine("Registering interesting in y");
        bool firstTime = true;
        await coordinator.ComeFrom("y");

        Console.WriteLine("After ComeFrom(y). FirstTime={0}", firstTime);

        if (firstTime)
        {
            firstTime = false;
            await coordinator.Label("y");
        }
        Console.WriteLine("Finished");
    }

    private static async void SimpleOtherMethod(Coordinator coordinator)
    {
        Console.WriteLine("Start of SimpleOtherMethod");

        int count = 0;
        await coordinator.ComeFrom("x");

        Console.WriteLine("After ComeFrom x in SimpleOtherMethod. count={0}. Returning.",
                          count);
        count++;
    }
}

The reason for the "Simple" prefix on the method names is that there's another example in the same file, with a more complex control flow.

Here's the output - then we can look at why we're getting it...

Start of SimpleOtherMethod
After ComeFrom x in SimpleOtherMethod. count=0. Returning.
First call to Label(x)
After ComeFrom x in SimpleOtherMethod. count=1. Returning.
Second call to Label(x)
After ComeFrom x in SimpleOtherMethod. count=2. Returning.
Registering interesting in y
After ComeFrom(y). FirstTime=True
After ComeFrom(y). FirstTime=False
Finished

So, the control flow is a bit like this:

Start SimpleEntryPoint
- Call into SimpleOtherMethod
  - Log "Start of SimpleOtherMethod"
  - Initialize the "count" variable with value 0
  - Register interest in x; ComeFrom remembers the continuation but keeps going.
  - Log "After ComeFrom x in SimpleOtherMethod. count=0. Returning."
  - Increment count to 1.
  - Return.
- Return takes us back to SimpleEntryPoint...
Log "First call to Label(x)"
Call Label("x")...
- ... which takes us back into SimpleOtherMethod (remember, the method we thought we'd finished executing?) just after ComeFrom
  - Log AfterComeFrom x in SimpleOtherMethod. count=1. Returning.
  - Increment count to 2.
  - Return.
- Return takes us back to SimpleEntryPoint...
Log "Second call to Label(x)"
Call Label("x")...
- ... which takes us back into SimpleOtherMethod again
  - Log AfterComeFrom x in SimpleOtherMethod. count=2. Returning.
  - Increment count to 3.
  - Return.
- Return takes us back to SimpleEntryPoint...
Log "Registering interest in y"
Initialize the "firstTime" variable with value true.
Register interest in y; ComeFrom remembers the continuation and keeps going
Log "After ComeFrom(y). FirstTime=True"
Check the value of firstTime... It's true, so:
- Set firstTime to false
- Call Label("y")
... which takes us back to earlier in the method (just after ComeFrom), like a normal looping construct...
Log "After ComeFrom(y). FirstTime=False"
Check the value of firstTime... It's false, so:
- Log "Finished"
- Exit!

Doing all of this has a few interesting challenges. Let's look at them one at a time... and I would strongly advise you not to try to pay too much attention to the details.

Noting a continuation and continuing regardless...

Just as a quick reminder before we get cracking, it's worth remembering that all of this is entirely synchronous, despite being implemented with async. There's only a single user thread involved here. As with previous parts, we maintain a stack of actions to call, and basically keep calling from the top until we're done - but the actions we call can create extra stack entries, of course.

ComeFrom has unusual semantics in terms of async. We want to remember the continuation and keep executing as if we didn't need to wait. We can easily do one side or the other. If we wanted to just keep going without needing to know about the continuation, we could just return true from IsCompleted. If we just want to remember the continuation, we can make the awaiter's IsCompleted property return false, and remember the continuation when it's passed to OnCompleted. How do we do both?

Well, effectively we want to remember the continuation and then call it immediately. But we can't just call it directly from OnCompleted, as otherwise each ComeFrom call would end up in a "real" execution stack from, whereas our execution stack is stored as a Stack<Action>. So instead, we need to remember the continuation and immediately put it at the top of the stack.

However, that only works if as soon as the generated code returns from the async method containing the ComeFrom call, we go back into the state machine. If we'd just called SimpleOtherMethod directly in SimpleEntryPoint, we would have continued within SimpleEntryPoint with the new stack entry just waiting around. This is why we need the Executor method: that does exactly the same thing, effectively shuffling the stack around. When it's given something to execute, it puts its own continuation on the action stack, then the action it's been asked to execute, then returns. The top level code will then pick up the original action, and we're away.

So, here's the code for Execute, which is the simplest part of the coordinator:

public ExecuteAwaiter Execute(Action<Coordinator> action)
{
    return new ExecuteAwaiter(() => action(this), this);
}

public class ExecuteAwaiter
{
    private readonly Action action;
    private readonly Coordinator coordinator;

    internal ExecuteAwaiter(Action action, Coordinator coordinator)
    {
        this.action = action;
        this.coordinator = coordinator;
    }

    public ExecuteAwaiter GetAwaiter()
    {
        return this;
    }

    // Always yield
    public bool IsCompleted { get { return false; } }

    public void OnCompleted(Action callerContinuation)
    {
        // We want to execute the action continuation, then get back here,
        // allowing any extra continuations put on the stack *within* the action
        // to be executed.
        coordinator.stack.Push(callerContinuation);
        coordinator.stack.Push(action);
    }

    public void GetResult()
    {
    }
}

All the awaitables in this project return themselves as the awaiter - when you don't need any other state, it's an easy step to take.

That's all we need to say about Execute, but how exactly are we capturing the continuation in ComeFrom?

Capturing continuations

Once we've got the action stack shuffling under our belts, there are two more problems with ComeFrom:

What happens if we ComeFrom the same label twice?
How do we really capture a continuation?

The first point didn't come up in the sample I've shown here, but it does come up in the more complex example - imagine if SimpleOtherMethod had two ComeFrom calls; when we jump back to the first one, we'll execute the second one again. I made a simple policy decision to only allow a single "return point" for any label - if a ComeFrom call tries to register the existing continuation point for a label, we ignore it; otherwise we throw an exception. So we only need to care about a single continuation for any label, which makes life easier.

The second point is trickier. If you remember back to earlier posts in this series, we saw that the state machine generated for async only really contains a single entry point (MoveNext) which is used for all continuations. A variable in the state machine is responsible for remembering where we were within it between calls. So in order to really make the continuation remember the point at which it needs to continue, we need to remember that state. We need to store an object for the continuation, which contains the delegate to invoke, and the state of the state machine when we were first passed the continuation. I've created a class for this, unimaginatively called Continuation, which looks like this:

/// <summary>
/// This hack allows a continuation to be executed more than once,
/// contrary to the C# spec. It does this using reflection to store the
/// value of the "state" field within the generated class. NEVER, EVER, EVER
/// try to use this in real code. It's purely for fun.
/// </summary>
internal sealed class Continuation : IEquatable<Continuation>
{
    private readonly int savedState;
    private readonly object target;
    private readonly FieldInfo field;
    private readonly Action action;

    internal Continuation(Action action)
    {
        target = action.Target;
        field = target.GetType().GetField("<>1__state", BindingFlags.Instance | BindingFlags.NonPublic);
        savedState = (int) field.GetValue(target);
        this.action = action;
    }

    internal void Execute()
    {
        field.SetValue(target, savedState);
        action();
    }

    // Snip Equals/GetHashCode
}

Yes, we use reflection to fish out the <>1__state variable initially, and poke the state machine with the same value when we next want to execute the continuation. All highly implementation-specific, of course.

Now the ComeFrom method is reasonably straightforward - all we need is a dictionary mapping labels to continuations. Oh, and the same action stack shuffling as for Execute:

// In the coordinator
private readonly Dictionary<string, Continuation> labels = new Dictionary<string, Continuation>();

public ComeFromAwaiter ComeFrom(string label)
{
    return new ComeFromAwaiter(label, this);
}

public struct ComeFromAwaiter
{
    private readonly string label;
    private readonly Coordinator coordinator;

    internal ComeFromAwaiter(string label, Coordinator coordinator)
    {
        this.label = label;
        this.coordinator = coordinator;
    }

    public ComeFromAwaiter GetAwaiter()
    {
        return this;
    }

    // We *always* want to be given the continuation
    public bool IsCompleted { get { return false; } }

    public void OnCompleted(Action action)
    {
        Continuation newContinuation = new Continuation(action);
        Continuation oldContinuation;
        if (!coordinator.labels.TryGetValue(label, out oldContinuation))
        {
            // First time coming from this label. Always succeeds.
            coordinator.labels[label] = newContinuation;
        }
        else
        {
            // Current semantics are to prohibit two different ComeFrom calls for the same label.
            // An alternative would be to just replace the existing continuation with the new one,
            // in which case we wouldn't need any of this - we could just use
            // coordinator.labels[label] = newContinuation;
            // unconditionally.
            if (!oldContinuation.Equals(newContinuation))
            {
                throw new InvalidOperationException("Additional continuation detected for label " + label);
            }
            // Okay, we've seen this one before. Nothing to see here, move on.
        }
        // We actually want to continue from where we were: we're only really marking the
        // ComeFrom point.
        coordinator.stack.Push(action);
    }

    public void GetResult()
    {
    }
}

There's one interesting point here which is somewhat subtle, and screwed me up for a bit...

The default value of a struct is always valid...

You may have noticed that ComeFromAwaiter is a struct. That's pretty unusual for me. However, it's also absolutely critical. Without it, we'd get a NullReferenceException when we execute the continuation the second time.

Normally, the flow of async methods looks a bit like this, for an await expression taking the "long" route (i.e. IsCompleted is false):

Call GetAwaiter() and assign the result to an awaiter field
Call IsCompleted (which returns false in this scenario)
Set the state variable to remember where we'd got to
Call OnCompleted
Return
... When we continue...
Set state to 0 (running)
Call GetResult() on the awaiter
Set the awaiter field to default(TypeOfAwaiter)
Continue

Now that's fine when we're only continuing once - but if we need to jump into the middle of that sequence a second time, we're going to call GetAwaiter() on the awaiter field after it's been set to the default value of the awaiter type. If the default value is null, we'll go bang. So we must use a struct.

Fortunately, our GetResult() call doesn't need any of the state in the awaiter - it's purely there to satisfy the normal flow of things. So we're quite happy with a "default" ComeFrom awaiter.

Finally, labels...

We've now done all the hard work. The final piece of the puzzle is Label, which just needs to check whether there's a continuation to jump to, and shuffle the action stack in the way we're now painfully accustomed to:

public LabelAwaiter Label(string label)
{
    Continuation continuation;
    labels.TryGetValue(label, out continuation);
    return new LabelAwaiter(continuation, this);
}

public class LabelAwaiter
{
    private readonly Continuation continuation;
    private readonly Coordinator coordinator;

    internal LabelAwaiter(Continuation continuation, Coordinator coordinator)
    {
        this.continuation = continuation;
        this.coordinator = coordinator;
    }

    public LabelAwaiter GetAwaiter()
    {
        return this;
    }

    // If there's no continuation to execute, just breeze through.
    public bool IsCompleted { get { return continuation == null; } }

    public void OnCompleted(Action action)
    {
        // We want to execute the ComeFrom continuation, then get back here.
        coordinator.stack.Push(action);
        coordinator.stack.Push(continuation.Execute);
    }

    public void GetResult()
    {
    }
}

Almost painfully simple, really.

So that looks like all the code that's used, right? Not quite...

Reusable builders?

As we saw in the sample code, we can end up finishing the same async method multiple times (SimpleOtherMethod completes three times). That's going to call SetResult on the AsyncVoidMethodBuilder three times... which feels like it should go bang. Indeed, when I revisited my code earlier I wondered why it didn't go bang - it's the sort of illegal state transition the framework is usually pretty good at picking up on.

Then I remembered - this isn't the framework's AsyncVoidMethodBuilder - it's mine. And my SetResult method in this project does absolutely nothing. How convenient!

Make it stop, make it stop!

Okay thiat was a pretty quick tour of some horrible code. You'll never have to do anything like this with async in sane code, but it certainly made me painfully familiar with how it all worked. Just to recap on the oddities involved:

We needed to capture a continuation and then immediately keep going, almost as if the awaiter had said the awaitable had completed already. This involved shenanigans with the execution model and an extra method (Execute)
We needed to remember the state of a continuation, which we did with reflection.
We needed to make awaiter.GetResult() a valid call after awaiter had been reset to the default value for the type
We needed to ensure that the builder created in the skeleton method could have SetResult called on it multiple times

That's all on continuations and co-routines, I promise.

Next time (hopefully soon) I'll look at an example of how composition works so neatly in async, and then show how we can unit test async methods - at least sometimes.

Jon Skeet: Coding Blog : Evil Code, C# 5