jmx.org.apache.cassandra.service.CassandraJMXServiceImpl Maven / Gradle / Ivy

package jmx.org.apache.cassandra.service;

import javax.management.JMX;
import javax.management.MBeanServerConnection;
import javax.management.ObjectName;
import javax.management.remote.JMXConnector;
import java.util.Set;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import com.instaclustr.operations.FunctionWithEx;
import jmx.org.apache.cassandra.CassandraJMXConnectionInfo;
import jmx.org.apache.cassandra.CassandraObjectNames;
import jmx.org.apache.cassandra.CassandraObjectNames.V2;
import jmx.org.apache.cassandra.CassandraObjectNames.V4;
import jmx.org.apache.cassandra.JMXUtils;
import jmx.org.apache.cassandra.service.cassandra2.Cassandra2StorageServiceMBean;
import jmx.org.apache.cassandra.service.cassandra3.ColumnFamilyStoreMBean;
import jmx.org.apache.cassandra.service.cassandra3.StorageServiceMBean;
import jmx.org.apache.cassandra.service.cassandra4.Cassandra4ColumnFamilyStoreMBean;
import jmx.org.apache.cassandra.service.cassandra4.Cassandra4StorageServiceMBean;

public class CassandraJMXServiceImpl implements CassandraJMXService {

    private final CassandraJMXConnectionInfo jmxConnectionInfo;

    public CassandraJMXServiceImpl(CassandraJMXConnectionInfo jmxConnectionInfo) {
        this.jmxConnectionInfo = jmxConnectionInfo;
    }

    @Override
    public  T doWithCassandra4StorageServiceMBean(final FunctionWithEx func) throws Exception {
        return doWithMBean(func,
                           Cassandra4StorageServiceMBean.class,
                           V4.STORAGE_SERVICE_MBEAN_NAME,
                           jmxConnectionInfo);
    }

    @Override
    public  T doWithCassandra3StorageServiceMBean(final FunctionWithEx func) throws Exception {
        return doWithStorageServiceMBean(func);
    }

    @Override
    public  T doWithCassandra2StorageServiceMBean(final FunctionWithEx func) throws Exception {
        return doWithMBean(func,
                           Cassandra2StorageServiceMBean.class,
                           V2.STORAGE_SERVICE_MBEAN_NAME,
                           jmxConnectionInfo);
    }

    @Override
    public  T doWithStorageServiceMBean(final FunctionWithEx func) throws Exception {
        return doWithMBean(func,
                           StorageServiceMBean.class,
                           CassandraObjectNames.V3.STORAGE_SERVICE_MBEAN_NAME,
                           jmxConnectionInfo);
    }

    @Override
    public  T doWithCassandra3ColumnFamilyStoreMBean(final FunctionWithEx func,
                                                        final String keyspace,
                                                        final String columnFamily) throws Exception {
        return doWithMBean(func, ColumnFamilyStoreMBean.class,
                           getColumnFamilyMBeanObjectNameQuery(keyspace, columnFamily),
                           jmxConnectionInfo);
    }

    @Override
    public  T doWithCassandra4ColumnFamilyStoreMBean(final FunctionWithEx func,
                                                        final String keyspace,
                                                        final String columnFamily) throws Exception {
        return doWithMBean(func,
                           Cassandra4ColumnFamilyStoreMBean.class,
                           getColumnFamilyMBeanObjectNameQuery(keyspace, columnFamily),
                           jmxConnectionInfo);
    }

    @Override
    public Multimap getCFSMBeans() throws Exception {

        Multimap cfsMBeans = HashMultimap.create();

        try (JMXConnector jmxConnector = JMXUtils.getJmxConnector(jmxConnectionInfo)) {

            jmxConnector.connect();

            final MBeanServerConnection mBeanServerConnection = jmxConnector.getMBeanServerConnection();

            final ObjectName query = new ObjectName("org.apache.cassandra.db:type=ColumnFamilies,*");
            final Set cfObjects = mBeanServerConnection.queryNames(query, null);

            for (final ObjectName name : cfObjects) {
                String keyspace = name.getKeyProperty("keyspace");
                ColumnFamilyStoreMBean cfsProxy = JMX.newMBeanProxy(mBeanServerConnection, name, ColumnFamilyStoreMBean.class);
                cfsMBeans.put(keyspace, cfsProxy);
            }

            return cfsMBeans;
        }
    }

    @Override
    public CassandraJMXConnectionInfo getCassandraJmxConnectionInfo() {
        return jmxConnectionInfo;
    }

    private String getColumnFamilyMBeanObjectNameQuery(final String keyspace, final String columnFamily) {
        final String type = columnFamily.contains(".") ? "IndexColumnFamilies" : "ColumnFamilies";
        return "org.apache.cassandra.db:type=*" + type + ",keyspace=" + keyspace + ",columnfamily=" + columnFamily;
    }
}