JMXServer

package com.mprew.commons.jmx;

import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.net.InetSocketAddress;
import java.net.MalformedURLException;
import java.net.Socket;
import java.net.SocketAddress;
import java.rmi.NoSuchObjectException;
import java.rmi.registry.LocateRegistry;
import java.rmi.registry.Registry;
import java.rmi.server.UnicastRemoteObject;
import java.util.HashMap;

import javax.management.remote.JMXConnectorServer;
import javax.management.remote.JMXConnectorServerFactory;
import javax.management.remote.JMXServiceURL;

/**
 * JMX server connection for main to self publish itself via RMI.
 *
 * @author gwatson
 */
public class JmxServer {

    private Registry rmiRegistry;
    private final int port;
    private JMXConnectorServer connector;

    public JmxServer(int port) {
        this.port = port;
    }

    /**
     * Start our JMX service.
     */
    public synchronized void start() throws MprewJmxException {
        startRmiRegistry();
        startJmxService();
    }

    /**
     * Start the RMI registry.
     */
    private void startRmiRegistry() throws MprewJmxException {
        if (!isRegistryCreated()) {
            createRegistry();
        }
    }

    /**
     * Start our JMX service.
     */
    private void startJmxService() throws MprewJmxException {
        if (connector != null) {
            return;
        }
        JMXServiceURL url = null;
        String urlString = "service:jmx:rmi://localhost:" + (port + 1) + "/jndi/rmi://:" + port + "/jmxrmi";
        try {
            url = new JMXServiceURL(urlString);
        } catch (MalformedURLException e) {
            throw new MprewJmxException("Malformed service url created " + urlString, e);
        }
        try {
            connector =
                    JMXConnectorServerFactory.newJMXConnectorServer(url, new HashMap<String, Object>(),
                            ManagementFactory.getPlatformMBeanServer());
        } catch (IOException e) {
            throw new MprewJmxException("Could not make our Jmx connector server", e);
        }
        try {
            connector.start();
        } catch (IOException e) {
            try {
                // try to start the rmi service again
                createRegistry();
                connector.start();
            } catch (Exception ignored) {
                // ignore exception
            }
            connector = null;
            throw new MprewJmxException("Could not start our Jmx connector server", e);
        }
    }

    private boolean isRegistryCreated() {
        /**
         * After a ton of trial and error, we determined that making the connection to the port is the only reliable way
         * to see if there is an existing registry on the port. Although there is a
         * {@link LocateRegistry#getRegistry(int)} method, there is no easy way to detect if the registry that is in
         * place is a stub or not. Sometimes it is a stub delegating to a real registry and other times it is an empty
         * stub. We also had problems where the registry would be closed but the socket is still open and not closed by
         * the finalizer yet.
         *
         * NOTE: this could be in a different VM or on the way down in unit tests in which case may need to create the
         * registry later.
         */
        Socket socket = new Socket();
        SocketAddress sockAddr = new InetSocketAddress(port);
        try {
            socket.connect(sockAddr);
            // try to connect, if it works then don't start the registry again
            return true;
        } catch (IOException e) {
            // if we get an exception then create a new registry
        } finally {
            try {
                socket.close();
            } catch (IOException e) {
                // ignored
            }
        }
        return false;
    }

    private void createRegistry() throws MprewJmxException {
        try {
            rmiRegistry = LocateRegistry.createRegistry(port);
        } catch (IOException e) {
            throw new MprewJmxException("Unable to create RMI registry on port " + port, e);
        }
    }

    /**
     * Close the JMX server.
     */
    public synchronized void close() throws MprewJmxException {
        if (connector != null) {
            try {
                connector.stop();
            } catch (IOException e) {
                throw new MprewJmxException("Could not stop our Jmx connector server", e);
            } finally {
                connector = null;
            }
        }
        if (rmiRegistry != null) {
            try {
                UnicastRemoteObject.unexportObject(rmiRegistry, true);
            } catch (NoSuchObjectException e) {
                // oh well, we tried
            } finally {
                rmiRegistry = null;
            }
        }
    }
}