[To dev/1.3] Wait for consensus before region requests

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java

@@ -149,6 +149,7 @@
 import org.apache.iotdb.db.schemaengine.schemaregion.read.resp.reader.ISchemaReader;
 import org.apache.iotdb.db.schemaengine.template.ClusterTemplateManager;
 import org.apache.iotdb.db.schemaengine.template.TemplateInternalRPCUpdateType;
+import org.apache.iotdb.db.service.ConsensusReadiness;
 import org.apache.iotdb.db.service.DataNode;
 import org.apache.iotdb.db.service.RegionMigrateService;
 import org.apache.iotdb.db.service.metrics.FileMetrics;
@@ -319,7 +320,7 @@ public class DataNodeInternalRPCServiceImpl implements IDataNodeRPCService.Iface
   private final SchemaEngine schemaEngine = SchemaEngine.getInstance();
   private final StorageEngine storageEngine = StorageEngine.getInstance();
 
-  private final DataNodeRegionManager regionManager = DataNodeRegionManager.getInstance();
+  private final DataNodeRegionManager regionManager;
 
   private final DataNodeSpaceQuotaManager spaceQuotaManager =
       DataNodeSpaceQuotaManager.getInstance();
@@ -330,8 +331,15 @@ public class DataNodeInternalRPCServiceImpl implements IDataNodeRPCService.Iface
   private static final long TEST_CONNECTION_TIMEOUT_MS =
       CommonDescriptor.getInstance().getConfig().getDnConnectionTimeoutInMS();
   private static final int TEST_CONNECTION_RETRY_NUM = 1;
+  private static final long DEFAULT_CONSENSUS_WAIT_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(30);
+  private static final String CONSENSUS_NOT_INITIALIZED_LOG =
+      "Consensus is not initialized; rejecting the region topology request after waiting up to {} ms";
+  private static final String CONSENSUS_NOT_INITIALIZED_MESSAGE =
+      "Consensus is not initialized; region topology request rejected after waiting up to %d ms";
 
   private final CommonConfig commonConfig = CommonDescriptor.getInstance().getConfig();
