Fix ThreadPoolMetrics NPE (#14404)

Signed-off-by: OneSizeFitQuorum <[email protected]>
apache · Dec 12, 2024 · 8144e5c · 8144e5c
1 parent 4be9909
commit 8144e5c
Showing 1 changed file with 77 additions and 62 deletions.
diff --git a/...ore/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadPoolMetrics.java b/...ore/node-commons/src/main/java/org/apache/iotdb/commons/concurrent/ThreadPoolMetrics.java
@@ -48,87 +48,102 @@ public synchronized void registerThreadPool(IThreadPoolMBean pool, String name)
       notRegisteredPoolMap.put(name, pool);
     } else {
       registeredPoolMap.put(name, pool);
-      metricService.createAutoGauge(
-          SystemMetric.THREAD_POOL_ACTIVE_THREAD_COUNT.toString(),
-          MetricLevel.IMPORTANT,
-          registeredPoolMap,
-          map -> registeredPoolMap.get(name).getActiveCount(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.createAutoGauge(
-          SystemMetric.THREAD_POOL_CORE_SIZE.toString(),
-          MetricLevel.IMPORTANT,
-          registeredPoolMap,
-          map -> registeredPoolMap.get(name).getCorePoolSize(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.createAutoGauge(
-          SystemMetric.THREAD_POOL_WAITING_TASK_COUNT.toString(),
-          MetricLevel.IMPORTANT,
-          registeredPoolMap,
-          map -> registeredPoolMap.get(name).getQueueLength(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.createAutoGauge(
-          SystemMetric.THREAD_POOL_DONE_TASK_COUNT.toString(),
-          MetricLevel.IMPORTANT,
-          registeredPoolMap,
-          map -> registeredPoolMap.get(name).getCompletedTaskCount(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.createAutoGauge(
-          SystemMetric.THREAD_POOL_LARGEST_POOL_SIZE.toString(),
-          MetricLevel.IMPORTANT,
-          registeredPoolMap,
-          map -> registeredPoolMap.get(name).getLargestPoolSize(),
-          SystemTag.POOL_NAME.toString(),
-          name);
+      registerThreadPoolMetrics(name);
     }
   }
 
+  private void registerThreadPoolMetrics(String name) {
+    metricService.createAutoGauge(
+        SystemMetric.THREAD_POOL_ACTIVE_THREAD_COUNT.toString(),
+        MetricLevel.IMPORTANT,
+        registeredPoolMap,
+        map -> registeredPoolMap.get(name).getActiveCount(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.createAutoGauge(
+        SystemMetric.THREAD_POOL_CORE_SIZE.toString(),
+        MetricLevel.IMPORTANT,
+        registeredPoolMap,
+        map -> registeredPoolMap.get(name).getCorePoolSize(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.createAutoGauge(
+        SystemMetric.THREAD_POOL_WAITING_TASK_COUNT.toString(),
+        MetricLevel.IMPORTANT,
+        registeredPoolMap,
+        map -> registeredPoolMap.get(name).getQueueLength(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.createAutoGauge(
+        SystemMetric.THREAD_POOL_DONE_TASK_COUNT.toString(),
+        MetricLevel.IMPORTANT,
+        registeredPoolMap,
+        map -> registeredPoolMap.get(name).getCompletedTaskCount(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.createAutoGauge(
+        SystemMetric.THREAD_POOL_LARGEST_POOL_SIZE.toString(),
+        MetricLevel.IMPORTANT,
+        registeredPoolMap,
+        map -> registeredPoolMap.get(name).getLargestPoolSize(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+  }
+
   public synchronized void unRegisterThreadPool(String name) {
     if (metricService == null) {
       notRegisteredPoolMap.remove(name);
     } else {
       registeredPoolMap.remove(name);
-      metricService.remove(
-          MetricType.GAUGE,
-          SystemMetric.THREAD_POOL_ACTIVE_THREAD_COUNT.toString(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.remove(
-          MetricType.GAUGE,
-          SystemMetric.THREAD_POOL_CORE_SIZE.toString(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.remove(
-          MetricType.GAUGE,
-          SystemMetric.THREAD_POOL_WAITING_TASK_COUNT.toString(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.remove(
-          MetricType.GAUGE,
-          SystemMetric.THREAD_POOL_DONE_TASK_COUNT.toString(),
-          SystemTag.POOL_NAME.toString(),
-          name);
-      metricService.remove(
-          MetricType.GAUGE,
-          SystemMetric.THREAD_POOL_LARGEST_POOL_SIZE.toString(),
-          SystemTag.POOL_NAME.toString(),
-          name);
+      unRegisterThreadPoolMetrics(name);
     }
   }
 
+  private void unRegisterThreadPoolMetrics(String name) {
+    metricService.remove(
+        MetricType.GAUGE,
+        SystemMetric.THREAD_POOL_ACTIVE_THREAD_COUNT.toString(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.remove(
+        MetricType.GAUGE,
+        SystemMetric.THREAD_POOL_CORE_SIZE.toString(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.remove(
+        MetricType.GAUGE,
+        SystemMetric.THREAD_POOL_WAITING_TASK_COUNT.toString(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.remove(
+        MetricType.GAUGE,
+        SystemMetric.THREAD_POOL_DONE_TASK_COUNT.toString(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+    metricService.remove(
+        MetricType.GAUGE,
+        SystemMetric.THREAD_POOL_LARGEST_POOL_SIZE.toString(),
+        SystemTag.POOL_NAME.toString(),
+        name);
+  }
+
   @Override
   public synchronized void bindTo(AbstractMetricService metricService) {
     this.metricService = metricService;
-    notRegisteredPoolMap.forEach((name, pool) -> registerThreadPool(pool, name));
+    for (Map.Entry<String, IThreadPoolMBean> entry : notRegisteredPoolMap.entrySet()) {
+      registerThreadPoolMetrics(entry.getKey());
+    }
+    registeredPoolMap.putAll(notRegisteredPoolMap);
     notRegisteredPoolMap.clear();
   }
 
   @Override
   public synchronized void unbindFrom(AbstractMetricService metricService) {
-    registeredPoolMap.forEach((name, pool) -> unRegisterThreadPool(name));
+    for (Map.Entry<String, IThreadPoolMBean> entry : registeredPoolMap.entrySet()) {
+      unRegisterThreadPoolMetrics(entry.getKey());
+    }
+    registeredPoolMap.clear();
+    notRegisteredPoolMap.clear();
   }
 
   private static class ThreadPoolMetricsHolder {