Using Java’s Proxy Class to Transparently Manage Transactions

After we got a rough idea how Spring works internally, let’s go a bit further. Spring allows it to transparently open and close Database connections, add caching mechanisms and so on. Let’s find out how Spring does this.

There are many different approaches to create such features. We will examine Java’s java.lang.reflect.Proxy class here – it is rather simple, but not as powerful as the big alternative: bytecode manipulation. We’ll talk about the pros and cons in a few minutes.

Again, i made a small but working Java project on github so that you can experiment with a working source code.

The Scenario

Let’s assume we have a Database wrapper. We use an interface for this (it makes testing – especially mocking – easier):

public interface Database {
    void beginTransaction();
    void commitTransaction();
    void rollbackTransaction();
}

public class ClientClass implements ClientInterface1, ClientInterface2 {
    @Override
    public void runWithException() {
        throw new TestException();
    }
     
    @Override
    public int runWithoutException(final int result) {
        return result;
    }
}
     
public interface ClientInterface1 {
    int runWithoutException(int result);
}
     
public interface ClientInterface2 {
    void runWithException();
}

1. Starts beginTransaction
2. Calls the original Method
3. Checks if an exception was thrown
4. If there was an exception: runs rollbackTransaction
5. Otherwise: runs commitTransaction

It is easy to implement such a class for ClientClass, but we want a generic way which works with all Java objects.

The Test Setup

Our integration test has to fake the database, so we use JMock here. Its setup is simple:

public class ProxyTest {
    @Rule
    public JUnitRuleMockery context = new JUnitRuleMockery();
    private Database database;
    private TransactionProxy proxy;
     
    private void expectDatabaseCalls(final boolean expectCommit) {
        final Sequence sequence = context.sequence("dbsequence");
     
        context.checking(new Expectations() {{
            oneOf(database).beginTransaction();
            inSequence(sequence);
     
            if (expectCommit) {
                    oneOf(database).commitTransaction();
            } else {
                    oneOf(database).rollbackTransaction();
            }
            inSequence(sequence);
        }});
    }

TransactionProxy is the last key part. There is some boilerplate-Code, but it is not complicated: it expects that the beginTransaction method of the database mock is called. After that, either the commitTransaction or the rollbackTransaction must be called (depending on the parameter). The mock is created in the setup part of the test:

    @Before
    public void setUp() {
        database = context.mock(Database.class);
        proxy = new TransactionProxy(database);
    }

We now have everything we need to define our expectations:

    @Test
    public void successfulOperationsAreCommited() {
        expectDatabaseCalls(true);
        final ClientInterface1 interface1 = (ClientInterface1) proxy
                .createProxy(new ClientClass());
 
        final int expected = 10;
        final int actual = interface1.runWithoutException(expected);
        assertThat(actual, is(expected));
    }
 
    @Test(expected = TestException.class)
    public void unsuccessfulOperationsAreRolledBack() {
        expectDatabaseCalls(false);
        final ClientInterface2 interface1 = (ClientInterface2) proxy
                .createProxy(new ClientClass());
        interface1.runWithException();
    }
}

Our wrapper receives a new ClientClass instance. The wrapper is than casted to one of the interfaces (we will talk about this soon). Depending on the called method, we expect a different result and different invocations on the database.

How can such a universal wrapper be build? It must not know anything about the ClientClass! This is where Javas Proxy class kicks in.

The Wrapper

Java knows a Proxy class which is able to create some kind of “virtual object” for interfaces. Let’s look at the method signature:

Object java.lang.reflect.Proxy.newProxyInstance(ClassLoader loader,
        Class<?>[] interfaces, InvocationHandler h)
        throws IllegalArgumentException

 The first argument is a class loader. It will be used to define a new class on the fly. The second argument is a list of interfaces, which the newly created class will implement. newProxyInstance creates a class which implements all the provided interfaces and returns a new object of this class. But what happens if we invoke a method on this object? The third parameter – an InvocationHandler – takes care of this. We’ll see it in action very soon. This is the first part of our proxy class:

public class TransactionProxy {
    private final Database database;
     
    public TransactionProxy(final Database database) {
        this.database = database;
    }
     
    public Object createProxy(final Object object) {
        final Class<?> clazz = object.getClass();
        final Class<?>[] interfaces = clazz.getInterfaces();
        final ClassLoader classLoader = clazz.getClassLoader();
     
        final Object proxy = Proxy.newProxyInstance(classLoader, interfaces,
                createInvocationHandler(object));
        return proxy;
    }

The createProxy method wraps the object into our proxy object. It uses the same ClassLoader that loaded the wrapped object. It also collects all interfaces which are implemented by the object. Finally it creates a new proxy object which implements the same interfaces. And this is how the invocations are handled:

    private InvocationHandler createInvocationHandler(final Object target) {
        return new InvocationHandler() {
            @Override
            public Object invoke(final Object proxy, final Method method, final Object[] arguments)
                    throws Throwable {
     
                database.beginTransaction();
     
                final Object result;
                try {
                    result = method.invoke(target, arguments);
                } catch (final InvocationTargetException e) {
                    database.rollbackTransaction();
                    throw e.getCause();
                }
     
                database.commitTransaction();
                return result;
            }
        };
    }
}

The invoke method is called for every call to one of the interfaces’ methods. The InvokationHandler receives the proxy object (which we don’t need), the called method of the interface and the arguments. It must return a value which can be castet to the return type of the implemented method. Yep, we’re talking about reflections again . Our wrapped object (target) implements the same interfaces. This is why we can call method.invoke on the target object – we just forward the call. The result is just forwarded too. Java takes care to cast the result into the right type and throws an exception if it is not possible – but this is not important for us since we don’t modify the result.

What happens if an error occurs? Since we are using reflections, the thrown Exception is wrapped in an InvocationTargetException. Its getCause method contains the real exception object. A client will not expect a InvocationTargetException so we unwrap the original one after we rolled back the transaction.

Well – that’s it.

About the Interfaces

We used Interfaces all the time and there is a very good reason for this: the Proxy class of Java only supports interfaces. If the wrapper is initialized with a real class’ Class object, we will face an IllegalArgumentException. This forces the client to define every method in an interface. At least we can use multiple interfaces at once. But there is another drawback: the resulting object implements all interfaces of the client class, but it does not extend the client class itself. So the following method signature will only work if T is an interface:

public <T> T createProxy(final T object)

The following call works:

final ClientInterface2 interface1 = proxy
         .createProxy(new ClientClass());

 The following test does not:

@Test(expected = ClassCastException.class)
public void proxyCanNotBeCasted() {
    @SuppressWarnings("unused")
    final ClientClass error = proxy
            .createProxy(new ClientClass());
}

If this is a problem there are other solutions but they rely on bytecode manipulation. The way we explored here is a 100% pure Java solution.

I invite you to play around with this project. Try other clients, implement another wrapper. Try to understand what the wrapper does and when it is called. Look at stack traces to see that the wrapper method is one of the stack frames. Note that there is some overhead as our method is not the only one in the trace.

If you fully understand what happens here we can dive into JMock. It uses two proxys for the same interface to achieve a really, really cool syntax. I will blog about it soon (I hope ).