Intercept the call to an async method using DynamicProxy

Question

Below is the code from the Intercept method on a custom type that implements IInterceptor of the Castle Dynamic Proxy library. This snippet is from an AOP based logging proof-of-concept console app that is posted here.

    public void Intercept(IInvocation invocation)
    {
        if (Log.IsDebugEnabled) Log.Debug(CreateInvocationLogString("Called", invocation));
        try
        {
            invocation.Proceed();
            if (Log.IsDebugEnabled)
                if (invocation.Method.ReturnType != typeof(void))
                    Log.Debug("Returning with: " + invocation.ReturnValue);
        }
        catch (Exception ex)
        {
            if (Log.IsErrorEnabled) Log.Error(CreateInvocationLogString("ERROR", invocation), ex);
            throw;
        }
    }

This is working as expected on regular method calls, but not when tried with async methods (using the async/await keywords from C# 5.0). And I believe, I understand the reasons behind this as well.

For the async/await to work, the compiler adds the functional body of the method into a state machine behind the scenes and the control will return to the caller, as soon as the first awaitable expression that cannot be completed synchronously, is encountered.

Also, we can interrogate the return type and figure out whether we are dealing with an async method like this:

            if (invocation.Method.ReturnType == typeof(Task) || 
                (invocation.Method.ReturnType.IsGenericType && 
                 invocation.Method.ReturnType.GetGenericTypeDefinition() == typeof(Task<>)))
                Log.Info("Asynchronous method found...");

This works for only those async methods that returns either Task or Task<> and not void but I am fine with that.

What changes have to made within the Intercept method so that the awaiter would return to there rather than the original caller?

Answer 1

Presumably the "problem" is that it's just logging that it's returning a task - and you want the value within that task?

Assuming that's the case, you still have to return the task to the caller, immediately - without waiting for it to complete. If you break that, you're fundamentally messing things up.

However, before you return the task to the caller, you should add a continuation (via Task.ContinueWith) which will log the result (or failure) when the task completes. That will still give the result information, but of course you'll be logging it potentially after some other logging. You may also want to log immediately before returning, leading to a log something like this:

Called FooAsync
Returned from FooAsync with a task
Task from FooAsync completed, with return value 5

The business of getting the result out of the task (if it completed successfully) would have to be done with reflection, which is a bit of a pain - or you could use dynamic typing. (Either way it will be a bit of a performance hit.)

Answer 2

Thanks to Jon's answer, this is what I ended up with:

public void Intercept(IInvocation invocation)
{
    if (Log.IsDebugEnabled) Log.Debug(CreateInvocationLogString("Called", invocation));
    try
    {
        invocation.Proceed();

        if (Log.IsDebugEnabled)
        {
            var returnType = invocation.Method.ReturnType;
            if (returnType != typeof(void))
            {
                var returnValue = invocation.ReturnValue;
                if (returnType == typeof(Task))
                {
                    Log.Debug("Returning with a task.");
                }
                else if (returnType.IsGenericType && returnType.GetGenericTypeDefinition() == typeof(Task<>))
                {
                    Log.Debug("Returning with a generic task.");
                    var task = (Task)returnValue;
                    task.ContinueWith((antecedent) =>
                                          {
                                              var taskDescriptor = CreateInvocationLogString("Task from", invocation);
                                              var result =
                                                  antecedent.GetType()
                                                            .GetProperty("Result")
                                                            .GetValue(antecedent, null);
                                              Log.Debug(taskDescriptor + " returning with: " + result);
                                          });
                }
                else
                {
                    Log.Debug("Returning with: " + returnValue);
                }
            }
        }
    }
    catch (Exception ex)
    {
        if (Log.IsErrorEnabled) Log.Error(CreateInvocationLogString("ERROR", invocation), ex);
        throw;
    }
}

Answer 3

Trying to clarify with a generic and clean solution for:

• Intercepting async methods adding custom code as a continuation task.

I think the best solution is to use the dynamic keyword to bypass the compiler type checking and resolve the difference between Task and Task<T> at run time:

public void Intercept(IInvocation invocation)
{
    invocation.Proceed();
    var method = invocation.MethodInvocationTarget;
    var isAsync = method.GetCustomAttribute(typeof(AsyncStateMachineAttribute)) != null;
    if (isAsync && typeof(Task).IsAssignableFrom(method.ReturnType))
    {
        invocation.ReturnValue = InterceptAsync((dynamic)invocation.ReturnValue);
    }
}

private static async Task InterceptAsync(Task task)
{
    await task.ConfigureAwait(false);
    // do the logging here, as continuation work for Task...
}

private static async Task<T> InterceptAsync<T>(Task<T> task)
{
    T result = await task.ConfigureAwait(false);
    // do the logging here, as continuation work for Task<T>...
    return result;
}

Answer 4

Below is my async interceptor adapter implementation that correctly handles async methods.

abstract class AsyncInterceptor : IInterceptor
{
    class TaskCompletionSourceMethodMarkerAttribute : Attribute
    {

    }

    private static readonly MethodInfo _taskCompletionSourceMethod = typeof(AsyncInterceptor)
        .GetMethods(BindingFlags.NonPublic | BindingFlags.Instance)
        .Single(x => x.GetCustomAttributes(typeof(TaskCompletionSourceMethodMarkerAttribute)).Any());


    protected virtual Task<Object> InterceptAsync(Object target, MethodBase method, object[] arguments, Func<Task<Object>> proceed)
    {
        return proceed();
    }

    protected virtual void Intercept(Object target, MethodBase method, object[] arguments, Action proceed)
    {
        proceed();
    }

    [TaskCompletionSourceMethodMarker]
    Task<TResult> TaskCompletionSource<TResult>(IInvocation invocation)
    {
        var tcs = new TaskCompletionSource<TResult>();

        var task = InterceptAsync(invocation.InvocationTarget, invocation.Method, invocation.Arguments, () =>
        {
            var task2 = (Task)invocation.Method.Invoke(invocation.InvocationTarget, invocation.Arguments);
            var tcs2 = new TaskCompletionSource<Object>();
            task2.ContinueWith(x =>
            {
                if (x.IsFaulted)
                {
                    tcs2.SetException(x.Exception);
                    return;
                }
                dynamic dynamicTask = task2;
                Object result = dynamicTask.Result;
                tcs2.SetResult(result);
            });
            return tcs2.Task;
        });

        task.ContinueWith(x =>
        {
            if (x.IsFaulted)
            {
                tcs.SetException(x.Exception);
                return;
            }

            tcs.SetResult((TResult)x.Result);
        });

        return tcs.Task;
    }
    void IInterceptor.Intercept(IInvocation invocation)
    {
        if (!typeof(Task).IsAssignableFrom(invocation.Method.ReturnType))
        {
            Intercept(invocation.InvocationTarget, invocation.Method, invocation.Arguments, invocation.Proceed);
            return;
        }
        var returnType = invocation.Method.ReturnType.IsGenericType ? invocation.Method.ReturnType.GetGenericArguments()[0] : typeof(object);
        var method = _taskCompletionSourceMethod.MakeGenericMethod(returnType);
        invocation.ReturnValue = method.Invoke(this, new object[] { invocation });
    }
}

and sample usage:

class TestInterceptor : AsyncInterceptor
{
    protected override async Task<Object> InterceptAsync(object target, MethodBase method, object[] arguments, Func<Task<object>> proceed)
    {
        await Task.Delay(5000);
        var result = await proceed();
        return DateTime.Now.Ticks % 2 == 0 ? 10000 :result;
    }
}

Answer 5

My 2 cents:

It has been correctly established that for async methods the purpose of the interceptor would be to "enhance" the task returned by the invocation, via a continuation.

Now, it is precisely this task continuation the one that has to be returned to make the job of the interceptor complete.

So, based on the above discussions and examples, this would work perfectly well for regular methods as well as "raw" async Task methods.

public virtual void Intercept(IInvocation invocation)
{
    try
    {
        invocation.Proceed();
        var task = invocation.ReturnValue as Task;
        if (task != null)
        {
            invocation.ReturnValue = task.ContinueWith(t => {
                if (t.IsFaulted)
                    OnException(invocation, t.Exception);
            });
        }
    }
    catch (Exception ex)
    {
        OnException(invocation, ex);
    }
}

public virtual void OnException(IInvocation invocation, Exception exception)
{
    ...
}

1. But when dealing with async Task<T> methods, the above would incorrectly change the type of the task returned by the interception, from Task<T> to regular Task

2. Notice that we are calling Task.ContinueWith() and not Task<TResult>.ContinueWith(), which is the method we want to call.

This would be the resulting exception when ultimately awaiting the such an interception:

System.InvalidCastException: Unable to cast object of type 'System.Threading.Tasks.ContinuationTaskFromTask' to type 'System.Threading.Tasks.Task`1

Answer 6

Having a need to intercept methods returning Task<TResult>, I've created an extension to Castle.Core that simplifies the process.

Castle.Core.AsyncInterceptor

The package is available to download on NuGet.

The solution is largely based on this answer from @silas-reinagel, but simplifies it by providing a new interface to implement IAsyncInterceptor. There are also further abstractions to that make interception similar to implementing Interceptor.

See the readme of the project for further details.

Answer 7

   void IInterceptor.Intercept(IInvocation invocation) {
       try {
           invocation.Proceed();
           var task = invocation.ReturnValue as Task;
           if (task != null && task.IsFaulted) throw task.Exception;
       }
       catch {
           throw;
       }
   }

Answer 8

Instead of:

tcs2.SetException(x.Exception);

You should use:

x.Exception.Handle(ex => { tcs2.SetException(ex); return true; });

to bubble up the real exception...