+  private final ConsensusReadiness consensusReadiness;
+  private final long consensusWaitTimeoutMs;
 
   private final ExecutorService schemaExecutor =
       new WrappedThreadPoolExecutor(
@@ -347,10 +355,43 @@ public class DataNodeInternalRPCServiceImpl implements IDataNodeRPCService.Iface
 
   private static final String SYSTEM = "system";
 
-  public DataNodeInternalRPCServiceImpl() {
+  DataNodeInternalRPCServiceImpl(
+      ConsensusReadiness consensusReadiness, long consensusWaitTimeoutMs) {
+    this(consensusReadiness, consensusWaitTimeoutMs, DataNodeRegionManager.getInstance());
+  }
+
+  DataNodeInternalRPCServiceImpl(
+      ConsensusReadiness consensusReadiness,
+      long consensusWaitTimeoutMs,
+      DataNodeRegionManager regionManager) {
     super();
     partitionFetcher = ClusterPartitionFetcher.getInstance();
     schemaFetcher = ClusterSchemaFetcher.getInstance();
+    this.consensusReadiness = consensusReadiness;
+    this.consensusWaitTimeoutMs = consensusWaitTimeoutMs;
+    this.regionManager = regionManager;
+  }
+
+  public DataNodeInternalRPCServiceImpl(ConsensusReadiness consensusReadiness) {
+    this(consensusReadiness, DEFAULT_CONSENSUS_WAIT_TIMEOUT_MS);
+  }
+
+  private TSStatus waitForConsensusStarted() {
+    if (consensusReadiness.isAllConsensusStarted()) {
+      return null;
+    }
+    try {
+      if (consensusReadiness.awaitAllConsensusStarted(
+          consensusWaitTimeoutMs, TimeUnit.MILLISECONDS)) {
+        return null;
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+    LOGGER.warn(CONSENSUS_NOT_INITIALIZED_LOG, consensusWaitTimeoutMs);
+    return RpcUtils.getStatus(
+        TSStatusCode.CONSENSUS_NOT_INITIALIZED,
+        String.format(CONSENSUS_NOT_INITIALIZED_MESSAGE, consensusWaitTimeoutMs));
   }
 
   @Override
@@ -525,11 +566,19 @@ private TLoadResp createTLoadResp(TSStatus resultStatus) {
 
   @Override
   public TSStatus createSchemaRegion(TCreateSchemaRegionReq req) {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     return regionManager.createSchemaRegion(req.getRegionReplicaSet(), req.getStorageGroup());
   }
 
   @Override
   public TSStatus createDataRegion(TCreateDataRegionReq req) {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     return regionManager.createDataRegion(req.getRegionReplicaSet(), req.getStorageGroup());
   }
 
@@ -2097,6 +2146,10 @@ public TSStatus updateTemplate(final TUpdateTemplateReq req) {
 
   @Override
   public TSStatus deleteRegion(TConsensusGroupId tconsensusGroupId) {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     ConsensusGroupId consensusGroupId =
         ConsensusGroupId.Factory.createFromTConsensusGroupId(tconsensusGroupId);
     if (consensusGroupId instanceof DataRegionId) {
@@ -2125,6 +2178,12 @@ public TRegionLeaderChangeResp changeRegionLeader(TRegionLeaderChangeReq req) {
     LOGGER.info("[ChangeRegionLeader] {}", req);
     TRegionLeaderChangeResp resp = new TRegionLeaderChangeResp();
 
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      resp.setStatus(consensusStatus);
+      return resp;
+    }
+
     TSStatus successStatus = new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
     TConsensusGroupId tgId = req.getRegionId();
     ConsensusGroupId regionId = ConsensusGroupId.Factory.createFromTConsensusGroupId(tgId);
@@ -2194,6 +2253,10 @@ private boolean isLeader(ConsensusGroupId regionId) {
 
   @Override
   public TSStatus createNewRegionPeer(TCreatePeerReq req) {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     ConsensusGroupId regionId =
         ConsensusGroupId.Factory.createFromTConsensusGroupId(req.getRegionId());
     List<Peer> peers =
@@ -2214,6 +2277,10 @@ public TSStatus createNewRegionPeer(TCreatePeerReq req) {
 
   @Override
   public TSStatus addRegionPeer(TMaintainPeerReq req) {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     TConsensusGroupId regionId = req.getRegionId();
     String selectedDataNodeIP = req.getDestNode().getInternalEndPoint().getIp();
     boolean submitSucceed = RegionMigrateService.getInstance().submitAddRegionPeerTask(req);
@@ -2232,6 +2299,10 @@ public TSStatus addRegionPeer(TMaintainPeerReq req) {
 
   @Override
   public TSStatus removeRegionPeer(TMaintainPeerReq req) {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     TConsensusGroupId regionId = req.getRegionId();
     String selectedDataNodeIP = req.getDestNode().getInternalEndPoint().getIp();
     boolean submitSucceed = RegionMigrateService.getInstance().submitRemoveRegionPeerTask(req);
@@ -2250,6 +2321,10 @@ public TSStatus removeRegionPeer(TMaintainPeerReq req) {
 
   @Override
   public TSStatus deleteOldRegionPeer(TMaintainPeerReq req) {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     TConsensusGroupId regionId = req.getRegionId();
     String selectedDataNodeIP = req.getDestNode().getInternalEndPoint().getIp();
     boolean submitSucceed = RegionMigrateService.getInstance().submitDeleteOldRegionPeerTask(req);
@@ -2269,6 +2344,10 @@ public TSStatus deleteOldRegionPeer(TMaintainPeerReq req) {
   // TODO: return which DataNode fail
   @Override
   public TSStatus resetPeerList(TResetPeerListReq req) throws TException {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     return RegionMigrateService.getInstance().resetPeerList(req);
   }
 
@@ -2279,6 +2358,10 @@ public TRegionMigrateResult getRegionMaintainResult(long taskId) throws TExcepti
 
   @Override
   public TSStatus notifyRegionMigration(TNotifyRegionMigrationReq req) throws TException {
+    TSStatus consensusStatus = waitForConsensusStarted();
+    if (consensusStatus != null) {
+      return consensusStatus;
+    }
     RegionMigrateService.getInstance().notifyRegionMigration(req);
     return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/ConsensusReadiness.java

@@ -0,0 +1,30 @@
+/*
+ * 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.db.service;
+
+import java.util.concurrent.TimeUnit;
+
+/** Read-only context for the startup state of SchemaRegion and DataRegion consensus. */
+public interface ConsensusReadiness {
+
+  boolean isAllConsensusStarted();
+
+  boolean awaitAllConsensusStarted(long timeout, TimeUnit unit) throws InterruptedException;
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNode.java

@@ -128,6 +128,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
@@ -167,8 +168,10 @@ public class DataNode extends ServerCommandLine implements DataNodeMBean {
   private static final String REGISTER_INTERRUPTION =
       "Unexpected interruption when waiting to register to the cluster";
 
-  private boolean schemaRegionConsensusStarted = false;
-  private boolean dataRegionConsensusStarted = false;
+  private volatile boolean schemaRegionConsensusStarted = false;
+  private volatile boolean dataRegionConsensusStarted = false;
+  private final ConsensusReadinessContext consensusReadinessContext =
+      new ConsensusReadinessContext();
   private long schemaEngineRecoveryTimeInMs;
   private static Thread watcherThread;
 
@@ -292,6 +295,7 @@ protected void start() {
             "DataRegion consensus start successfully, which takes {} ms.",
             (dataRegionEndTime - dataRegionStartTime));
         dataRegionConsensusStarted = true;
+        consensusReadinessContext.markDataRegionConsensusStarted();
       }
 
     } catch (StartupException | IOException e) {
@@ -721,13 +725,15 @@ private void active() throws StartupException {
           "SchemaRegion consensus start successfully, which takes {} ms.",
           (schemaRegionEndTime - startTime));
       schemaRegionConsensusStarted = true;
+      consensusReadinessContext.markSchemaRegionConsensusStarted();
       if (!isUsingPipeConsensus()) {
         DataRegionConsensusImpl.getInstance().start();
         long dataRegionEndTime = System.currentTimeMillis();
         logger.info(
             "DataRegion consensus start successfully, which takes {} ms.",
             (dataRegionEndTime - schemaRegionEndTime));
         dataRegionConsensusStarted = true;
+        consensusReadinessContext.markDataRegionConsensusStarted();
       }
     } catch (IOException e) {
       throw new StartupException(e);
@@ -810,7 +816,9 @@ private void setUp() throws StartupException {
   /** Set up RPC and protocols after DataNode is available */
   private void setUpRPCService() throws StartupException {
     // Start InternalRPCService to indicate that the current DataNode can accept cluster scheduling
-    registerManager.register(DataNodeInternalRPCService.getInstance());
+    DataNodeInternalRPCService internalRPCService = DataNodeInternalRPCService.getInstance();
+    internalRPCService.setConsensusReadiness(consensusReadinessContext);
+    registerManager.register(internalRPCService);
 
     // Notice: During the period between starting the internal RPC service
     // and starting the client RPC service , some requests may fail because
@@ -1235,4 +1243,38 @@ private DataNodeHolder() {
       // Empty constructor
     }
   }
+
+  static class ConsensusReadinessContext implements ConsensusReadiness {
+
+    private final CountDownLatch allConsensusStarted = new CountDownLatch(1);
+    private volatile boolean schemaRegionConsensusStarted;
+    private volatile boolean dataRegionConsensusStarted;
+
+    void markSchemaRegionConsensusStarted() {
+      schemaRegionConsensusStarted = true;
+      signalIfReady();
+    }
+
+    void markDataRegionConsensusStarted() {
+      dataRegionConsensusStarted = true;
+      signalIfReady();
+    }
+
+    private void signalIfReady() {
+      if (schemaRegionConsensusStarted && dataRegionConsensusStarted) {
+        allConsensusStarted.countDown();
+      }
+    }
+
+    @Override
+    public boolean isAllConsensusStarted() {
+      return allConsensusStarted.getCount() == 0;
+    }
+
+    @Override
+    public boolean awaitAllConsensusStarted(long timeout, TimeUnit unit)
+        throws InterruptedException {
+      return allConsensusStarted.await(timeout, unit);
+    }
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/DataNodeInternalRPCService.java

@@ -33,12 +33,43 @@
 import org.apache.iotdb.mpp.rpc.thrift.IDataNodeRPCService.Processor;
 import org.apache.iotdb.rpc.DeepCopyRpcTransportFactory;
 
+import java.util.Objects;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
 public class DataNodeInternalRPCService extends ThriftService
     implements DataNodeInternalRPCServiceMBean {
 
+  private static final ConsensusReadiness NOT_READY =
+      new ConsensusReadiness() {
+        @Override
+        public boolean isAllConsensusStarted() {
+          return false;
+        }
+
+        @Override
+        public boolean awaitAllConsensusStarted(long timeout, TimeUnit unit) {
+          return false;
+        }
+      };
+
   private final AtomicReference<DataNodeInternalRPCServiceImpl> impl = new AtomicReference<>();
+  private final AtomicReference<ConsensusReadiness> consensusReadiness =
+      new AtomicReference<>(NOT_READY);
+
+  private final ConsensusReadiness delegatingConsensusReadiness =
+      new ConsensusReadiness() {
+        @Override
+        public boolean isAllConsensusStarted() {
+          return consensusReadiness.get().isAllConsensusStarted();
+        }
+
+        @Override
+        public boolean awaitAllConsensusStarted(long timeout, TimeUnit unit)
+            throws InterruptedException {
+          return consensusReadiness.get().awaitAllConsensusStarted(timeout, unit);
+        }
+      };
 
   private DataNodeInternalRPCService() {}
 
@@ -49,9 +80,9 @@ public ServiceType getID() {
 
   @Override
   public void initTProcessor() {
-    impl.compareAndSet(null, new DataNodeInternalRPCServiceImpl());
+    DataNodeInternalRPCServiceImpl service = getImpl();
     initSyncedServiceImpl(null);
-    processor = new Processor<>(impl.get());
+    processor = new Processor<>(service);
   }
 
   @Override
@@ -90,10 +121,18 @@ public int getBindPort() {
   }
 
   public DataNodeInternalRPCServiceImpl getImpl() {
-    impl.compareAndSet(null, new DataNodeInternalRPCServiceImpl());
+    impl.compareAndSet(null, new DataNodeInternalRPCServiceImpl(delegatingConsensusReadiness));
     return impl.get();
   }
 
+  void setConsensusReadiness(ConsensusReadiness consensusReadiness) {
+    this.consensusReadiness.set(Objects.requireNonNull(consensusReadiness));
+  }
+
+  ConsensusReadiness getConsensusReadiness() {
+    return delegatingConsensusReadiness;
+  }
+
   private static class DataNodeInternalRPCServiceHolder {
     private static final DataNodeInternalRPCService INSTANCE = new DataNodeInternalRPCService();