jmx state metrics

instrumentation/jmx-metrics/library/src/main/java/io/opentelemetry/instrumentation/jmx/engine/BeanAttributeExtractor.java

@@ -253,7 +253,8 @@ private Object extractAttributeValue(MBeanServerConnection connection, ObjectNam
   }
 
   @Nullable
-  Number extractNumericalAttribute(MBeanServerConnection connection, ObjectName objectName) {
+  protected Number extractNumericalAttribute(
+      MBeanServerConnection connection, ObjectName objectName) {
     Object value = extractAttributeValue(connection, objectName);
     if (value instanceof Number) {
       return (Number) value;

instrumentation/jmx-metrics/library/src/main/java/io/opentelemetry/instrumentation/jmx/engine/MetricAttribute.java

@@ -22,6 +22,10 @@ public MetricAttribute(String name, MetricAttributeExtractor extractor) {
     this.extractor = extractor;
   }
 
+  public boolean isStateAttribute() {
+    return extractor == null;
+  }
+
   public String getAttributeName() {
     return name;
   }

instrumentation/jmx-metrics/library/src/main/java/io/opentelemetry/instrumentation/jmx/yaml/JmxRule.java

@@ -9,6 +9,7 @@
 import io.opentelemetry.instrumentation.jmx.engine.BeanAttributeExtractor;
 import io.opentelemetry.instrumentation.jmx.engine.BeanGroup;
 import io.opentelemetry.instrumentation.jmx.engine.MetricAttribute;
+import io.opentelemetry.instrumentation.jmx.engine.MetricAttributeExtractor;
 import io.opentelemetry.instrumentation.jmx.engine.MetricDef;
 import io.opentelemetry.instrumentation.jmx.engine.MetricExtractor;
 import io.opentelemetry.instrumentation.jmx.engine.MetricInfo;
@@ -17,7 +18,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import javax.management.MBeanServerConnection;
 import javax.management.MalformedObjectNameException;
 import javax.management.ObjectName;
 
@@ -143,8 +146,7 @@ public MetricDef buildMetricDef() throws Exception {
     }
 
     Set<String> attrNames = mapping.keySet();
-    MetricExtractor[] metricExtractors = new MetricExtractor[attrNames.size()];
-    int n = 0;
+    List<MetricExtractor> metricExtractors = new ArrayList<>();
     for (String attributeName : attrNames) {
       BeanAttributeExtractor attrExtractor = BeanAttributeExtractor.fromName(attributeName);
       // This is essentially the same as 'attributeName' but with escape characters removed
@@ -162,44 +164,85 @@ public MetricDef buildMetricDef() throws Exception {
         metricInfo = m.buildMetricInfo(prefix, niceAttributeName, getUnit(), getMetricType());
       }
 
-      List<MetricAttribute> attributeList;
       List<MetricAttribute> ownAttributes = getAttributeList();
-      if (ownAttributes != null && m != null && m.getAttributeList() != null) {
-        // MetricAttributes have been specified at two levels, need to combine them
-        attributeList = combineMetricAttributes(ownAttributes, m.getAttributeList());
-      } else if (ownAttributes != null) {
-        attributeList = ownAttributes;
-      } else if (m != null && m.getAttributeList() != null) {
-        // Get the attributes from the metric
-        attributeList = m.getAttributeList();
+      List<MetricAttribute> metricAttributes = m != null ? m.getAttributeList() : null;
+      List<MetricAttribute> attributeList =
+          combineMetricAttributes(ownAttributes, metricAttributes);
+
+      // higher priority to metric level mapping, then jmx rule as fallback
+      StateMapping stateMapping = getEffectiveStateMapping(m, this);
+
+      if (stateMapping.isEmpty()) {
+        MetricExtractor metricExtractor =
+            new MetricExtractor(
+                attrExtractor, metricInfo, attributeList.toArray(new MetricAttribute[0]));
+        metricExtractors.add(metricExtractor);
       } else {
-        // There are no attributes at all
-        attributeList = new ArrayList<MetricAttribute>();
-      }
 
-      MetricExtractor metricExtractor =
-          new MetricExtractor(
-              attrExtractor,
-              metricInfo,
-              attributeList.toArray(new MetricAttribute[attributeList.size()]));
-      metricExtractors[n++] = metricExtractor;
+        // generate one metric extractor per state metric key
+        // each metric extractor will have the state attribute replaced with a constant
+        for (String key : stateMapping.getStateKeys()) {
+          List<MetricAttribute> stateMetricAttributes =
+              attributeList.stream()
+                  .map(
+                      ma -> {
+                        if (!ma.isStateAttribute()) {
+                          return ma;
+                        } else {
+                          return new MetricAttribute(
+                              ma.getAttributeName(), MetricAttributeExtractor.fromConstant(key));
+                        }
+                      })
+                  .collect(Collectors.toList());
+
+          BeanAttributeExtractor stateMetricExtractor =
+              new BeanAttributeExtractor(attrExtractor.getAttributeName()) {
+
+                @Override
+                protected Number extractNumericalAttribute(
+                    MBeanServerConnection connection, ObjectName objectName) {
+                  String rawStateValue = attrExtractor.extractValue(connection, objectName);
+                  String mappedStateValue = stateMapping.getStateValue(rawStateValue);
+                  return key.equals(mappedStateValue) ? 1 : 0;
+                }
+              };
+
+          metricExtractors.add(
+              new MetricExtractor(
+                  stateMetricExtractor,
+                  metricInfo,
+                  stateMetricAttributes.toArray(new MetricAttribute[0])));
+        }
+      }
     }
 
-    return new MetricDef(group, metricExtractors);
+    return new MetricDef(group, metricExtractors.toArray(new MetricExtractor[0]));
   }
 
   private static List<MetricAttribute> combineMetricAttributes(
       List<MetricAttribute> ownAttributes, List<MetricAttribute> metricAttributes) {
+
     Map<String, MetricAttribute> set = new HashMap<>();
-    for (MetricAttribute ownAttribute : ownAttributes) {
-      set.put(ownAttribute.getAttributeName(), ownAttribute);
+    if (ownAttributes != null) {
+      for (MetricAttribute ownAttribute : ownAttributes) {
+        set.put(ownAttribute.getAttributeName(), ownAttribute);
+      }
     }
-
-    // Let the metric level defined attributes override own attributes
-    for (MetricAttribute metricAttribute : metricAttributes) {
-      set.put(metricAttribute.getAttributeName(), metricAttribute);
+    if (metricAttributes != null) {
+      // Let the metric level defined attributes override own attributes
+      for (MetricAttribute metricAttribute : metricAttributes) {
+        set.put(metricAttribute.getAttributeName(), metricAttribute);
+      }
     }
 
-    return new ArrayList<MetricAttribute>(set.values());
+    return new ArrayList<>(set.values());
+  }
+
+  private static StateMapping getEffectiveStateMapping(Metric m, JmxRule rule) {
+    if (m == null || m.getStateMapping().isEmpty()) {
+      return rule.getStateMapping();
+    } else {
+      return m.getStateMapping();
+    }
   }
 }

instrumentation/jmx-metrics/library/src/main/java/io/opentelemetry/instrumentation/jmx/yaml/MetricStructure.java

@@ -20,7 +20,7 @@
  * <li>the metric attributes
  * <li>the unit
  *
- *     <p>Known subclasses are JmxRule and Metric.
+ *     <p>Known subclasses are {@link JmxRule} and {@link Metric}.
  */
 abstract class MetricStructure {
 
@@ -30,14 +30,22 @@ abstract class MetricStructure {
   //      KEY1: SPECIFICATION1
   //      KEY2: SPECIFICATION2
   //    unit: UNIT
-
-  private Map<String, String> metricAttribute; // unused, for YAML parser only
+  //    stateMapping:
+  //      state1: [a,b]
+  //      state2: c
+  //      state3: '*'
+
+  private Map<String, String> metricAttribute;
+  private StateMapping stateMapping = StateMapping.empty();
+  private static final String STATE_MAPPING_WILDCARD = "*";
   private String unit;
 
   private MetricInfo.Type metricType;
   private List<MetricAttribute> metricAttributes;
 
-  public void setType(String t) {
+  MetricStructure() {}
+
+  void setType(String t) {
     // Do not complain about case variations
     t = t.trim().toUpperCase(Locale.ROOT);
     this.metricType = MetricInfo.Type.valueOf(t);
@@ -51,6 +59,36 @@ public void setUnit(String unit) {
     this.unit = validateUnit(unit.trim());
   }
 
+  void setStateMapping(Map<String, Object> stateMapping) {
+    StateMapping.Builder builder = StateMapping.builder();
+    for (Map.Entry<String, Object> entry : stateMapping.entrySet()) {
+      String stateKey = entry.getKey();
+      Object stateValue = entry.getValue();
+      if (stateValue instanceof String) {
+        addMappedValue(builder, (String) stateValue, stateKey);
+      } else if (stateValue instanceof List) {
+        for (Object listEntry : (List<?>) stateValue) {
+          if (!(listEntry instanceof String)) {
+            throw new IllegalArgumentException("unexpected state list value: " + stateKey);
+          }
+          addMappedValue(builder, (String) listEntry, stateKey);
+        }
+      } else {
+        throw new IllegalArgumentException("unexpected state value: " + stateValue);
+      }
+    }
+    this.stateMapping = builder.build();
+  }
+
+  private static void addMappedValue(
+      StateMapping.Builder builder, String stateValue, String stateKey) {
+    if (stateValue.equals(STATE_MAPPING_WILDCARD)) {
+      builder.withDefaultState(stateKey);
+    } else {
+      builder.withMappedValue(stateValue, stateKey);
+    }
+  }
+
   @CanIgnoreReturnValue
   private String validateUnit(String unit) {
     requireNonEmpty(unit, "Metric unit is empty");
@@ -67,9 +105,7 @@ private String validateUnit(String unit) {
   public void setMetricAttribute(Map<String, String> map) {
     this.metricAttribute = map;
     // pre-build the MetricAttributes
-    List<MetricAttribute> attrList = new ArrayList<>();
-    addMetricAttributes(attrList, map);
-    this.metricAttributes = attrList;
+    this.metricAttributes = addMetricAttributes(map);
   }
 
   // Used only for testing
@@ -91,47 +127,66 @@ protected void requireNonEmpty(String s, String msg) {
     }
   }
 
-  private static void addMetricAttributes(
-      List<MetricAttribute> list, Map<String, String> metricAttributeMap) {
-    if (metricAttributeMap != null) {
-      for (String key : metricAttributeMap.keySet()) {
-        String target = metricAttributeMap.get(key);
-        if (target == null) {
-          throw new IllegalStateException(
-              "nothing specified for metric attribute key '" + key + "'");
+  private List<MetricAttribute> addMetricAttributes(Map<String, String> metricAttributeMap) {
+
+    List<MetricAttribute> list = new ArrayList<>();
+    boolean hasStateAttribute = false;
+    for (String key : metricAttributeMap.keySet()) {
+      String target = metricAttributeMap.get(key);
+      if (target == null) {
+        throw new IllegalStateException("nothing specified for metric attribute key '" + key + "'");
+      }
+      MetricAttribute attribute = buildMetricAttribute(key, target.trim());
+      if (attribute.isStateAttribute()) {
+        if (hasStateAttribute) {
+          throw new IllegalArgumentException("only one state attribute is allowed");
         }
-        list.add(buildMetricAttribute(key, target.trim()));
+        hasStateAttribute = true;
       }
+      list.add(attribute);
+    }
+    if (hasStateAttribute && stateMapping.isEmpty()) {
+      throw new IllegalArgumentException("statekey() usage without stateMapping");
     }
+    return list;
   }
 
   private static MetricAttribute buildMetricAttribute(String key, String target) {
     // The recognized forms of target are:
     //  - param(STRING)
     //  - beanattr(STRING)
     //  - const(STRING)
+    //  - statekey()
     // where STRING is the name of the corresponding parameter key, attribute name,
     // or the direct value to use
     int k = target.indexOf(')');
 
     // Check for one of the cases as above
     if (target.startsWith("param(")) {
       if (k > 0) {
+        String jmxAttribute = target.substring(6, k).trim();
         return new MetricAttribute(
-            key, MetricAttributeExtractor.fromObjectNameParameter(target.substring(6, k).trim()));
+            key, MetricAttributeExtractor.fromObjectNameParameter(jmxAttribute));
       }
     } else if (target.startsWith("beanattr(")) {
       if (k > 0) {
-        return new MetricAttribute(
-            key, MetricAttributeExtractor.fromBeanAttribute(target.substring(9, k).trim()));
+        String jmxAttribute = target.substring(9, k).trim();
+        return new MetricAttribute(key, MetricAttributeExtractor.fromBeanAttribute(jmxAttribute));
       }
     } else if (target.startsWith("const(")) {
       if (k > 0) {
-        return new MetricAttribute(
-            key, MetricAttributeExtractor.fromConstant(target.substring(6, k).trim()));
+        String constantValue = target.substring(6, k).trim();
+        return new MetricAttribute(key, MetricAttributeExtractor.fromConstant(constantValue));
       }
+    } else if (target.equals("statekey()")) {
+      return new MetricAttribute(key, null);
     }
 
-    throw new IllegalArgumentException("Invalid metric attribute specification for '" + key + "'");
+    String msg = "Invalid metric attribute specification for '" + key + "': " + target;
+    throw new IllegalArgumentException(msg);
+  }
+
+  public StateMapping getStateMapping() {
+    return stateMapping;
   }
 }

instrumentation/jmx-metrics/library/src/main/java/io/opentelemetry/instrumentation/jmx/yaml/RuleParser.java

@@ -97,9 +97,13 @@ private static Map<String, Metric> parseMappings(@Nullable Map<String, Object> m
     Map<String, Metric> mappings = new LinkedHashMap<>();
     if (mappingYaml != null) {
       mappingYaml.forEach(
-          (name, metricYaml) ->
-              mappings.put(
-                  name, metricYaml == null ? null : parseMetric((Map<String, Object>) metricYaml)));
+          (name, metricYaml) -> {
+            Metric m = null;
+            if (metricYaml != null) {
+              m = parseMetric((Map<String, Object>) metricYaml);
+            }
+            mappings.put(name, m);
+          });
     }
     return mappings;
   }
@@ -124,10 +128,18 @@ private static Metric parseMetric(Map<String, Object> metricYaml) {
   @SuppressWarnings("unchecked")
   private static void parseMetricStructure(
       Map<String, Object> metricStructureYaml, MetricStructure out) {
+
     String type = (String) metricStructureYaml.remove("type");
     if (type != null) {
       out.setType(type);
     }
+
+    Map<String, Object> stateMapping =
+        (Map<String, Object>) metricStructureYaml.remove("stateMapping");
+    if (stateMapping != null) {
+      out.setStateMapping(stateMapping);
+    }
+
     Map<String, String> metricAttribute =
         (Map<String, String>) metricStructureYaml.remove("metricAttribute");
     if (metricAttribute != null) {