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

Answer 1

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 2

Instead of:

tcs2.SetException(x.Exception);

You should use:

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

to bubble up the real exception...

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

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 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

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.)