Fix a problem: benchmark won't quit when DataNode accidently shutdown (…

…#384)
thulab · Dec 15, 2023 · 410b1c2 · 410b1c2
1 parent 9e8b915
commit 410b1c2
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 35 deletions.
diff --git a/.../src/main/java/cn/edu/tsinghua/iot/benchmark/measurement/persistence/csv/CSVRecorder.java b/.../src/main/java/cn/edu/tsinghua/iot/benchmark/measurement/persistence/csv/CSVRecorder.java
@@ -25,6 +25,7 @@
 import cn.edu.tsinghua.iot.benchmark.measurement.enums.SystemMetrics;
 import cn.edu.tsinghua.iot.benchmark.measurement.persistence.TestDataPersistence;
 import cn.edu.tsinghua.iot.benchmark.mode.enums.BenchmarkMode;
+import cn.edu.tsinghua.iot.benchmark.utils.NamedThreadFactory;
 import cn.edu.tsinghua.iot.benchmark.utils.ZipUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -141,7 +142,7 @@ public CSVRecorder() {
         LOGGER.error("", ioException);
       }
     }
-    service = Executors.newSingleThreadExecutor();
+    service = Executors.newSingleThreadExecutor(new NamedThreadFactory("CSVRecorder"));
   }
 
   /** write header of csv file */

diff --git a/iotdb-1.1/src/main/java/cn/edu/tsinghua/iot/benchmark/iotdb110/IoTDBClusterSession.java b/iotdb-1.1/src/main/java/cn/edu/tsinghua/iot/benchmark/iotdb110/IoTDBClusterSession.java
@@ -30,6 +30,7 @@
 
 import cn.edu.tsinghua.iot.benchmark.tsdb.DBConfig;
 import cn.edu.tsinghua.iot.benchmark.tsdb.TsdbException;
+import cn.edu.tsinghua.iot.benchmark.utils.NamedThreadFactory;
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
@@ -189,22 +190,7 @@ public IoTDBClusterSession(DBConfig dbConfig) {
   @Override
   public void init() throws TsdbException {
     // do nothing
-    this.service = Executors.newSingleThreadExecutor();
-  }
-
-  @Override
-  public void close() throws TsdbException {
-    if (sessionWrapper != null) {
-      try {
-        sessionWrapper.close();
-      } catch (IoTDBConnectionException ignored) {
-        // should never happen
-      }
-    }
-    if (ioTDBConnection != null) {
-      ioTDBConnection.close();
-    }
-    this.service.shutdown();
+    this.service = Executors.newSingleThreadExecutor(new NamedThreadFactory("ClusterSession"));
   }
 
   @Override

diff --git a/iotdb-1.1/src/main/java/cn/edu/tsinghua/iot/benchmark/iotdb110/IoTDBSession.java b/iotdb-1.1/src/main/java/cn/edu/tsinghua/iot/benchmark/iotdb110/IoTDBSession.java
@@ -29,7 +29,7 @@
 import org.apache.iotdb.tsfile.write.record.Tablet;
 
 import cn.edu.tsinghua.iot.benchmark.tsdb.DBConfig;
-import cn.edu.tsinghua.iot.benchmark.tsdb.TsdbException;
+import cn.edu.tsinghua.iot.benchmark.utils.NamedThreadFactory;
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
@@ -183,7 +183,7 @@ public void init() {
       } else {
         sessionWrapper.open();
       }
-      this.service = Executors.newSingleThreadExecutor();
+      this.service = Executors.newSingleThreadExecutor(new NamedThreadFactory("Session"));
     } catch (IoTDBConnectionException e) {
       LOGGER.error("Failed to add session", e);
     }
@@ -208,20 +208,4 @@ public void cleanup() {
       LOGGER.error("Failed to execute statement:" + e.getMessage());
     }
   }
-
-  @Override
-  public void close() throws TsdbException {
-    try {
-      if (sessionWrapper != null) {
-        sessionWrapper.close();
-      }
-      if (ioTDBConnection != null) {
-        ioTDBConnection.close();
-      }
-      this.service.shutdown();
-    } catch (IoTDBConnectionException ioTDBConnectionException) {
-      LOGGER.error("Failed to close session.");
-      throw new TsdbException(ioTDBConnectionException);
-    }
-  }
 }
diff --git a/iotdb-1.1/src/main/java/cn/edu/tsinghua/iot/benchmark/iotdb110/IoTDBSessionBase.java b/iotdb-1.1/src/main/java/cn/edu/tsinghua/iot/benchmark/iotdb110/IoTDBSessionBase.java
@@ -497,4 +497,20 @@ Status waitWriteTaskToFinishAndGetStatus() {
     }
     return new Status(true);
   }
+
+  @Override
+  public void close() throws TsdbException {
+    try {
+      if (sessionWrapper != null) {
+        sessionWrapper.close();
+      }
+      if (ioTDBConnection != null) {
+        ioTDBConnection.close();
+      }
+    } catch (IoTDBConnectionException e) {
+      throw new TsdbException(e);
+    } finally {
+      this.service.shutdown();
+    }
+  }
 }