Limit MBeans query to certain scopes

src/main/java/org/datadog/jmxfetch/Configuration.java

@@ -1,6 +1,14 @@
 package org.datadog.jmxfetch;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.Map.Entry;
+import java.util.Set;
 
 public class Configuration {
 
@@ -9,7 +17,9 @@ public class Configuration {
     private Filter exclude;
 
     /**
-     * A simple class to access configuration elements more easily
+     * Access configuration elements more easily
+     *
+     * Also provides helper methods to extract common information among filters.
      */
     public Configuration(LinkedHashMap<String, Object> conf) {
         this.conf = conf;
@@ -20,7 +30,6 @@ public Configuration(LinkedHashMap<String, Object> conf) {
     public LinkedHashMap<String, Object> getConf() {
         return conf;
     }
-
     public Filter getInclude() {
         return include;
     }
@@ -33,4 +42,205 @@ public String toString() {
         return "include: " + this.include + " - exclude: " + this.exclude;
     }
 
+    private Boolean hasInclude(){
+        return getInclude() != null;
+    }
+
+    /**
+     * Filter a configuration list to keep the ones with `include` filters.
+     *
+     * @param configurationList         the configuration list to filter
+     *
+     * @return                          a configuration list
+     */
+    private static LinkedList<Configuration> getIncludeConfigurationList(LinkedList<Configuration> configurationList){
+        LinkedList<Configuration> includeConfigList = new LinkedList<Configuration>(configurationList);
+        Iterator<Configuration> confItr = includeConfigList.iterator();
+
+        while(confItr.hasNext()) {
+            Configuration conf = confItr.next();
+            if (!conf.hasInclude()) {
+                confItr.remove();
+            }
+        }
+        return includeConfigList;
+    }
+
+    /**
+     * Extract `include` filters from the configuration list and index then by domain name.
+     *
+     * @param configurationList         the configuration list to process
+     *
+     * @return                          filters by domain name
+     */
+    private static HashMap<String, LinkedList<Filter>> getIncludeFiltersByDomain(LinkedList<Configuration> configurationList){
+        HashMap<String, LinkedList<Filter>> includeFiltersByDomain = new HashMap<String, LinkedList<Filter>>();
+
+        for (Configuration conf : configurationList) {
+            Filter filter = conf.getInclude();
+            LinkedList<Filter> filters = new LinkedList<Filter>();
+
+            // Convert bean name, to a proper filter, i.e. a hash
+            if (!filter.isEmptyBeanName()) {
+                ArrayList<String> beanNames = filter.getBeanNames();
+
+                for (String beanName : beanNames) {
+                    String[] splitBeanName = beanName.split(":");
+                    String domain = splitBeanName[0];
+                    String rawBeanParameters = splitBeanName[1];
+                    LinkedList<String> beanParameters = new LinkedList<String>(Arrays.asList(new String(rawBeanParameters).replace("=", ":").split(",")));
+                    HashMap<String, String> beanParametersHash = JMXAttribute.getBeanParametersHash(beanParameters);
+                    beanParametersHash.put("domain", domain);
+                    filters.add(new Filter(beanParametersHash));
+                }
+            } else {
+                filters.add(filter);
+            }
+
+            for (Filter f: filters) {
+                //  Retrieve the existing filters for the domain, add the new filters
+                LinkedList<Filter> domainFilters;
+                String domainName = f.getDomain();
+
+                if (includeFiltersByDomain.containsKey(domainName)) {
+                    domainFilters = includeFiltersByDomain.get(domainName);
+                } else {
+                    domainFilters = new LinkedList<Filter>();
+                }
+
+                domainFilters.add(f);
+                includeFiltersByDomain.put(domainName, domainFilters);
+            }
+        }
+        return includeFiltersByDomain;
+    }
+
+    /**
+     * Extract, among filters, bean key parameters in common.
+     *
+     * @param filtersByDomain       filters by domain name
+     *
+     * @return                      common bean key parameters by domain name
+     */
+    private static HashMap<String, Set<String>> getCommonBeanKeysByDomain(HashMap<String, LinkedList<Filter>> filtersByDomain){
+        HashMap<String, Set<String>> beanKeysIntersectionByDomain = new HashMap<String,Set<String>>();
+
+        for (Entry<String, LinkedList<Filter>> filtersEntry : filtersByDomain.entrySet()) {
+            String domainName = filtersEntry.getKey();
+            LinkedList<Filter> mFilters= filtersEntry.getValue();
+
+            // Compute keys intersection
+            Set<String> keysIntersection = new HashSet<String>(mFilters.getFirst().keySet());
+
+            for (Filter f: mFilters) {
+                keysIntersection.retainAll(f.keySet());
+            }
+
+            // Remove special parameters
+            for(String param : JMXAttribute.getExcludedBeanParams()){
+                keysIntersection.remove(param);
+            }
+
+            beanKeysIntersectionByDomain.put(domainName, keysIntersection);
+        }
+
+        return beanKeysIntersectionByDomain;
+    }
+
+    /**
+     * Build a map of common bean keys->values, with the specified bean keys, among the given filters.
+     *
+     * @param beanKeysByDomain      bean keys by domain name
+     * @param filtersByDomain       filters by domain name
+     *
+     * @return                      bean pattern (keys->values) by domain name
+     */
+    private static HashMap<String, LinkedHashMap<String, String>> getCommonScopeByDomain(HashMap<String, Set<String>> beanKeysByDomain, HashMap<String, LinkedList<Filter>> filtersByDomain){
+        // Compute a common scope a among filters by domain name
+        HashMap<String, LinkedHashMap<String, String>> commonScopeByDomain = new HashMap<String, LinkedHashMap<String, String>>();
+
+        for (Entry<String, Set<String>> commonParametersByDomainEntry : beanKeysByDomain.entrySet()) {
+            String domainName = commonParametersByDomainEntry.getKey();
+            Set<String> commonParameters = commonParametersByDomainEntry.getValue();
+            LinkedList<Filter> filters = filtersByDomain.get(domainName);
+            LinkedHashMap<String, String> commonScope = new LinkedHashMap<String, String>();
+
+            for (String parameter : commonParameters) {
+                // Check if all values associated with the parameters are the same
+                String commonValue = null;
+                Boolean hasCommonValue = true;
+
+                for (Filter f : filters) {
+                    ArrayList<String> parameterValues = f.getParameterValues(parameter);
+
+                    if (parameterValues.size() != 1 || (commonValue != null && !commonValue.equals(parameterValues.get(0)))) {
+                        hasCommonValue = false;
+                        break;
+                 }
+
+                commonValue = parameterValues.get(0);
+
+             }
+             if (hasCommonValue) {
+                commonScope.put(parameter, commonValue);
+             }
+         }
+        commonScopeByDomain.put(domainName, commonScope);
+     }
+
+    return commonScopeByDomain;
+    }
+
+    /**
+     * Stringify a bean pattern.
+     *
+     * @param domain                domain name
+     * @param beanScope             map of bean keys-> values
+     *
+     * @return                      string pattern identifying the bean scope
+     */
+    private static String beanScopeToString(String domain, LinkedHashMap<String, String> beanScope){
+        String result = "";
+
+        // Domain
+        domain = (domain != null) ? domain : "*";
+        result += domain + ":";
+
+        // Scope parameters
+        for (Entry<String, String> beanScopeEntry : beanScope.entrySet()) {
+            String param = beanScopeEntry.getKey();
+            String value = beanScopeEntry.getValue();
+
+            result += param + "=" + value + ",";
+        }
+        result += "*";
+
+        return result;
+    }
+
+    /**
+     * Find, among the configuration list, a potential common bean pattern by domain name.
+     *
+     * @param configurationList         the configuration list to process
+     *
+     * @return                          common bean pattern strings
+     */
+    public static LinkedList<String> getGreatestCommonScopes(LinkedList<Configuration> configurationList){
+        LinkedList<Configuration> includeConfigList = getIncludeConfigurationList(configurationList);
+        HashMap<String, LinkedList<Filter>> includeFiltersByDomain = getIncludeFiltersByDomain(includeConfigList);
+        HashMap<String, Set<String>> parametersIntersectionByDomain = getCommonBeanKeysByDomain(includeFiltersByDomain);
+        HashMap<String, LinkedHashMap<String, String>> commonBeanScopeByDomain = getCommonScopeByDomain(parametersIntersectionByDomain, includeFiltersByDomain);
+
+        LinkedList<String> result = new LinkedList<String>();
+
+        for (Entry<String, LinkedHashMap<String, String>> beanScopeEntry: commonBeanScopeByDomain.entrySet()) {
+            String domain = beanScopeEntry.getKey();
+            LinkedHashMap<String, String> beanScope = beanScopeEntry.getValue();
+
+            result.add(beanScopeToString(domain, beanScope));
+        }
+
+        return result;
+    }
+
 }

src/main/java/org/datadog/jmxfetch/Connection.java

@@ -48,8 +48,10 @@ public MBeanAttributeInfo[] getAttributesForBean(ObjectName bean_name)
         return mbs.getMBeanInfo(bean_name).getAttributes();
     }
 
-    public Set<ObjectInstance> queryMBeans() throws IOException {
-        return mbs.queryMBeans(null, null);
+    public Set<ObjectInstance> queryMBeans(ObjectName name) throws IOException {
+        String scope = (name != null) ? name.toString() : "*:*";
+        LOGGER.debug("Querying beans on scope: " + scope);
+        return mbs.queryMBeans(name, null);
     }
 
     protected void createConnection() throws IOException {

src/main/java/org/datadog/jmxfetch/Filter.java

@@ -1,14 +1,14 @@
 package org.datadog.jmxfetch;
 
-import java.util.LinkedHashMap;
+import java.util.HashMap;
 import java.util.ArrayList;
 import java.util.Set;
 import java.lang.ClassCastException;
 import java.util.regex.Pattern;
 
 
 class Filter {
-    LinkedHashMap<String, Object> filter;
+    HashMap<String, Object> filter;
     Pattern domainRegex;
     ArrayList<Pattern> beanRegexes = null;
 
@@ -23,11 +23,11 @@ class Filter {
 
     @SuppressWarnings("unchecked")
     public Filter(Object filter) {
-        LinkedHashMap<String, Object> castFilter;
+        HashMap<String, Object> castFilter;
         if (filter != null) {
-            castFilter = (LinkedHashMap<String, Object>) filter;
+            castFilter = (HashMap<String, Object>) filter;
         } else {
-            castFilter = new LinkedHashMap<String, Object>();
+            castFilter = new HashMap<String, Object>();
         }
         this.filter = castFilter;
     }
@@ -40,7 +40,6 @@ public Set<String> keySet() {
         return filter.keySet();
     }
 
-
     @SuppressWarnings({ "unchecked", "serial" })
     private static ArrayList<String> toStringArrayList(final Object toCast) {
         // Return object as an ArrayList wherever it's defined as

src/main/java/org/datadog/jmxfetch/Instance.java

@@ -9,7 +9,6 @@
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 
 import javax.management.MBeanAttributeInfo;
@@ -31,6 +30,7 @@ public class Instance {
     private final static int DEFAULT_REFRESH_BEANS_PERIOD = 600;
 
     private Set<ObjectInstance> beans;
+    private LinkedList<String> beanScopes;
     private LinkedList<Configuration> configurationList = new LinkedList<Configuration>();
     private LinkedList<JMXAttribute> matchingAttributes;
     private HashSet<JMXAttribute> failingAttributes;
@@ -263,9 +263,27 @@ private void getMatchingAttributes() {
         LOGGER.info("Found " + matchingAttributes.size() + " matching attributes");
     }
 
-    private void refreshBeansList() throws IOException {
+    public LinkedList<String> getBeansScopes(){
+        if(this.beanScopes == null){
+            this.beanScopes = Configuration.getGreatestCommonScopes(configurationList);
+        }
+        return this.beanScopes;
+    }
 
-        this.beans = connection.queryMBeans();
+    private void refreshBeansList() throws IOException {
+        this.beans = new HashSet<ObjectInstance>();
+        try {
+            LinkedList<String> beanScopes = getBeansScopes();
+            for (String scope : beanScopes) {
+                ObjectName name = new ObjectName(scope);
+                this.beans.addAll(connection.queryMBeans(name));
+            }
+        }
+        catch (Exception e) {
+            e.printStackTrace();
+            LOGGER.error("Unable to compute a common bean scope, querying all beans as a fallback");
+            this.beans = connection.queryMBeans(null);
+        }
         this.lastRefreshTime = System.currentTimeMillis();
     }
 

src/main/java/org/datadog/jmxfetch/JMXAttribute.java

@@ -64,7 +64,7 @@ public abstract class JMXAttribute {
         this.defaultTagsList = renameConflictingParameters(beanParametersList);
     }
 
-    private static HashMap<String, String> getBeanParametersHash(LinkedList<String> beanParameters) {
+    public static HashMap<String, String> getBeanParametersHash(LinkedList<String> beanParameters) {
         HashMap<String, String> beanParams = new HashMap<String, String>();
         for (String param : beanParameters) {
             String[] paramSplit = param.split(":");
@@ -355,4 +355,8 @@ String getBeanName() {
     String getAttributeName() {
         return attributeName;
     }
+
+    public static List<String> getExcludedBeanParams(){
+        return EXCLUDED_BEAN_PARAMS;
+    }
 }