Cancel only the inner Task when awaiting an async function?

Question

I am trying to extend a framework, where our users of the framework can use C# as a language with non-blocking functions (aka "coroutines"). My first design used yield return statements and an IEnumerator as function return value.

However, that is quite a bit error prone and awkward to use if you try to call another yielding function (and even more if the other function should return a value).

So I am toying with the idea of using async and await to provide the corountines. The resulting syntax would be so much more fun.

In our framework, I need to ensure that no two of these non-blocking scripts ever run in parallel, so I wrote my own SynchonizationContext to schedule all actions by myself. That works like a charm.

Now the more interesting stuff: Another core part of our design is, that users can register some kind of "check functions" that will not continue the "current" executing function if they are failing. The check functions should be executed every time the non-blocking function resumes.

So for example:

async Task DoSomething()
{
    var someObject = SomeService.Get(...);
    using (var check = new SanityCheckScope())
    {
        check.StopWhen(() => someObject.Lifetime < 0);

        ...

        await SomeOtherStuff();
        // when execution comes back in here, the StopWhen - condition above should be
        // evaluated and the Task should not be scheduled if its true.

        ...
    }
}

Whenever SomeOtherStuff or any asynchronous function called by it resumes, the registered condition should be checked and DoSomething should stop if any condition is true.

To implement this, my SynchonizationContext checks the registered functions (passed via CallContext.LogicalSetData) and just does not schedule the task if one returns true.

Now comes the tricky problem. Suppose the function SomeOtherStuff looks like this:

async Task SomeOtherStuff()
{
    using (var check = new SanityCheckScope())
    {
        // register my own check functions
        await Whatever();
    }
}

In the example SomeOtherStuff registers its own check function. If they are true after the await Whatever(), naturally only the SomeOtherStuff function should be stopped. (Lets assume that if SomeOtherStuff returns a Task<XXX>, then it would be OK for me if default(XXX) would be used as return value.)

How could I approach this? I can't get to the Action that comes after the await SomeOtherStuff(); in my original function. There seem to be no hook or callback at the begin of an await either (nor would it make sense anyway).

Another idea would be to throw an exception instead of "not scheduling" the task. Instead (or additional to) the using-block, I would write a try/catch. However, then I got the problem that if any check function of the original DoSomething fails after the inner await Whatever(); how would I stop DoSomething (together with SomeOtherStuff)?

So before I scratch all this and go back to my IEnumerator and yield return.. Anyone got an idea whether this can be done by the async / await framework?

Answer 1

My recommendation is to just use TPL Dataflow and be done with it. The more you abuse a language feature like async (or IEnumerable for that matter), the less maintainable your code is.

That said...

You do (sort of) have some hooks around await: custom awaitables. Stephen Toub has a blog on them and Jon Skeet covers some details in his eduasync series.

However, you cannot return a custom awaitable from an async method, so this approach would mean all your awaits would have to "opt-in" to the "check" behavior:

async Task DoSomething()
{
  var someObject = SomeService.Get(...);
  using (var check = new SanityCheckScope())
  {
    check.StopWhen(() => someObject.Lifetime < 0);
    await SomeOtherStuff().WithChecks();
  }
}

It's not clear to me what the semantics should be when a check fails. From your question, it sounds like once a check fails you just want that method to abort, so the check is more like a code contract than a monitor or condition variable. It also sounds like the scope should only be applied to that method and not methods called from that method.

You may be able to get something working by using a custom awaiter which encapsulates the returned Task and uses TaskCompletionSource<T> to expose a different Task. The custom awaiter would perform the check before executing the continuation of the async method.

Another thing to watch out for is how your "scope" works. Consider this situation:

async Task DoSomething()
{
  using (var check = new SanityCheckScope())
  {
    check.StopWhen(...);

    await SomeOtherStuff();

    await MoreOtherStuff();
  }
}

Or even more fun:

async Task DoSomething()
{
  using (var check = new SanityCheckScope())
  {
    check.StopWhen(...);

    await SomeOtherStuff();

    await Task.WhenAll(MoreOtherStuff(), MoreOtherStuff());
  }
}

I suspect your scope will have to make use of AsyncLocal<T> (described on my blog) in order to work properly.

Answer 2

Have you considered implementing your own SynchronizationContext? You could either implement your own (assuming there is no meaningful SynchronizationContext already installed) or you could perhaps install a wrapper around an existing context.

The await syntax will post the continuation to the SynchronizationContext (assuming the await is not configured not to do so) which would allow you to intercept it. Of course any asynchronous callback will be posted to the SynchronizationContext so it may not be straightforward to detect the exact moment the continuation is called, so this could cause the sanity check to happen for all callbacks posted to the context.

Your SanityCheckScope can ensure that a checking scope is registered in a nested fashion with your SynchronizationContext where, obviously, the IDisposable part of your class deregisters and reinstates the parent scope. This would only work if you don't have a situation where you can have multiple parallel child scopes.

The only way to stop execution is, as far as I can see, to throw some kind of exception.