A general-purpose decorator - correctly awaiting an awaitable type obtained from Reflection

Question

Preamble: I was looking at how best to implement a general-purpose decorator in C# that copes with async methods. Most examples use DispatcherProxy, which only offers a synchronous Invoke. Unless NET improves this and adds async support, currently closed, we're left with calling this in a way that works the same way as a normal specific decorator. The problem seems to be that because of the decision to make awaitables be a pattern, there's no clean way to detect this circumstance and write the code correctly, but I'm hoping someone more familiar with the depths of Reflection can shine a light.

A specific easy example first:

interface IMyWork
{
  Task DoStuff();
}

class MyWork : IMyWork
{
  public async Task DoStuff() { await Something(); }
}

// A decorator like any other
class MyDecoratedWork : IMyWork
{
  private readonly MyWork _real;

  public Task DoStuff()
  {
    using var scope = Decorate();
    await _real.Something();
    // Disposal is in the same context as Decorate was - things like AsyncLocal are restored
  }

  private IDisposable Decorate() { /* return some scoped decoration /* }
}

That's easy if you want to write decorators manually, but if you've got lots of these to write, it becomes appealing to have a general-purpose decorator. Examples tend to look like this:

class DispatchProxyScopedDecorator<TDecorated> : DispatchProxy
{
  private TDecorated? _decorated;
  private IScopeProvider? _scopeProvider;

  private void SetParameters(TDecorated decorated, IScopeProviderscopeProvider)
  {
    _decorated = decorated;
    _scopeProvider = scopeProvider;
  }

  protected override object? Invoke(MethodInfo? targetMethod, object?[]? args)
  {
    // The important bit
  }

  public static TDecorated Create(TDecorated decorated, IScopeProvider scopeProvider)
  {
    object created = Create<TDecorated, DispatchProxyScopedDecorator<TDecorated>>()!;
    ((DispatchProxyScopedDecorator<TDecorated>)created).SetParameters(decorated, scopeProvider);

    return (TDecorated)created;
  }
}

The synchronous case is easy:

protected override object? Invoke(MethodInfo? targetMethod, object?[]? args)
{
  using var scope = _scopeProvider.Decorate();
  return targetMethod?.Invoke(_decorated, args);
}

but Invoke is a synchronous method, which means to correctly decorate a method call that returns an awaitable, the return type has to be intercepted with a new awaitable.

The only requirement appears to be implementing a method GetAwaiter (things like AsyncStateMachineAttribute don't necessarily apply).

var scope = _scopeProvider.Decorate();

var result = targetMethod.Invoke(_decorated, args);

if (result is null)
{
  return null;
}

// This isn't complete - also needs to hunt through assemblies for extension methods!
var awaiterMethod = result.GetType().GetMethod("GetAwaiter");

If this isn't null, I've got an awaitable. My ideal in terms of simple-to-maintain code would be to just drop this into an async method and let the compiler do the hard work:

async Task InterceptAsync(??? awaitable, MyScope scope)
{
  try
  {
    await awaitable;
  }
  finally
  {
    scope.Dispose();
  }
}

if (awaiterMethod is not null)
{
  return InterceptAsync(result, scope);
}

Unfortunately, this doesn't compile for obvious reasons - since an awaitable is not a single type, there's nothing I can supply to ??? that the compiler will accept (other than maybe dynamic, but that's deprecated).

There are obvious downsides as with any decorator (performance, potential change of return type), but is this even feasible? Am I missing some other approach? My immediate thought is that if the original code isn't returning a Task, the above would violate the expected method signature - but does that mean every single different awaitable requires its own implementation?

Thoughts appreciated.

Answer 1

If all you want to do is call Dispose before you execute the await continuation code in the code consuming your proxy , maybe you could just check for the mandatory INotifyCompletion and then register the scope.Dispose() there.

Something like this:

var scope = _scopeProvider.Decorate();

var result = targetMethod.Invoke(_decorated, args);

if (result is null)
{
  return null;
}

var awaiterMethod = result.GetType().GetMethod("GetAwaiter");

var awaiter = awaiterMethod.Invoke(result,null)
if(awaiter is INotifyCompletion notifycompl){
   notifycomp.OnCompleted(()=>scope.Dispose());
}

return result;

You would be returning the initial task-like object too, not a Task wrapped one.

This is the general idea:

// the awaitable we intercept
async Task AsyncMethod()
{
    await Task.Delay(1000);
    Console.WriteLine("Inside the action/work behind behind the awaitable awaitable");
    throw new Exception();
}

async Task Main()
{
    var task = AsyncMethod();
    // our decorator work
    task.GetAwaiter().OnCompleted(() => Console.WriteLine("In Proxy Code: Scope.Disposed()"));

    try
    {
        await task; // the actual use of our proxy
    }
    catch (Exception ex)
    {
        Console.WriteLine("In Real Code");
    }
}

Output:

Inside the action/work behind behind the awaitable awaitable
In Proxy Code: Scope.Disposed()
In Real Code

This wouldn't work if there is a single continuation action maintained internally by the custom awaitable like in this answer, but the wording of INotifyCompletion from the docs and in the comments in the source code:

Represents an operation that schedules continuations when it completes.

leads me to believe that most implementations should adhere to multiple continuations which will favor our approach here.