More on JMX...

My previous entry about JMX got some real feedback from people that are working on JMX. In appreciation of that, I wrote a little utility in the hope of contributing back to JMX.

One of the problems I had was the lack of "weak MBean" support ---- I didn't want the MBeanServer to hold a strong reference to my MBean, because that prevents those objects from GC-ed. Eamonn McManus suggested a technique to handle this, but I didn't like that I have to write such class for each MBean that I have.

So instead, I wrote a little utility code that adds weak MBean support to any MBean. It almost makes it look like the MBeanServer supports such functionality by itself. In this way, I didn't need to write a weak MBean wrapper for my MBeans.

The trick behind this is java.lang.reflect.Proxy class, which allows me to generate implementations of any MBean interface at runtime.

package dalma.container;

import javax.management.JMException;
import javax.management.MBeanRegistration;
import javax.management.MBeanServer;
import javax.management.ObjectName;
import javax.management.StandardMBean;
import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;

/**
 * Proxy MBean that avoid strong reference to the real MBean object.
 *
 * @author Kohsuke Kawaguchi
 */
final class MBeanProxy implements InvocationHandler, MBeanRegistration {

    /**
     * Creates a proxy MBean and registers it to the server,
     * overriding the existing mbean if necessary.
     * 
     * @param server
     *   MBean will be registered to this server.
     * @param name
     *   The name under which the MBean will be registered
     * @param mbeanInterface
     *   MBean interface to be exposed
     * @param object
     *   MBean instance to be exposed
     */
    public static <T> void register( MBeanServer server, ObjectName name, Class<T> mbeanInterface, T object ) throws JMException {
        Object proxy = mbeanInterface.cast(
            Proxy.newProxyInstance(mbeanInterface.getClassLoader(),
            new Class[]{mbeanInterface,MBeanRegistration.class},
            new MBeanProxy(object) ));

        if(server.isRegistered(name)) {
            try {
                server.unregisterMBean(name);
            } catch (JMException e) {
                // if we fail to unregister, try to register ours anyway.
                // maybe a GC kicked in in-between.
            }
        }

        // since the proxy class has random names like '$Proxy1',
        // we need to use StandardMBean to designate a management interface
        server.registerMBean(new StandardMBean(proxy,mbeanInterface),name);

    }

    /**
     * The real MBean object.
     */
    private final ReferenceImpl real;
    private MBeanServer server;
    private ObjectName name;

    private MBeanProxy(Object realObject) {
        this.real = new ReferenceImpl(realObject);
    }

    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        Object o = real.get();

        if(method.getDeclaringClass()==MBeanRegistration.class) {
            o = this;
        }

        if(o==null) {
            unregister();
            throw new IllegalStateException(name+" no longer exists");
        }

        try {
            return method.invoke(o,args);
        } catch (InvocationTargetException e) {
            if(e.getCause()!=null)
                throw e.getCause();
            else
                throw e;
        }
    }

    private void unregister() {
        try {
            server.unregisterMBean(name);
        } catch (JMException e) {
            throw new Error(e); // is this even possible?
        }
    }

    public ObjectName preRegister(MBeanServer server, ObjectName name) throws Exception {
        this.server = server;
        this.name = name;
        return name;
    }

    public void postRegister(Boolean registrationDone) {
        // noop
    }

    public void preDeregister() throws Exception {
        // noop
    }

    public void postDeregister() {
        server = null;
        name = null;
    }

    private class ReferenceImpl extends WeakReference<Object> {
        public ReferenceImpl(Object referent) {
            super(referent,queue);
        }

        public MBeanProxy getProxy() {
            return MBeanProxy.this;
        }
    }

    private static final ReferenceQueue<Object> queue = new ReferenceQueue<Object>();

    static {
        Runnable r = new Runnable() {
            public void run() {
                while (true) {
                    ReferenceImpl ref;
                    try {
                        ref = (ReferenceImpl) (WeakReference<Object>)queue.remove();
                    } catch (InterruptedException e) {
                        return;
                    }
                    ref.getProxy().unregister();
                }
            }
        };
        Thread t = new Thread(r, "Dalma JMX bean clenaer");
        t.setDaemon(true);
        t.start();
    }
}