[iotdb] branch master updated: [ISSUES-6057] Add a method(getAllDataRegionID) to StorageEngineV2 (#6059)

[ha...@apache.org]

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

haonan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git


The following commit(s) were added to refs/heads/master by this push:
     new bb1a88540d [ISSUES-6057] Add a method(getAllDataRegionID) to StorageEngineV2 (#6059)
bb1a88540d is described below

commit bb1a88540d801ad56d0fc1dbaec8063cef321a6e
Author: QiangShaowei <ba...@163.com>
AuthorDate: Mon May 30 12:48:33 2022 +0800

    [ISSUES-6057] Add a method(getAllDataRegionID) to StorageEngineV2 (#6059)
---
 .../apache/iotdb/db/engine/StorageEngineV2.java    | 11 ++--
 .../iotdb/db/engine/StorageEngineV2Test.java       | 71 ++++++++++++++++++++++
 2 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/server/src/main/java/org/apache/iotdb/db/engine/StorageEngineV2.java b/server/src/main/java/org/apache/iotdb/db/engine/StorageEngineV2.java
index 94ea120d20..dfe4824466 100644
--- a/server/src/main/java/org/apache/iotdb/db/engine/StorageEngineV2.java
+++ b/server/src/main/java/org/apache/iotdb/db/engine/StorageEngineV2.java
@@ -22,7 +22,6 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
 import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
 import org.apache.iotdb.commons.concurrent.ThreadName;
-import org.apache.iotdb.commons.consensus.ConsensusGroupId;
 import org.apache.iotdb.commons.consensus.DataRegionId;
 import org.apache.iotdb.commons.exception.MetadataException;
 import org.apache.iotdb.commons.exception.ShutdownException;
@@ -105,7 +104,7 @@ public class StorageEngineV2 implements IService {
       FilePathUtils.regularizePath(config.getSystemDir()) + "storage_groups";
 
   /** DataRegionId -> DataRegion */
-  private final ConcurrentHashMap<ConsensusGroupId, DataRegion> dataRegionMap =
+  private final ConcurrentHashMap<DataRegionId, DataRegion> dataRegionMap =
       new ConcurrentHashMap<>();
 
   /** number of ready data region */
@@ -500,7 +499,7 @@ public class StorageEngineV2 implements IService {
    * @param insertRowNode
    */
   // TODO:(New insert)
-  public void insert(ConsensusGroupId dataRegionId, InsertRowNode insertRowNode)
+  public void insert(DataRegionId dataRegionId, InsertRowNode insertRowNode)
       throws StorageEngineException, MetadataException {
     if (enableMemControl) {
       try {
@@ -521,7 +520,7 @@ public class StorageEngineV2 implements IService {
 
   /** insert an InsertTabletNode to a storage group */
   // TODO:(New insert)
-  public void insertTablet(ConsensusGroupId dataRegionId, InsertTabletNode insertTabletNode)
+  public void insertTablet(DataRegionId dataRegionId, InsertTabletNode insertTabletNode)
       throws StorageEngineException, BatchProcessException {
     if (enableMemControl) {
       try {
@@ -601,6 +600,10 @@ public class StorageEngineV2 implements IService {
     return dataRegionMap.get(regionId);
   }
 
+  public List<DataRegionId> getAllDataRegionIds() {
+    return new ArrayList<>(dataRegionMap.keySet());
+  }
+
   public void setDataRegion(DataRegionId regionId, DataRegion newRegion) {
     if (dataRegionMap.containsKey(regionId)) {
       DataRegion oldRegion = dataRegionMap.get(regionId);
diff --git a/server/src/test/java/org/apache/iotdb/db/engine/StorageEngineV2Test.java b/server/src/test/java/org/apache/iotdb/db/engine/StorageEngineV2Test.java
new file mode 100644
index 0000000000..a53f186093
--- /dev/null
+++ b/server/src/test/java/org/apache/iotdb/db/engine/StorageEngineV2Test.java
@@ -0,0 +1,71 @@
+/*
+ * 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.engine;
+
+import org.apache.iotdb.commons.consensus.DataRegionId;
+import org.apache.iotdb.db.engine.storagegroup.DataRegion;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.api.mockito.PowerMockito;
+import org.powermock.core.classloader.annotations.PowerMockIgnore;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+import org.testcontainers.shaded.com.google.common.collect.Lists;
+
+import java.util.List;
+
+@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"})
+@RunWith(PowerMockRunner.class)
+@PrepareForTest(DataRegion.class)
+public class StorageEngineV2Test {
+
+  private StorageEngineV2 storageEngineV2;
+
+  @Before
+  public void setUp() {
+    storageEngineV2 = StorageEngineV2.getInstance();
+  }
+
+  @After
+  public void after() {
+    storageEngineV2 = null;
+  }
+
+  @Test
+  public void testGetAllDataRegionIds() throws Exception {
+    DataRegionId id1 = new DataRegionId(1);
+    DataRegion rg1 = PowerMockito.mock(DataRegion.class);
+    DataRegion rg2 = PowerMockito.mock(DataRegion.class);
+    DataRegionId id2 = new DataRegionId(2);
+    storageEngineV2.setDataRegion(id1, rg1);
+    storageEngineV2.setDataRegion(id2, rg2);
+
+    List<DataRegionId> actual = Lists.newArrayList(id1, id2);
+    List<DataRegionId> expect = storageEngineV2.getAllDataRegionIds();
+
+    Assert.assertEquals(expect.size(), actual.size());
+    Assert.assertTrue(actual.containsAll(expect));
+    rg1.syncDeleteDataFiles();
+    rg2.syncDeleteDataFiles();
+  }
+}