Who’s calling you?

I’ve been working on some performance testing profiling and in trying to diagnose a fix, we found one particular method was being called constantly. Searching for usages (static analysis) tells me who could possibly call the method we were inspecting, but we were interested in the runtime invocations. Since this was quite deep in the code, I used the power of dynamic proxies to do this.

I’ve rebuilt the code here:

Source:

package com.thekua.examples;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;

public class TraceBackProxy implements InvocationHandler {
    public static interface CallingMethodListener {
        void notify(String method);
    }

    private final Object wrapped;
    private final CallingMethodListener listener;

    private TraceBackProxy(Object wrapped, CallingMethodListener context) {
        this.wrapped = wrapped;
        this.listener = context;
    }

    public static Object wrap(Object target, CallingMethodListener context) {
        Class targetClass = target.getClass();
        return Proxy.newProxyInstance(targetClass.getClassLoader(),
                targetClass.getInterfaces(), new TraceBackProxy(target, context));
    }

    @Override
    public Object invoke(Object proxy, Method method, Object[] arguments) throws Throwable {
        String callingMethod = findCallingMethod(method);
        listener.notify(callingMethod);
        return method.invoke(wrapped, arguments);
    }

    private String findCallingMethod(Method method) {
        try {
            throw new RuntimeException();
        } catch(RuntimeException e) {
            StackTraceElement[] elements = e.getStackTrace();
            int callingMethodIndex = findIndexOfMethod(elements, method) + 1; // caller is next one down in stack
            return elements[callingMethodIndex].getMethodName();
        }
    }

    private int findIndexOfMethod(StackTraceElement[] elements, Method method) {
        for (int i = 0; i < elements.length; i++) {
            StackTraceElement current = elements[i];
            // does not cope with overloaded or duplicate method names
            if (current.getMethodName().equals(method.getName())) {
                return i;
            }
        }
        throw new IllegalStateException("Something went wrong and couldn't find method in stacktrace");
    }
}

Test:

package com.thekua.examples;

import org.junit.Test;

import static org.hamcrest.core.Is.is;
import static org.hamcrest.core.IsEqual.equalTo;
import static org.junit.Assert.assertThat;

public class TraceBackProxyTest {

    public static interface SomeRole {
        void doStuff();
    }

    public static class TestSubject implements SomeRole {
        public boolean called;

        @Override
        public void doStuff() {
            called = true;
        }
    }

    @Test
    public void shouldStillDelegate() {
        TestSubject target = new TestSubject();
        SomeRole action = (SomeRole)TraceBackProxy.wrap(target, new TestOnlyListener());

        action.doStuff();

        assertThat(target.called, is(true));
    }

    public static class TestOnlyListener implements TraceBackProxy.CallingMethodListener {
        String lastCalledMethod;

        @Override
        public void notify(String method) {
            lastCalledMethod = method;
        }
    }

    @Test
    public void shouldFindCallingMethod() {
        TestOnlyListener listener = new TestOnlyListener();
        SomeRole action = (SomeRole) TraceBackProxy.wrap(new TestSubject(), listener);

        action.doStuff();

        assertThat(listener.lastCalledMethod, equalTo("shouldFindCallingMethod"));
    }
}

Note that your mileage may vary since it probably won’t work when you have duplicate method names across classes, or overloaded methods on the same. It proved useful for me and hope it helps you.