Fix client manager metrics and JMX lifecycle races

iotdb-core/metrics/core/pom.xml

@@ -46,6 +46,11 @@
             <groupId>org.slf4j</groupId>
             <artifactId>slf4j-api</artifactId>
         </dependency>
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

iotdb-core/metrics/core/src/main/java/org/apache/iotdb/metrics/core/reporter/IoTDBJmxReporter.java

@@ -40,8 +40,9 @@
 import javax.management.ObjectName;
 
 import java.lang.management.ManagementFactory;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 public class IoTDBJmxReporter implements JmxReporter {
   private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBJmxReporter.class);
@@ -55,37 +56,38 @@ public class IoTDBJmxReporter implements JmxReporter {
   /** The objectNameFactory used to create objectName for metrics */
   private final ObjectNameFactory objectNameFactory;
 
-  /** The map that stores all registered metrics */
-  private final Map<ObjectName, ObjectName> registered;
+  /** Registrations owned by this reporter, guarded by the map's monitor. */
+  private final Map<ObjectName, Registration> registered;
+
+  private boolean started;
 
   /** The JMX MBeanServer */
   private final MBeanServer mBeanServer;
 
-  private void registerMBean(Object mBean, ObjectName objectName) throws JMException {
-    if (!mBeanServer.isRegistered(objectName)) {
-      ObjectInstance objectInstance = mBeanServer.registerMBean(mBean, objectName);
-      if (objectInstance != null) {
-        // the websphere mbeanserver rewrites the objectname to include
-        // cell, node & server info
-        // make sure we capture the new objectName for unregistration
-        registered.put(objectName, objectInstance.getObjectName());
-      } else {
-        registered.put(objectName, objectName);
+  private void registerMBean(IMetric metric, ObjectName objectName) throws JMException {
+    Registration previous = registered.get(objectName);
+    if (previous != null) {
+      if (previous.metric == metric && mBeanServer.isRegistered(previous.actualName)) {
+        return;
       }
+      unregisterMBean(previous);
+      registered.remove(objectName);
+    }
+    if (!mBeanServer.isRegistered(objectName)) {
+      ObjectInstance objectInstance = mBeanServer.registerMBean(metric, objectName);
+      // Some MBean servers rewrite ObjectNames. Keep the actual name together with its owner.
+      registered.put(
+          objectName,
+          new Registration(
+              metric, objectInstance == null ? objectName : objectInstance.getObjectName()));
     }
   }
 
-  private void unregisterMBean(ObjectName originalObjectName)
-      throws InstanceNotFoundException, MBeanRegistrationException {
-    ObjectName storedObjectName = registered.remove(originalObjectName);
-    if (storedObjectName != null) {
-      if (mBeanServer.isRegistered(storedObjectName)) {
-        mBeanServer.unregisterMBean(storedObjectName);
-      }
-    } else {
-      if (mBeanServer.isRegistered(originalObjectName)) {
-        mBeanServer.unregisterMBean(originalObjectName);
-      }
+  private void unregisterMBean(Registration registration) throws MBeanRegistrationException {
+    try {
+      mBeanServer.unregisterMBean(registration.actualName);
+    } catch (InstanceNotFoundException ignored) {
+      // An externally removed MBean is already unregistered.
     }
   }
 
@@ -94,59 +96,81 @@ public void registerMetric(IMetric metric, MetricInfo metricInfo) {
     String metricName = metric.getClass().getSimpleName();
     try {
       final ObjectName objectName = createName(metricName, metricInfo);
-      metric.setObjectName(objectName);
-      registerMBean(metric, objectName);
+      synchronized (registered) {
+        // Ignore callbacks from a stopped reporter or a superseded registry entry.
+        if (!started || metricManager.getAllMetrics().get(metricInfo) != metric) {
+          return;
+        }
+        metric.setObjectName(objectName);
+        registerMBean(metric, objectName);
+      }
     } catch (Exception e) {
       LOGGER.warn(MetricsCoreMessages.JMX_REGISTER_FAILED + metricName, e);
     }
   }
 
   @Override
   public void unregisterMetric(IMetric metric, MetricInfo metricInfo) {
+    if (metric == null) {
+      return;
+    }
     String metricName = metric.getClass().getSimpleName();
     try {
       final ObjectName objectName = createName(metricName, metricInfo);
-      unregisterMBean(objectName);
-    } catch (InstanceNotFoundException e) {
-      LOGGER.debug(MetricsCoreMessages.JMX_UNREGISTER_FAILED, e);
+      synchronized (registered) {
+        Registration registration = registered.get(objectName);
+        // A delayed callback for an old metric must not delete its replacement.
+        if (registration != null && registration.metric == metric) {
+          unregisterMBean(registration);
+          registered.remove(objectName);
+        }
+      }
     } catch (MBeanRegistrationException e) {
       LOGGER.warn(MetricsCoreMessages.JMX_UNREGISTER_FAILED, e);
     }
   }
 
   private ObjectName createName(String type, MetricInfo metricInfo) {
     String name = metricInfo.getName();
-    return objectNameFactory.createName(type, DOMAIN, name);
+    return objectNameFactory.createName(type, DOMAIN, name, metricInfo.getTags());
   }
 
-  void unregisterAll() throws InstanceNotFoundException, MBeanRegistrationException {
-    for (ObjectName name : registered.keySet()) {
-      unregisterMBean(name);
+  void unregisterAll() throws MBeanRegistrationException {
+    synchronized (registered) {
+      Iterator<Registration> iterator = registered.values().iterator();
+      while (iterator.hasNext()) {
+        unregisterMBean(iterator.next());
+        iterator.remove();
+      }
     }
-    // clear registered
-    registered.clear();
   }
 
-  private IoTDBJmxReporter(
+  IoTDBJmxReporter(
       AbstractMetricManager metricManager,
       MBeanServer mBeanServer,
       ObjectNameFactory objectNameFactory) {
     this.metricManager = metricManager;
     this.mBeanServer = mBeanServer;
     this.objectNameFactory = objectNameFactory;
-    this.registered = new ConcurrentHashMap<>();
+    this.registered = new HashMap<>();
   }
 
   @Override
   public boolean start() {
     try {
-      if (!registered.isEmpty()) {
+      boolean alreadyStarted;
+      synchronized (registered) {
+        alreadyStarted = started;
+        started = true;
+      }
+      if (alreadyStarted) {
         LOGGER.warn(MetricsCoreMessages.JMX_REPORTER_ALREADY_START);
         return false;
       }
       // register all existed metrics into JmxReporter
       metricManager.getAllMetrics().forEach((key, value) -> registerMetric(value, key));
     } catch (Exception e) {
+      stop();
       LOGGER.warn(MetricsCoreMessages.JMX_REPORTER_START_FAILED, e);
       return false;
     }
@@ -157,7 +181,10 @@ public boolean start() {
   @Override
   public boolean stop() {
     try {
-      unregisterAll();
+      synchronized (registered) {
+        started = false;
+        unregisterAll();
+      }
     } catch (Exception e) {
       LOGGER.warn(MetricsCoreMessages.JMX_REPORTER_STOP_FAILED, e);
       return false;
@@ -171,6 +198,16 @@ public ReporterType getReporterType() {
     return ReporterType.JMX;
   }
 
+  private static class Registration {
+    private final IMetric metric;
+    private final ObjectName actualName;
+
+    private Registration(IMetric metric, ObjectName actualName) {
+      this.metric = metric;
+      this.actualName = actualName;
+    }
+  }
+
   private static class IoTDBJmxReporterHolder {
     private static final IoTDBJmxReporter INSTANCE =
         new IoTDBJmxReporter(

iotdb-core/metrics/core/src/main/java/org/apache/iotdb/metrics/core/type/IoTDBAutoGauge.java

@@ -37,9 +37,10 @@ public IoTDBAutoGauge(T object, ToDoubleFunction<T> mapper) {
 
   @Override
   public double getValue() {
-    if (refObject.get() == null) {
+    T object = refObject.get();
+    if (object == null) {
       return 0d;
     }
-    return mapper.applyAsDouble(refObject.get());
+    return mapper.applyAsDouble(object);
   }
 }

iotdb-core/metrics/core/src/main/java/org/apache/iotdb/metrics/core/utils/IoTDBMetricObjNameFactory.java

@@ -30,7 +30,7 @@
 import java.util.Hashtable;
 
 public class IoTDBMetricObjNameFactory implements ObjectNameFactory {
-  private static final char[] QUOTABLE_CHARS = new char[] {',', '=', ':', '"'};
+  private static final char[] QUOTABLE_CHARS = new char[] {',', '=', ':', '"', '\n', '*', '?'};
   private static final Logger LOGGER = LoggerFactory.getLogger(IoTDBMetricObjNameFactory.class);
 
   private IoTDBMetricObjNameFactory() {
@@ -40,38 +40,14 @@ private IoTDBMetricObjNameFactory() {
   @Override
   public ObjectName createName(String type, String domain, String name) {
     try {
-      ObjectName objectName;
       Hashtable<String, String> properties = new Hashtable<>();
-
-      properties.put("name", name);
-      properties.put("type", type);
-      objectName = new ObjectName(domain, properties);
-
-      /*
-       * The only way we can find out if we need to quote the properties is by
-       * checking an ObjectName that we've constructed.
-       */
-      if (objectName.isDomainPattern()) {
-        domain = ObjectName.quote(domain);
-      }
-      if (objectName.isPropertyValuePattern("name")
-          || shouldQuote(objectName.getKeyProperty("name"))) {
-        properties.put("name", ObjectName.quote(name));
-      }
-      if (objectName.isPropertyValuePattern("type")
-          || shouldQuote(objectName.getKeyProperty("type"))) {
-        properties.put("type", ObjectName.quote(type));
-      }
-      objectName = new ObjectName(domain, properties);
-
-      return objectName;
+      // Quote before constructing the name; falling back to a name-only MBean loses its type.
+      properties.put("name", shouldQuote(name) ? ObjectName.quote(name) : name);
+      properties.put("type", shouldQuote(type) ? ObjectName.quote(type) : type);
+      return new ObjectName(domain, properties);
     } catch (MalformedObjectNameException e) {
-      try {
-        return new ObjectName(domain, "name", ObjectName.quote(name));
-      } catch (MalformedObjectNameException e1) {
-        LOGGER.warn(MetricsCoreMessages.JMX_UNABLE_TO_REGISTER, type, name, e1);
-        throw new RuntimeException(e1);
-      }
+      LOGGER.warn(MetricsCoreMessages.JMX_UNABLE_TO_REGISTER, type, name, e);
+      throw new IllegalArgumentException(e);
     }
   }
 

iotdb-core/metrics/core/src/main/java/org/apache/iotdb/metrics/core/utils/ObjectNameFactory.java

@@ -19,8 +19,12 @@
 
 package org.apache.iotdb.metrics.core.utils;
 
+import javax.management.MalformedObjectNameException;
 import javax.management.ObjectName;
 
+import java.util.Hashtable;
+import java.util.Map;
+
 public interface ObjectNameFactory {
   /**
    * Create objectName for a certain metric.
@@ -31,4 +35,45 @@ public interface ObjectNameFactory {
    * @return metric's objectName
    */
   public ObjectName createName(String type, String domain, String name);
+
+  /**
+   * Include all metric tags in the MBean identity. Tag keys use a separate namespace and reversible
+   * escaping so they cannot overwrite the metric name/type or collide after sanitization. Tag
+   * values are quoted to preserve literal wildcard characters. Untagged metrics keep their existing
+   * names.
+   */
+  default ObjectName createName(String type, String domain, String name, Map<String, String> tags) {
+    ObjectName base = createName(type, domain, name);
+    if (tags.isEmpty()) {
+      return base;
+    }
+    Hashtable<String, String> properties = base.getKeyPropertyList();
+    tags.forEach((key, value) -> properties.put(encodeTagKey(key), ObjectName.quote(value)));
+    try {
+      return new ObjectName(base.getDomain(), properties);
+    } catch (MalformedObjectNameException e) {
+      throw new IllegalArgumentException(e);
+    }
+  }
+
+  private static String encodeTagKey(String key) {
+    StringBuilder encoded = new StringBuilder("tag.");
+    for (int i = 0; i < key.length(); i++) {
+      char character = key.charAt(i);
+      if ((character >= 'a' && character <= 'z')
+          || (character >= 'A' && character <= 'Z')
+          || (character >= '0' && character <= '9')
+          || character == '_'
+          || character == '-'
+          || character == '.') {
+        encoded.append(character);
+      } else {
+        encoded.append('%');
+        for (int shift = 12; shift >= 0; shift -= 4) {
+          encoded.append(Character.forDigit((character >> shift) & 0xf, 16));
+        }
+      }
+    }
+    return encoded.toString();
+  }
 }