Iotdb 2977

integration/src/test/java/org/apache/iotdb/db/integration/IoTDBSnapshotIT.java

@@ -0,0 +1,259 @@
+/*
+ * 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.integration;
+
+import org.apache.iotdb.commons.consensus.DataRegionId;
+import org.apache.iotdb.commons.exception.IllegalPathException;
+import org.apache.iotdb.commons.exception.MetadataException;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.constant.TestConstant;
+import org.apache.iotdb.db.engine.StorageEngineV2;
+import org.apache.iotdb.db.engine.cache.ChunkCache;
+import org.apache.iotdb.db.engine.cache.TimeSeriesMetadataCache;
+import org.apache.iotdb.db.engine.snapshot.SnapshotLoader;
+import org.apache.iotdb.db.engine.snapshot.SnapshotTaker;
+import org.apache.iotdb.db.engine.snapshot.exception.DirectoryNotLegalException;
+import org.apache.iotdb.db.engine.storagegroup.DataRegion;
+import org.apache.iotdb.db.exception.StorageEngineException;
+import org.apache.iotdb.integration.env.EnvFactory;
+
+import org.apache.commons.io.FileUtils;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.HashMap;
+import java.util.Map;
+
+public class IoTDBSnapshotIT {
+  final String SG_NAME = "root.snapshotTest";
+
+  @Before
+  public void setUp() throws Exception {
+    EnvFactory.getEnv().initBeforeTest();
+    IoTDBDescriptor.getInstance().getConfig().setEnableCrossSpaceCompaction(false);
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    EnvFactory.getEnv().cleanAfterTest();
+    IoTDBDescriptor.getInstance().getConfig().setEnableCrossSpaceCompaction(true);
+  }
+
+  @Test
+  public void testTakeSnapshot()
+      throws SQLException, IllegalPathException, StorageEngineException, IOException,
+          DirectoryNotLegalException {
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("set storage group to " + SG_NAME);
+      for (int i = 0; i < 10; ++i) {
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j));
+        }
+        statement.execute("flush");
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j + 1));
+        }
+        statement.execute("flush");
+      }
+
+      DataRegion region = StorageEngineV2.getInstance().getDataRegion(new DataRegionId(0));
+      File snapshotDir = new File(TestConstant.OUTPUT_DATA_DIR, "snapshot");
+      if (snapshotDir.exists()) {
+        FileUtils.forceDelete(snapshotDir);
+      }
+
+      new SnapshotTaker(region).takeFullSnapshot(snapshotDir.getAbsolutePath(), true);
+
+      Assert.assertTrue(snapshotDir.exists());
+      Assert.assertTrue(snapshotDir.isDirectory());
+      File[] seqTsfiles =
+          snapshotDir.listFiles((dir, name) -> name.endsWith(".tsfile") && name.startsWith("seq"));
+      File[] unseqTsfiles =
+          snapshotDir.listFiles(
+              (dir, name) -> name.endsWith(".tsfile") && name.startsWith("unseq"));
+      File[] tsfileResources =
+          snapshotDir.listFiles((dir, name) -> name.endsWith(".tsfile.resource"));
+      Assert.assertNotNull(seqTsfiles);
+      Assert.assertNotNull(unseqTsfiles);
+      Assert.assertNotNull(tsfileResources);
+      Assert.assertEquals(10, seqTsfiles.length);
+      Assert.assertEquals(10, unseqTsfiles.length);
+      Assert.assertEquals(20, tsfileResources.length);
+    }
+  }
+
+  @Test(expected = DirectoryNotLegalException.class)
+  public void testTakeSnapshotInNotEmptyDir()
+      throws SQLException, IOException, IllegalPathException, StorageEngineException,
+          DirectoryNotLegalException {
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      statement.execute("set storage group to " + SG_NAME);
+      for (int i = 0; i < 10; ++i) {
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j));
+        }
+        statement.execute("flush");
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j + 1));
+        }
+        statement.execute("flush");
+      }
+
+      DataRegion region = StorageEngineV2.getInstance().getDataRegion(new DataRegionId(0));
+      File snapshotDir = new File(TestConstant.OUTPUT_DATA_DIR, "snapshot");
+      if (!snapshotDir.exists()) {
+        snapshotDir.mkdirs();
+      }
+
+      File tmpFile = new File(snapshotDir, "test");
+      tmpFile.createNewFile();
+
+      new SnapshotTaker(region).takeFullSnapshot(snapshotDir.getAbsolutePath(), true);
+    }
+  }
+
+  @Test
+  public void testLoadSnapshot()
+      throws SQLException, MetadataException, StorageEngineException, DirectoryNotLegalException,
+          IOException {
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      Map<String, Integer> resultMap = new HashMap<>();
+      statement.execute("set storage group to " + SG_NAME);
+      for (int i = 0; i < 10; ++i) {
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j));
+        }
+        statement.execute("flush");
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j + 1));
+        }
+        statement.execute("flush");
+      }
+      ResultSet resultSet = statement.executeQuery("select ** from root");
+      while (resultSet.next()) {
+        long time = resultSet.getLong("Time");
+        for (int i = 0; i < 10; ++i) {
+          String measurment = SG_NAME + ".d" + i + ".s";
+          int res = resultSet.getInt(SG_NAME + ".d" + i + ".s");
+          resultMap.put(time + measurment, res);
+        }
+      }
+
+      DataRegion region = StorageEngineV2.getInstance().getDataRegion(new DataRegionId(0));
+      File snapshotDir = new File(TestConstant.OUTPUT_DATA_DIR, "snapshot");
+      if (!snapshotDir.exists()) {
+        snapshotDir.mkdirs();
+      }
+      new SnapshotTaker(region).takeFullSnapshot(snapshotDir.getAbsolutePath(), true);
+      StorageEngineV2.getInstance()
+          .setDataRegion(
+              new DataRegionId(0),
+              new SnapshotLoader(snapshotDir.getAbsolutePath(), SG_NAME, "0")
+                  .loadSnapshotForStateMachine());
+
+      ChunkCache.getInstance().clear();
+      TimeSeriesMetadataCache.getInstance().clear();
+      resultSet = statement.executeQuery("select ** from root");
+      while (resultSet.next()) {
+        long time = resultSet.getLong("Time");
+        for (int i = 0; i < 10; ++i) {
+          String measurment = SG_NAME + ".d" + i + ".s";
+          int res = resultSet.getInt(SG_NAME + ".d" + i + ".s");
+          Assert.assertEquals(resultMap.get(time + measurment).intValue(), res);
+        }
+      }
+    }
+  }
+
+  @Test
+  public void testTakeAndLoadSnapshotWhenCompaction()
+      throws SQLException, MetadataException, StorageEngineException, InterruptedException,
+          DirectoryNotLegalException, IOException {
+    try (Connection connection = EnvFactory.getEnv().getConnection();
+        Statement statement = connection.createStatement()) {
+      Map<String, Integer> resultMap = new HashMap<>();
+      statement.execute("set storage group to " + SG_NAME);
+      for (int i = 0; i < 10; ++i) {
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j));
+        }
+        statement.execute("flush");
+        for (int j = 0; j < 10; ++j) {
+          statement.execute(
+              String.format("insert into %s.d%d(time, s) values (%d, %d)", SG_NAME, i, j, j + 1));
+        }
+        statement.execute("flush");
+      }
+
+      ResultSet resultSet = statement.executeQuery("select ** from root");
+      while (resultSet.next()) {
+        long time = resultSet.getLong("Time");
+        for (int i = 0; i < 10; ++i) {
+          String measurment = SG_NAME + ".d" + i + ".s";
+          int res = resultSet.getInt(SG_NAME + ".d" + i + ".s");
+          resultMap.put(time + measurment, res);
+        }
+      }
+
+      File snapshotDir = new File(TestConstant.OUTPUT_DATA_DIR, "snapshot");
+      if (!snapshotDir.exists()) {
+        snapshotDir.mkdirs();
+      }
+      IoTDBDescriptor.getInstance().getConfig().setEnableCrossSpaceCompaction(true);
+      statement.execute("merge");
+      DataRegion region = StorageEngineV2.getInstance().getDataRegion(new DataRegionId(0));
+      new SnapshotTaker(region).takeFullSnapshot(snapshotDir.getAbsolutePath(), true);
+      region.abortCompaction();
+      StorageEngineV2.getInstance()
+          .setDataRegion(
+              new DataRegionId(0),
+              new SnapshotLoader(snapshotDir.getAbsolutePath(), SG_NAME, "0")
+                  .loadSnapshotForStateMachine());
+      ChunkCache.getInstance().clear();
+      TimeSeriesMetadataCache.getInstance().clear();
+      resultSet = statement.executeQuery("select ** from root");
+      while (resultSet.next()) {
+        long time = resultSet.getLong("Time");
+        for (int i = 0; i < 10; ++i) {
+          String measurment = SG_NAME + ".d" + i + ".s";
+          int res = resultSet.getInt(SG_NAME + ".d" + i + ".s");
+          Assert.assertEquals(resultMap.get(time + measurment).intValue(), res);
+        }
+      }
+    }
+  }
+}

server/src/main/java/org/apache/iotdb/db/consensus/statemachine/DataRegionStateMachine.java

@@ -20,8 +20,12 @@
 package org.apache.iotdb.db.consensus.statemachine;
 
 import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.consensus.DataRegionId;
 import org.apache.iotdb.consensus.common.DataSet;
 import org.apache.iotdb.db.consensus.statemachine.visitor.DataExecutionVisitor;
+import org.apache.iotdb.db.engine.StorageEngineV2;
+import org.apache.iotdb.db.engine.snapshot.SnapshotLoader;
+import org.apache.iotdb.db.engine.snapshot.SnapshotTaker;
 import org.apache.iotdb.db.engine.storagegroup.DataRegion;
 import org.apache.iotdb.db.mpp.execution.fragment.FragmentInstanceManager;
 import org.apache.iotdb.db.mpp.plan.planner.plan.FragmentInstance;
@@ -39,7 +43,7 @@ public class DataRegionStateMachine extends BaseStateMachine {
   private static final FragmentInstanceManager QUERY_INSTANCE_MANAGER =
       FragmentInstanceManager.getInstance();
 
-  private final DataRegion region;
+  private DataRegion region;
 
   public DataRegionStateMachine(DataRegion region) {
     this.region = region;
@@ -53,11 +57,39 @@ public void stop() {}
 
   @Override
   public boolean takeSnapshot(File snapshotDir) {
-    return false;
+    try {
+      return new SnapshotTaker(region).takeFullSnapshot(snapshotDir.getAbsolutePath(), true);
+    } catch (Exception e) {
+      logger.error(
+          "Exception occurs when taking snapshot for {}-{} in {}",
+          region.getLogicalStorageGroupName(),
+          region.getDataRegionId(),
+          snapshotDir,
+          e);
+      return false;
+    }
   }
 
   @Override
-  public void loadSnapshot(File latestSnapshotRootDir) {}
+  public void loadSnapshot(File latestSnapshotRootDir) {
+    DataRegion newRegion =
+        new SnapshotLoader(
+                latestSnapshotRootDir.getAbsolutePath(),
+                region.getLogicalStorageGroupName(),
+                region.getDataRegionId())
+            .loadSnapshotForStateMachine();
+    if (newRegion == null) {
+      logger.error("Fail to load snapshot from {}", latestSnapshotRootDir);
+      return;
+    }
+    this.region = newRegion;
+    try {
+      StorageEngineV2.getInstance()
+          .setDataRegion(new DataRegionId(Integer.parseInt(region.getDataRegionId())), region);
+    } catch (Exception e) {
+      logger.error("Exception occurs when replacing data region in storage engine.", e);
+    }
+  }
 
   @Override
   protected TSStatus write(FragmentInstance fragmentInstance) {

server/src/main/java/org/apache/iotdb/db/engine/StorageEngineV2.java

@@ -601,6 +601,19 @@ public DataRegion getDataRegion(DataRegionId regionId) {
     return dataRegionMap.get(regionId);
   }
 
+  public void setDataRegion(DataRegionId regionId, DataRegion newRegion) {
+    if (dataRegionMap.containsKey(regionId)) {
+      DataRegion oldRegion = dataRegionMap.get(regionId);
+      oldRegion.syncCloseAllWorkingTsFileProcessors();
+      oldRegion.abortCompaction();
+    }
+    dataRegionMap.put(regionId, newRegion);
+  }
+
+  public TsFileFlushPolicy getFileFlushPolicy() {
+    return fileFlushPolicy;
+  }
+
   static class InstanceHolder {
 
     private static final StorageEngineV2 INSTANCE = new StorageEngineV2();