(iotdb) 01/01: Finish

[yongzao]

This is an automated email from the ASF dual-hosted git repository.

yongzao pushed a commit to branch partition-policy-table-concurrent-control
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit ba7f2bc46819ef38ed287dcec379c3d70a09693c
Author: YongzaoDan <532741...@qq.com>
AuthorDate: Mon Apr 22 17:30:01 2024 +0800

    Finish
---
 .../apache/iotdb/confignode/manager/load/LoadManager.java  |  2 +-
 .../manager/load/balancer/PartitionBalancer.java           |  8 ++++++--
 .../confignode/manager/partition/PartitionManager.java     | 14 ++++++++++----
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git 
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java
 
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java
index 895d89246b8..eeb4c39fd44 100644
--- 
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java
+++ 
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/LoadManager.java
@@ -119,7 +119,7 @@ public class LoadManager {
    */
   public Map<String, DataPartitionTable> allocateDataPartition(
       Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> 
unassignedDataPartitionSlotsMap)
-      throws NoAvailableRegionGroupException {
+      throws DatabaseNotExistsException, NoAvailableRegionGroupException {
     return 
partitionBalancer.allocateDataPartition(unassignedDataPartitionSlotsMap);
   }
 
diff --git 
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java
 
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java
index 51bf521c994..70f25323df3 100644
--- 
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java
+++ 
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/balancer/PartitionBalancer.java
@@ -43,6 +43,7 @@ import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -112,8 +113,8 @@ public class PartitionBalancer {
    */
   public Map<String, DataPartitionTable> allocateDataPartition(
       Map<String, Map<TSeriesPartitionSlot, TTimeSlotList>> 
unassignedDataPartitionSlotsMap)
-      throws NoAvailableRegionGroupException {
-    Map<String, DataPartitionTable> result = new HashMap<>();
+      throws DatabaseNotExistsException, NoAvailableRegionGroupException {
+    Map<String, DataPartitionTable> result = new TreeMap<>();
 
     for (Map.Entry<String, Map<TSeriesPartitionSlot, TTimeSlotList>> 
slotsMapEntry :
         unassignedDataPartitionSlotsMap.entrySet()) {
@@ -133,6 +134,9 @@ public class PartitionBalancer {
       }
 
       DataPartitionTable dataPartitionTable = new DataPartitionTable();
+      if (!dataPartitionPolicyTableMap.containsKey(database)) {
+        throw new DatabaseNotExistsException(database);
+      }
       DataPartitionPolicyTable allotTable = 
dataPartitionPolicyTableMap.get(database);
       try {
         allotTable.acquireLock();
diff --git 
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java
 
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java
index 500a6d82660..5daf89f8dd5 100644
--- 
a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java
+++ 
b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/partition/PartitionManager.java
@@ -389,7 +389,7 @@ public class PartitionManager {
       try {
         assignedDataPartition =
             
getLoadManager().allocateDataPartition(unassignedDataPartitionSlotsMap);
-      } catch (NoAvailableRegionGroupException e) {
+      } catch (DatabaseNotExistsException | NoAvailableRegionGroupException e) 
{
         status = getConsensusManager().confirmLeader();
         if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
           // The allocation might fail due to leadership change
@@ -398,9 +398,15 @@ public class PartitionManager {
         }
 
         LOGGER.error("Create DataPartition failed because: ", e);
-        resp.setStatus(
-            new 
TSStatus(TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode())
-                .setMessage(e.getMessage()));
+        if (e instanceof DatabaseNotExistsException) {
+          resp.setStatus(
+              new TSStatus(TSStatusCode.DATABASE_NOT_EXIST.getStatusCode())
+                  .setMessage(e.getMessage()));
+        } else {
+          resp.setStatus(
+              new 
TSStatus(TSStatusCode.NO_AVAILABLE_REGION_GROUP.getStatusCode())
+                  .setMessage(e.getMessage()));
+        }
         return resp;
       }