fix state metric detection

Author: SylvainJuge

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

@@ -151,7 +151,7 @@ AttributeInfo getAttributeInfo(MBeanServerConnection connection, ObjectName obje
 
           // Verify correctness of configuration by attempting to extract the metric value.
           // The value will be discarded, but its type will be checked.
-          Object sampleValue = extractAttributeValue(connection, objectName, logger);
+          Object sampleValue = getSampleValue(connection, objectName);
 
           // Only numbers can be used to generate metric values
           if (sampleValue instanceof Number) {
@@ -194,6 +194,11 @@ AttributeInfo getAttributeInfo(MBeanServerConnection connection, ObjectName obje
     return null;
   }
 
+  @Nullable
+  protected Object getSampleValue(MBeanServerConnection connection, ObjectName objectName) {
+    return extractAttributeValue(connection, objectName, logger);
+  }
+
   /**
    * Extracts the specified attribute value. In case the value is a CompositeData, drills down into
    * it to find the correct singleton value (usually a Number or a String).
@@ -203,7 +208,7 @@ AttributeInfo getAttributeInfo(MBeanServerConnection connection, ObjectName obje
    *     pattern
    * @param logger the logger to use, may be null. Typically we want to log any issues with the
    *     attributes during MBean discovery, but once the attribute is successfully detected and
-   *     confirmed to be eligble for metric evaluation, any further attribute extraction
+   *     confirmed to be eligible for metric evaluation, any further attribute extraction
    *     malfunctions will be silent to avoid flooding the log.
    * @return the attribute value, if found, or {@literal null} if an error occurred
    */

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

@@ -18,7 +18,6 @@
 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;
@@ -173,52 +172,71 @@ public MetricDef buildMetricDef() throws Exception {
       StateMapping stateMapping = getEffectiveStateMapping(m, this);
 
       if (stateMapping.isEmpty()) {
-        MetricExtractor metricExtractor =
+        metricExtractors.add(
             new MetricExtractor(
-                attrExtractor, metricInfo, attributeList.toArray(new MetricAttribute[0]));
-        metricExtractors.add(metricExtractor);
+                attrExtractor, metricInfo, attributeList.toArray(new MetricAttribute[0])));
       } else {
 
         // 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])));
-        }
+        metricExtractors.addAll(
+            createStateMappingExtractors(stateMapping, attributeList, attrExtractor, metricInfo));
       }
     }
 
     return new MetricDef(group, metricExtractors.toArray(new MetricExtractor[0]));
   }
 
+  private static List<MetricExtractor> createStateMappingExtractors(
+      StateMapping stateMapping,
+      List<MetricAttribute> attributeList,
+      BeanAttributeExtractor attrExtractor,
+      MetricInfo metricInfo) {
+
+    List<MetricExtractor> extractors = new ArrayList<>();
+    for (String key : stateMapping.getStateKeys()) {
+      List<MetricAttribute> stateMetricAttributes = new ArrayList<>();
+
+      for (MetricAttribute ma : attributeList) {
+        // replace state metric attribute with constant key value
+        if (!ma.isStateAttribute()) {
+          stateMetricAttributes.add(ma);
+        } else {
+          stateMetricAttributes.add(
+              new MetricAttribute(
+                  ma.getAttributeName(), MetricAttributeExtractor.fromConstant(key)));
+        }
+      }
+
+      BeanAttributeExtractor stateMetricExtractor =
+          new BeanAttributeExtractor(attrExtractor.getAttributeName()) {
+
+            @Override
+            protected Object getSampleValue(
+                MBeanServerConnection connection, ObjectName objectName) {
+              // metric actual type is sampled in the discovery process, so we have to
+              // make this extractor as extracting integers.
+              return 0;
+            }
+
+            @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;
+            }
+          };
+
+      extractors.add(
+          new MetricExtractor(
+              stateMetricExtractor,
+              metricInfo,
+              stateMetricAttributes.toArray(new MetricAttribute[0])));
+    }
+    return extractors;
+  }
+
   private static List<MetricAttribute> combineMetricAttributes(
       List<MetricAttribute> ownAttributes, List<MetricAttribute> metricAttributes) {
 

instrumentation/jmx-metrics/library/src/test/java/io/opentelemetry/instrumentation/jmx/engine/RuleParserTest.java

@@ -10,6 +10,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.entry;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import io.opentelemetry.instrumentation.jmx.yaml.JmxConfig;
@@ -21,12 +22,13 @@
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Map;
+import javax.management.MBeanAttributeInfo;
+import javax.management.MBeanInfo;
 import javax.management.MBeanServerConnection;
 import javax.management.ObjectName;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
-import org.mockito.Mockito;
 
 class RuleParserTest {
   private static RuleParser parser;
@@ -399,9 +401,17 @@ void testStateMetricConf() throws Exception {
     assertThat(metric.getMetricAttribute()).containsEntry("state", "statekey()");
 
     ObjectName objectName = new ObjectName(jmxRule.getBean());
-    MBeanServerConnection mockConnection = Mockito.mock(MBeanServerConnection.class);
+    MBeanServerConnection mockConnection = mock(MBeanServerConnection.class);
+
+    // mock attribute value
     when(mockConnection.getAttribute(objectName, "jmxStateAttribute")).thenReturn("STOPPED");
 
+    // mock attribute discovery
+    MBeanInfo mockBeanInfo = mock(MBeanInfo.class);
+    when(mockBeanInfo.getAttributes()).thenReturn(new MBeanAttributeInfo[] {
+        new MBeanAttributeInfo("jmxStateAttribute", "java.lang.String", "", true, false, false)});
+    when(mockConnection.getMBeanInfo(objectName)).thenReturn(mockBeanInfo);
+
     MetricDef metricDef = jmxRule.buildMetricDef();
     assertThat(metricDef.getMetricExtractors())
         .hasSize(3)
@@ -418,6 +428,14 @@ void testStateMetricConf() throws Exception {
 
               BeanAttributeExtractor attributeExtractor = me.getMetricValueExtractor();
               assertThat(attributeExtractor).isNotNull();
+              assertThat(attributeExtractor.getSampleValue(null, null))
+                  .describedAs("sampled value must be an integer")
+                  .isInstanceOf(Integer.class);
+
+              assertThat(attributeExtractor.getAttributeInfo(mockConnection, objectName))
+                  .describedAs("attribute info must be provided as a regular int metric")
+                  .isNotNull();
+
               int expectedValue = stateAttributeValue.equals("failed") ? 1 : 0;
               Number extractedValue =
                   attributeExtractor.extractNumericalAttribute(mockConnection, objectName);