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

package jmx.org.apache.cassandra.service;

import static javax.management.JMX.newMBeanProxy;

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

import com.google.common.collect.Multimap;
import com.instaclustr.operations.FunctionWithEx;
import jmx.org.apache.cassandra.CassandraJMXConnectionInfo;
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;
import org.awaitility.Awaitility;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public interface CassandraJMXService {

    static final Logger logger = LoggerFactory.getLogger(CassandraJMXService.class);

    // storage service mbean

     T doWithCassandra4StorageServiceMBean(FunctionWithEx func) throws Exception;

     T doWithCassandra3StorageServiceMBean(FunctionWithEx func) throws Exception;

     T doWithCassandra2StorageServiceMBean(FunctionWithEx func) throws Exception;

     T doWithStorageServiceMBean(FunctionWithEx func) throws Exception;

    // column family store mbean

     T doWithCassandra3ColumnFamilyStoreMBean(FunctionWithEx func, String keyspace, String columnFamily) throws Exception;

     T doWithCassandra4ColumnFamilyStoreMBean(FunctionWithEx func, String keyspace, String columnFamily) throws Exception;

    Multimap getCFSMBeans() throws Exception;

    default  T doWithMBean(FunctionWithEx func,
                                 Class mbeanClass,
                                 ObjectName objectName) throws Exception {
        return doWithMBean(func, mbeanClass, objectName, getCassandraJmxConnectionInfo());
    }

    default  T doWithMBean(FunctionWithEx func,
                                 Class mbeanClass,
                                 ObjectName objectName,
                                 CassandraJMXConnectionInfo jmxConnectionInfo) throws Exception {

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

            jmxConnector.connect();

            final MBeanServerConnection mBeanServerConnection = jmxConnector.getMBeanServerConnection();

            final U proxy = newMBeanProxy(mBeanServerConnection, objectName, mbeanClass);

            waitUntilRegistered(mBeanServerConnection, objectName);

            return func.apply(proxy);
        }
    }

    default  void waitUntilRegistered(final MBeanServerConnection mBeanServerConnection,
                                         final ObjectName objectName) {
        Awaitility.await()
            .pollInterval(1, TimeUnit.SECONDS)
            .timeout(1, TimeUnit.MINUTES)
            .until(() -> {
                boolean registered = mBeanServerConnection.isRegistered(objectName);

                if (!registered) {
                    logger.info(String.format("Waiting for %s to be registered.", objectName));
                }

                return registered;
            });
    }

    default  T doWithMBean(FunctionWithEx func,
                                 Class mbeanClass,
                                 String query,
                                 CassandraJMXConnectionInfo jmxConnectionInfo) throws Exception {

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

            jmxConnector.connect();

            MBeanServerConnection mBeanServerConnection = jmxConnector.getMBeanServerConnection();

            Set objectNames = mBeanServerConnection.queryNames(new ObjectName(query), null);

            if (objectNames.isEmpty()) {
                throw new IllegalStateException(String.format("Could not find ObjectName with query %s", query));
            }

            if (objectNames.size() != 1) {
                throw new IllegalStateException(String.format("There is more than one ObjectName returned by query %s. They are: %s",
                                                              query,
                                                              objectNames.stream().map(ObjectName::getCanonicalName)));
            }

            final U proxy = newMBeanProxy(mBeanServerConnection, objectNames.iterator().next(), mbeanClass);

            final ObjectName objectName = objectNames.iterator().next();

            waitUntilRegistered(mBeanServerConnection, objectName);

            return func.apply(proxy);
        }
    }

    CassandraJMXConnectionInfo getCassandraJmxConnectionInfo();
}