Support includeObjectNameAttributes (#1042)

* Support includeObjectNameAttributes and integration test

Signed-off-by: Vincent Daniel <[email protected]>

collector/src/main/java/io/prometheus/jmx/JmxScraper.java

@@ -172,7 +172,14 @@ private void scrapeBean(MBeanServerConnection beanConn, ObjectName mBeanName) {
                 continue;
             }
 
-            name2MBeanAttributeInfo.put(mBeanAttributeInfo.getName(), mBeanAttributeInfo);
+            if (objectNameAttributeFilter.includeObjectNameAttributesIsEmpty()) {
+                name2MBeanAttributeInfo.put(mBeanAttributeInfo.getName(), mBeanAttributeInfo);
+                continue;
+            }
+
+            if (objectNameAttributeFilter.include(mBeanName, mBeanAttributeInfo.getName())) {
+                name2MBeanAttributeInfo.put(mBeanAttributeInfo.getName(), mBeanAttributeInfo);
+            }
         }
 
         if (name2MBeanAttributeInfo.isEmpty()) {

collector/src/main/java/io/prometheus/jmx/ObjectNameAttributeFilter.java

@@ -37,19 +37,23 @@ public class ObjectNameAttributeFilter {
     /** Configuration constant to define a mapping of ObjectNames to attribute names */
     public static final String EXCLUDE_OBJECT_NAME_ATTRIBUTES = "excludeObjectNameAttributes";
 
+    public static final String INCLUDE_OBJECT_NAME_ATTRIBUTES = "includeObjectNameAttributes";
+
     /** Configuration constant to enable auto ObjectName attributes filtering */
     public static final String AUTO_EXCLUDE_OBJECT_NAME_ATTRIBUTES =
             "autoExcludeObjectNameAttributes";
 
     private final Map<ObjectName, Set<String>> configExcludeObjectNameAttributesMap;
     private final Map<ObjectName, Set<String>> dynamicExcludeObjectNameAttributesMap;
+    private final Map<ObjectName, Set<String>> includeObjectNameAttributesMap;
 
     private boolean autoExcludeObjectNameAttributes;
 
     /** Constructor */
     private ObjectNameAttributeFilter() {
         configExcludeObjectNameAttributesMap = new ConcurrentHashMap<>();
         dynamicExcludeObjectNameAttributesMap = new ConcurrentHashMap<>();
+        includeObjectNameAttributesMap = new ConcurrentHashMap<>();
     }
 
     /**
@@ -61,33 +65,48 @@ private ObjectNameAttributeFilter() {
      */
     private ObjectNameAttributeFilter initialize(Map<String, Object> yamlConfig)
             throws MalformedObjectNameException {
-        if (yamlConfig.containsKey(EXCLUDE_OBJECT_NAME_ATTRIBUTES)) {
-            Map<Object, Object> objectNameAttributeMap =
-                    (Map<Object, Object>) yamlConfig.get(EXCLUDE_OBJECT_NAME_ATTRIBUTES);
+        initializeObjectNameAttributes(
+                yamlConfig, EXCLUDE_OBJECT_NAME_ATTRIBUTES, configExcludeObjectNameAttributesMap);
+        initializeObjectNameAttributes(
+                yamlConfig, INCLUDE_OBJECT_NAME_ATTRIBUTES, includeObjectNameAttributesMap);
+        if (yamlConfig.containsKey(AUTO_EXCLUDE_OBJECT_NAME_ATTRIBUTES)) {
+            autoExcludeObjectNameAttributes =
+                    (Boolean) yamlConfig.get(AUTO_EXCLUDE_OBJECT_NAME_ATTRIBUTES);
+        } else {
+            autoExcludeObjectNameAttributes = true;
+        }
+
+        LOGGER.log(Level.FINE, "dynamicExclusion [%b]", autoExcludeObjectNameAttributes);
+
+        return this;
+    }
+
+    /**
+     * Method to initialize the configExcludeObjectNameAttributesMap and
+     * includeObjectNameAttributesMap
+     *
+     * @throws MalformedObjectNameException MalformedObjectNameException
+     */
+    private void initializeObjectNameAttributes(
+            Map<String, Object> yamlConfig,
+            String key,
+            Map<ObjectName, Set<String>> objectNameAttributesMap)
+            throws MalformedObjectNameException {
+        if (yamlConfig.containsKey(key)) {
+            Map<Object, Object> objectNameAttributeMap = (Map<Object, Object>) yamlConfig.get(key);
 
             for (Map.Entry<Object, Object> entry : objectNameAttributeMap.entrySet()) {
                 ObjectName objectName = new ObjectName((String) entry.getKey());
 
                 List<String> attributeNames = (List<String>) entry.getValue();
 
                 Set<String> attributeNameSet =
-                        configExcludeObjectNameAttributesMap.computeIfAbsent(
+                        objectNameAttributesMap.computeIfAbsent(
                                 objectName, o -> Collections.synchronizedSet(new HashSet<>()));
 
                 attributeNameSet.addAll(attributeNames);
             }
         }
-
-        if (yamlConfig.containsKey(AUTO_EXCLUDE_OBJECT_NAME_ATTRIBUTES)) {
-            autoExcludeObjectNameAttributes =
-                    (Boolean) yamlConfig.get(AUTO_EXCLUDE_OBJECT_NAME_ATTRIBUTES);
-        } else {
-            autoExcludeObjectNameAttributes = true;
-        }
-
-        LOGGER.log(Level.FINE, "dynamicExclusion [%b]", autoExcludeObjectNameAttributes);
-
-        return this;
     }
 
     /**
@@ -153,6 +172,21 @@ private boolean exclude(
         return result;
     }
 
+    public boolean include(ObjectName objectName, String attributeName) {
+        boolean result = false;
+
+        Set<String> attributeNameSet = includeObjectNameAttributesMap.get(objectName);
+        if (attributeNameSet != null) {
+            result = attributeNameSet.contains(attributeName);
+        }
+
+        return result;
+    }
+
+    public boolean includeObjectNameAttributesIsEmpty() {
+        return includeObjectNameAttributesMap.isEmpty();
+    }
+
     /**
      * Method to create an ObjectNameAttributeFilter
      *

collector/src/test/java/io/prometheus/jmx/ObjectNameAttributeFilterTest.java

@@ -30,28 +30,36 @@ public class ObjectNameAttributeFilterTest {
 
     static final String CONFIG_EXCLUDE_MAP_FIELD = "configExcludeObjectNameAttributesMap";
     static final String DYNAMIC_EXCLUDE_MAP_FIELD = "dynamicExcludeObjectNameAttributesMap";
+    static final String CONFIG_INCLUDE_MAP_FIELD = "includeObjectNameAttributesMap";
 
     @Test
     public void emptyConfig() throws Exception {
         ObjectNameAttributeFilter filter = initEmptyConfigFilter();
 
         assertEquals(0, excludeMapSize(filter, CONFIG_EXCLUDE_MAP_FIELD));
         assertEquals(0, excludeMapSize(filter, DYNAMIC_EXCLUDE_MAP_FIELD));
+        assertEquals(0, includeMapSize(filter));
     }
 
     @Test
     public void nonEmptyConfigShouldInitializeConfigExcludeMap() throws Exception {
         ObjectNameAttributeFilter filter = initNonEmptyConfigFilter();
 
-        Map<ObjectName, Set<String>> configMap =
+        Map<ObjectName, Set<String>> configExcludeMap =
                 getInternalMapValue(filter, CONFIG_EXCLUDE_MAP_FIELD);
         Map<ObjectName, Set<String>> dynamicMap =
                 getInternalMapValue(filter, DYNAMIC_EXCLUDE_MAP_FIELD);
+        Map<ObjectName, Set<String>> configIncludeMap =
+                getInternalMapValue(filter, CONFIG_INCLUDE_MAP_FIELD);
 
-        assertEquals(2, configMap.size());
-        assertEquals(1, configMap.get(new ObjectName("java.lang:type=OperatingSystem")).size());
-        assertEquals(2, configMap.get(new ObjectName("java.lang:type=Runtime")).size());
+        assertEquals(2, configExcludeMap.size());
+        assertEquals(
+                1, configExcludeMap.get(new ObjectName("java.lang:type=OperatingSystem")).size());
+        assertEquals(2, configExcludeMap.get(new ObjectName("java.lang:type=Runtime")).size());
         assertEquals(0, dynamicMap.size());
+        assertEquals(2, configIncludeMap.size());
+        assertEquals(1, configIncludeMap.get(new ObjectName("java.lang:type=Threading")).size());
+        assertEquals(1, configIncludeMap.get(new ObjectName("java.lang:type=ClassLoading")).size());
     }
 
     @Test
@@ -136,9 +144,22 @@ public void onlyKeepMBeansShouldNotRemoveConfigRecords() throws Exception {
         }
     }
 
+    @Test
+    public void includeShouldCheckConfigInclusionMap() throws Exception {
+        ObjectNameAttributeFilter filter = initNonEmptyConfigFilter();
+
+        boolean result = filter.include(new ObjectName("java.lang:type=Threading"), "ThreadCount");
+
+        assertTrue("java.lang:type=Threading<>ThreadCount should be included by config", result);
+    }
+
     private static ObjectNameAttributeFilter initEmptyConfigFilter() {
         return ObjectNameAttributeFilter.create(
-                new Yaml().load("---\nexcludeObjectNameAttributes: {}\n"));
+                new Yaml()
+                        .load(
+                                "---\n"
+                                        + "excludeObjectNameAttributes: {}\n"
+                                        + "includeObjectNameAttributes: {}\n"));
     }
 
     private static ObjectNameAttributeFilter initNonEmptyConfigFilter() {
@@ -151,7 +172,12 @@ private static ObjectNameAttributeFilter initNonEmptyConfigFilter() {
                                         + "    - \"ObjectName\"\n"
                                         + "  \"java.lang:type=Runtime\":\n"
                                         + "    - \"ClassPath\"\n"
-                                        + "    - \"SystemProperties\"\n"));
+                                        + "    - \"SystemProperties\"\n"
+                                        + "includeObjectNameAttributes:\n"
+                                        + "  \"java.lang:type=Threading\":\n"
+                                        + "    - \"ThreadCount\"\n"
+                                        + "  \"java.lang:type=ClassLoading\":\n"
+                                        + "    - \"LoadedClassCount\"\n"));
     }
 
     private static Set<ObjectName> getAliveMBeans() throws MalformedObjectNameException {
@@ -167,6 +193,10 @@ private int excludeMapSize(ObjectNameAttributeFilter filter, String mapName) thr
         return getInternalMapValue(filter, mapName).size();
     }
 
+    private int includeMapSize(ObjectNameAttributeFilter filter) throws Exception {
+        return getInternalMapValue(filter, CONFIG_INCLUDE_MAP_FIELD).size();
+    }
+
     private Map getInternalMapValue(ObjectNameAttributeFilter filter, String mapName)
             throws Exception {
         return getInternalFieldValue(filter, mapName, Map.class);

docs/README.md

@@ -134,6 +134,7 @@ includeObjectNames | A list of [ObjectNames](https://docs.oracle.com/javase/8/do
 excludeObjectNames | A list of [ObjectNames](https://docs.oracle.com/javase/8/docs/api/javax/management/ObjectName.html) to not query. Takes precedence over `includeObjectNames`. Defaults to none.
 autoExcludeObjectNameAttributes | Whether to auto exclude [ObjectName](https://docs.oracle.com/javase/8/docs/api/javax/management/ObjectName.html) attributes that can't be converted to standard metrics types. Defaults to `true`. 
 excludeObjectNameAttributes | A Map of [ObjectNames](https://docs.oracle.com/javase/8/docs/api/javax/management/ObjectName.html) with a list of attribute names to exclude. ObjectNames must be in canonical form. Both ObjectNames and attribute names are matched as a Strings (no regex.) Optional.
+includeObjectNameAttributes | A Map of [ObjectNames](https://docs.oracle.com/javase/8/docs/api/javax/management/ObjectName.html) with a list of attribute names to include. ObjectNames must be in canonical form. Both ObjectNames and attribute names are matched as a Strings (no regex.) Optional.
 rules      | A list of rules to apply in order, processing stops at the first matching rule. Attributes that aren't matched aren't collected. If not specified, defaults to collecting everything in the default format.
 pattern           | Regex pattern to match against each bean attribute. The pattern is not anchored. Capture groups can be used in other options. Defaults to matching everything.
 attrNameSnakeCase | Converts the attribute name to snake case. This is seen in the names matched by the pattern and the default format. For example, anAttrName to an\_attr\_name. Defaults to false.