[To dev/1.3] Repair lost data partitions by providing a new SQL statement

integration-test/src/test/java/org/apache/iotdb/confignode/it/partition/DataPartitionTableIntegrityCheckProcedureIT.java

@@ -0,0 +1,167 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.confignode.it.partition;
+
+import org.apache.iotdb.commons.enums.RepairDataPartitionTableProgressState;
+import org.apache.iotdb.it.env.EnvFactory;
+import org.apache.iotdb.it.framework.IoTDBTestRunner;
+import org.apache.iotdb.itbase.category.ClusterIT;
+import org.apache.iotdb.itbase.category.LocalStandaloneIT;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.iotdb.consensus.ConsensusFactory.RATIS_CONSENSUS;
+
+@RunWith(IoTDBTestRunner.class)
+@Category({LocalStandaloneIT.class, ClusterIT.class})
+public class DataPartitionTableIntegrityCheckProcedureIT {
+  private static final Logger LOGGER =
+      LoggerFactory.getLogger(DataPartitionTableIntegrityCheckProcedureIT.class);
+
+  @Before
+  public void setUp() {
+    EnvFactory.getEnv()
+        .getConfig()
+        .getCommonConfig()
+        .setConfigNodeConsensusProtocolClass(RATIS_CONSENSUS)
+        .setSchemaRegionConsensusProtocolClass(RATIS_CONSENSUS)
+        .setDataRegionConsensusProtocolClass(RATIS_CONSENSUS)
+        .setDataReplicationFactor(1);
+    EnvFactory.getEnv().initClusterEnvironment(1, 1);
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    EnvFactory.getEnv().cleanClusterEnvironment();
+  }
+
+  @Test
+  public void testConcurrentSubmitDataPartitionTableIntegrityCheckProcedure()
+      throws InterruptedException {
+    final int threadCount = 10;
+    final CountDownLatch startLatch = new CountDownLatch(1);
+    final CountDownLatch finishLatch = new CountDownLatch(threadCount);
+    final ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+
+    final AtomicInteger successCount = new AtomicInteger(0);
+    final AtomicInteger failCount = new AtomicInteger(0);
+    final List<String> failureMessages = Collections.synchronizedList(new ArrayList<>());
+
+    // Concurrently submit the DataPartitionTableIntegrityCheckProcedure
+    for (int i = 0; i < threadCount; i++) {
+      final int threadId = i;
+      executor.submit(
+          () -> {
+            try {
+              startLatch.await();
+
+              try (final Connection connection = EnvFactory.getEnv().getConnection();
+                  final Statement stmt = connection.createStatement()) {
+                stmt.execute("REPAIR DATA PARTITION TABLE");
+                successCount.incrementAndGet();
+                LOGGER.info("Thread {} submitted integrity check successfully", threadId);
+              }
+            } catch (final SQLException e) {
+              failCount.incrementAndGet();
+              failureMessages.add("Thread " + threadId + " failed: " + e.getMessage());
+              LOGGER.info(
+                  "Thread {} failed to submit integrity check: {}", threadId, e.getMessage());
+            } catch (final Exception e) {
+              failCount.incrementAndGet();
+              failureMessages.add("Thread " + threadId + " failed unexpectedly: " + e.getMessage());
+              LOGGER.error("Thread {} unexpected error: {}", threadId, e.getMessage(), e);
+            } finally {
+              finishLatch.countDown();
+            }
+          });
+    }
+
+    startLatch.countDown();
+
+    final boolean completed = finishLatch.await(60, TimeUnit.SECONDS);
+    Assert.assertTrue("Not all threads completed within timeout", completed);
+
+    executor.shutdown();
+    Assert.assertTrue(
+        "Executor did not terminate", executor.awaitTermination(10, TimeUnit.SECONDS));
+
+    LOGGER.info("Success count: {}, Fail count: {}", successCount.get(), failCount.get());
+    LOGGER.info("Failure messages: {}", failureMessages);
+
+    Assert.assertEquals(
+        "Only one procedure should be submitted successfully", 1, successCount.get());
+    Assert.assertEquals(
+        "The other concurrent submissions should be rejected", threadCount - 1, failCount.get());
+  }
+
+  @Test
+  public void testShowRepairDataPartitionTableProgress() throws Exception {
+    try (final Connection connection = EnvFactory.getEnv().getConnection();
+        final Statement statement = connection.createStatement()) {
+      assertRepairProgress(statement, RepairDataPartitionTableProgressState.IDLE.name(), 0.0, 0.0);
+
+      statement.execute("REPAIR DATA PARTITION TABLE");
+      assertRepairProgress(statement, null, 0.0, 100.0);
+    }
+  }
+
+  private static void assertRepairProgress(
+      final Statement statement,
+      final String expectedStatus,
+      final double minProgress,
+      final double maxProgress)
+      throws SQLException {
+    try (final ResultSet resultSet =
+        statement.executeQuery("SHOW REPAIR DATA PARTITION TABLE PROGRESS")) {
+      Assert.assertTrue(resultSet.next());
+      if (expectedStatus != null) {
+        Assert.assertEquals(expectedStatus, resultSet.getString("Status"));
+      } else {
+        Assert.assertNotEquals(
+            RepairDataPartitionTableProgressState.UNKNOWN.name(), resultSet.getString("Status"));
+      }
+      final double progress = resultSet.getDouble("Progress(%)");
+      Assert.assertTrue(progress >= minProgress);
+      Assert.assertTrue(progress <= maxProgress);
+      Assert.assertNotNull(resultSet.getString("Message"));
+      Assert.assertFalse(resultSet.next());
+    }
+  }
+}

iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4

@@ -176,6 +176,7 @@ keyWords
     | PRIVILEGES
     | PRIVILEGE_VALUE
     | PROCESSLIST
+    | PROGRESS
     | PROCESSOR
     | PROPERTY
     | PRUNE
@@ -225,6 +226,7 @@ keyWords
     | SUBSCRIPTIONS
     | SUBSTRING
     | SYSTEM
+    | TABLE
     | TAGS
     | TAIL
     | TASK

iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4

@@ -89,6 +89,7 @@ utilityStatement
     | showQueries | showCurrentTimestamp | killQuery | grantWatermarkEmbedding
     | revokeWatermarkEmbedding | loadConfiguration | loadTimeseries | loadFile
     | removeFile | unloadFile
+    | repairDataPartitionTable | showRepairDataPartitionTableProgress
     ;
 
 /**
@@ -1088,6 +1089,16 @@ stopRepairData
     : STOP REPAIR DATA (ON (LOCAL | CLUSTER))?
     ;
 
+// Repair Data Partition Table
+repairDataPartitionTable
+    : REPAIR DATA PARTITION TABLE
+    ;
+
+// Show Repair Data Partition Table Progress
+showRepairDataPartitionTableProgress
+    : SHOW REPAIR DATA PARTITION TABLE PROGRESS
+    ;
+
 // Explain
 explain
     : EXPLAIN (ANALYZE VERBOSE?)? selectStatement?
@@ -1432,4 +1443,4 @@ subStringExpression
 
 signedIntegerLiteral
     : (PLUS|MINUS)?INTEGER_LITERAL
-    ;
+    ;

iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4

@@ -806,6 +806,10 @@ SYSTEM
     : S Y S T E M
     ;
 
+TABLE
+    : T A B L E
+    ;
+
 TAGS
     : T A G S
     ;
@@ -1082,6 +1086,10 @@ REPAIR
     : R E P A I R
     ;
 
+PROGRESS
+    : P R O G R E S S
+    ;
+
 SCHEMA_REPLICATION_FACTOR
     : S C H E M A '_' R E P L I C A T I O N '_' F A C T O R
     ;
@@ -1277,4 +1285,4 @@ fragment V: [vV];
 fragment W: [wW];
 fragment X: [xX];
 fragment Y: [yY];
-fragment Z: [zZ];
+fragment Z: [zZ];

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/CnToDnSyncRequestType.java

@@ -36,6 +36,12 @@ public enum CnToDnSyncRequestType {
   DELETE_OLD_REGION_PEER,
   RESET_PEER_LIST,
 
+  // Data Partition Table Maintenance
+  COLLECT_EARLIEST_TIMESLOTS,
+  GENERATE_DATA_PARTITION_TABLE,
+  GENERATE_DATA_PARTITION_TABLE_HEART_BEAT,
+  GET_DATA_PARTITION_TABLE_GENERATOR_PROGRESS,
+
   // PartitionCache
   INVALIDATE_PARTITION_CACHE,
   INVALIDATE_PERMISSION_CACHE,

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/sync/SyncDataNodeClientPool.java

@@ -32,6 +32,7 @@
 import org.apache.iotdb.mpp.rpc.thrift.TCreateDataRegionReq;
 import org.apache.iotdb.mpp.rpc.thrift.TCreatePeerReq;
 import org.apache.iotdb.mpp.rpc.thrift.TCreateSchemaRegionReq;
+import org.apache.iotdb.mpp.rpc.thrift.TGenerateDataPartitionTableReq;
 import org.apache.iotdb.mpp.rpc.thrift.TInvalidateCacheReq;
 import org.apache.iotdb.mpp.rpc.thrift.TInvalidatePermissionCacheReq;
 import org.apache.iotdb.mpp.rpc.thrift.TMaintainPeerReq;
@@ -131,6 +132,19 @@ private void buildActionMap() {
         (req, client) -> client.resetPeerList((TResetPeerListReq) req));
     actionMapBuilder.put(
         CnToDnSyncRequestType.SHOW_CONFIGURATION, (req, client) -> client.showConfiguration());
+    actionMapBuilder.put(
+        CnToDnSyncRequestType.COLLECT_EARLIEST_TIMESLOTS,
+        (req, client) -> client.getEarliestTimeslots());
+    actionMapBuilder.put(
+        CnToDnSyncRequestType.GENERATE_DATA_PARTITION_TABLE,
+        (req, client) -> client.generateDataPartitionTable((TGenerateDataPartitionTableReq) req));
+    actionMapBuilder.put(
+        CnToDnSyncRequestType.GENERATE_DATA_PARTITION_TABLE_HEART_BEAT,
+        (req, client) ->
+            client.generateDataPartitionTableHeartbeat((TGenerateDataPartitionTableReq) req));
+    actionMapBuilder.put(
+        CnToDnSyncRequestType.GET_DATA_PARTITION_TABLE_GENERATOR_PROGRESS,
+        (req, client) -> client.getDataPartitionTableGeneratorProgress());
     actionMap = actionMapBuilder.build();
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java

@@ -44,6 +44,7 @@
 import org.apache.iotdb.commons.conf.ConfigurationFileUtils;
 import org.apache.iotdb.commons.conf.IoTDBConstant;
 import org.apache.iotdb.commons.conf.TrimProperties;
+import org.apache.iotdb.commons.enums.RepairDataPartitionTableProgressState;
 import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.exception.MetadataException;
 import org.apache.iotdb.commons.path.PartialPath;
@@ -211,6 +212,7 @@
 import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
 import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
 import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowRepairDataPartitionTableProgressResp;
 import org.apache.iotdb.confignode.rpc.thrift.TShowSubscriptionReq;
 import org.apache.iotdb.confignode.rpc.thrift.TShowSubscriptionResp;
 import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq;
@@ -400,15 +402,15 @@ protected void setLoadManager() {
   }
 
   public void close() throws IOException {
-    if (consensusManager.get() != null) {
-      consensusManager.get().close();
-    }
     if (partitionManager != null) {
       partitionManager.getRegionMaintainer().shutdown();
     }
     if (procedureManager != null) {
       procedureManager.stopExecutor();
     }
+    if (consensusManager.get() != null) {
+      consensusManager.get().close();
+    }
   }
 
   @Override
@@ -1044,6 +1046,28 @@ public TDataPartitionTableResp getOrCreateDataPartition(
     return resp;
   }
 
+  @Override
+  public TSStatus dataPartitionTableIntegrityCheck() {
+    TSStatus status = confirmLeader();
+    if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+      return status;
+    }
+
+    return partitionManager.dataPartitionTableIntegrityCheck();
+  }
+
+  @Override
+  public TShowRepairDataPartitionTableProgressResp showRepairDataPartitionTableProgress() {
+    TSStatus status = confirmLeader();
+    if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
+      return new TShowRepairDataPartitionTableProgressResp(
+              status, RepairDataPartitionTableProgressState.UNKNOWN.name(), 0.0)
+          .setMessage(status.getMessage());
+    }
+
+    return partitionManager.showRepairDataPartitionTableProgress();
+  }
+
   private void printNewCreatedDataPartition(
       GetOrCreateDataPartitionPlan getOrCreateDataPartitionPlan, TDataPartitionTableResp resp) {
     final String lineSeparator = System.lineSeparator();

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/IManager.java

@@ -136,6 +136,7 @@
 import org.apache.iotdb.confignode.rpc.thrift.TShowModelResp;
 import org.apache.iotdb.confignode.rpc.thrift.TShowPipeReq;
 import org.apache.iotdb.confignode.rpc.thrift.TShowPipeResp;
+import org.apache.iotdb.confignode.rpc.thrift.TShowRepairDataPartitionTableProgressResp;
 import org.apache.iotdb.confignode.rpc.thrift.TShowSubscriptionReq;
 import org.apache.iotdb.confignode.rpc.thrift.TShowSubscriptionResp;
 import org.apache.iotdb.confignode.rpc.thrift.TShowTopicReq;
@@ -448,6 +449,10 @@ TSchemaNodeManagementResp getNodePathsPartition(
   TDataPartitionTableResp getOrCreateDataPartition(
       GetOrCreateDataPartitionPlan getOrCreateDataPartitionPlan);
 
+  TSStatus dataPartitionTableIntegrityCheck();
+
+  TShowRepairDataPartitionTableProgressResp showRepairDataPartitionTableProgress();
+
   /**
    * Operate Permission.
    *