How to write interceptor for methods returning IEnumerable

Question

I wrote simple interceptor (with Castle.DynamicProxy) to handle database connection life cycle. I.e. all services have Connection property which opens new one on first use. Connection is being closed automatically after each method call:

class CloseConnectionInterceptor : IInterceptor
{
    public void Intercept(IInvocation invocation)
    {
        try
        {
            invocation.Proceed();
        }
        finally
        {
            var hasConnection = invocation.InvocationTarget as IHasConnection;
            if (hasConnection.IsConnectionOpen)
            {
                hasConnection.Connection.Dispose();
                hasConnection.Connection = null;
            }
        }
    }
}

interface IHasConnection
{
    bool IsConnectionOpen { get; }
    Connection Connection { get; set; }
}

It works well except for methods returning iterator:

[Test]
public void ConnectionClosedByProxy()
{
    // arrange
    var service = new MyService();
    var proxy = new ProxyGenerator()
        .CreateInterfaceProxyWithTarget<IMyService>(
            service, new CloseConnectionInterceptor());

    // act
    proxy.GetList();

    // assert: works!!!
    Assert.IsFalse(service.IsConnectionOpen);
}

[Test]
public void IteratorLeavesOpenConnection()
{
    // arrange
    var service = new MyService();
    var proxy = new ProxyGenerator()
        .CreateInterfaceProxyWithTarget<IMyService>(
            service, new CloseConnectionInterceptor());

    // act
    proxy.GetEnumerable().ToList();

    // assert: bad, bad, bad!
    Assert.IsTrue(service.IsConnectionOpen);
}

See full example here: https://gist.github.com/4087483

If there is "using (new Connection())" statement inside my GetEnumerable method then it works as expected - connection is being closed after last access to iterator. Is it possible to catch this moment in interceptor? Or should I proxy not only method but also resulting IEnumerable?

Answer 1

I put seed of the answer in my question :). Here is modified interceptor which proxies also resulting IEnumerable:

class CloseConnectionInterceptor : IInterceptor
{
    public void Intercept(IInvocation invocation)
    {
        Type genericType = invocation.Method.ReturnType.IsGenericType
            ? invocation.Method.ReturnType.GetGenericTypeDefinition()
            : null;

        if (genericType == typeof(IEnumerable<>))
        {
            invocation.Proceed();

            var method = GetType()
                .GetMethod("HandleIteratorInvocation")
                .MakeGenericMethod(
                    invocation.Method.ReturnType.GetGenericArguments()[0]);

            invocation.ReturnValue = method.Invoke(
                null,
                new[] { invocation.InvocationTarget, invocation.ReturnValue });
        }
        else
        {
            HandleNonIteratorInvocation(invocation);
        }
    }

    public static IEnumerable<T> HandleIteratorInvocation<T>(
        IHasConnection hasConnection, IEnumerable enumerable)
    {
        try
        {
            foreach (var element in enumerable)
                yield return (T)element;
        }
        finally
        {
            CloseOpenConnection(hasConnection);
        }
    }

    private static void HandleNonIteratorInvocation(IInvocation invocation)
    {
        try
        {
            invocation.Proceed();
        }
        finally
        {
            CloseOpenConnection(invocation.InvocationTarget as IHasConnection);
        }
    }

    private static void CloseOpenConnection(IHasConnection hasConnection)
    {
        if (hasConnection.IsConnectionOpen)
        {
            hasConnection.Connection.Dispose();
            hasConnection.Connection = null;
        }
    }
}