Merge pull request #93 from apache/jira39

[IOTDB-39]add autoRepair option for NativeRestorableIOWriter

tsfile/src/main/java/org/apache/iotdb/tsfile/write/writer/NativeRestorableIOWriter.java

@@ -53,6 +53,15 @@ public class NativeRestorableIOWriter extends TsFileIOWriter {
   }
 
   public NativeRestorableIOWriter(File file) throws IOException {
+    this(file, true);
+  }
+
+  /**
+   * @param file a given tsfile path you want to (continue to) write
+   * @param autoRepair whether auto repair the tsfile (if it is broken)
+   * @throws IOException if write failed, or the file is broken but autoRepair==false.
+   */
+  public NativeRestorableIOWriter(File file, boolean autoRepair) throws IOException {
     super();
     long fileSize;
     if (!file.exists()) {
@@ -85,20 +94,28 @@ public class NativeRestorableIOWriter extends TsFileIOWriter {
     boolean newGroup = true;
 
     TsFileSequenceReader reader = new TsFileSequenceReader(file.getAbsolutePath(), false);
-    if (fileSize <= magicStringBytes.length) {
-      LOGGER.debug("{} only has magic header, does not worth to recover.", file.getAbsolutePath());
-      reader.close();
-      this.out.truncate(0);
-      startFile();
-      truncatedPosition = magicStringBytes.length;
-      return;
+    if (fileSize < magicStringBytes.length) {
+        reader.close();
+        out.close();
+        throw new IOException(String
+            .format("%s is not using TsFile format, and will be ignored...", file.getAbsolutePath()));
     }
     String magic = reader.readHeadMagic(true);
     if (!magic.equals(new String(magicStringBytes))) {
+      reader.close();
+      out.close();
       throw new IOException(String
           .format("%s is not using TsFile format, and will be ignored...", file.getAbsolutePath()));
     }
-    if (reader.readTailMagic().equals(magic)) {
+
+    if (fileSize == magicStringBytes.length) {
+      if (!autoRepair) {
+        reader.close();
+        out.close();
+        throw new IOException(String
+            .format("%s only has header, but does not allowed to be repaired...", file.getAbsolutePath()));
+      }
+    } else if (reader.readTailMagic().equals(magic)) {
       LOGGER.debug("{} is an complete TsFile.", file.getAbsolutePath());
       canWrite = false;
       reader.close();
@@ -192,10 +209,17 @@ public class NativeRestorableIOWriter extends TsFileIOWriter {
     } finally {
       //something wrong or all data is complete. We will discard current FileMetadata
       // so that we can continue to write data into this tsfile.
-      LOGGER.info("File {} has {} bytes, and will be truncated from {}.",
-          file.getAbsolutePath(), file.length(), truncatedPosition);
-      out.truncate(truncatedPosition);
       reader.close();
+      if (autoRepair) {
+        LOGGER.info("File {} has {} bytes, and will be truncated from {}.",
+            file.getAbsolutePath(), file.length(), truncatedPosition);
+        out.truncate(truncatedPosition);
+      }
+    }
+    if (!autoRepair) {
+      out.close();
+      throw new IOException(String
+          .format("%s is incomplete but does not allowed to be repaired...", file.getAbsolutePath()));
     }
   }
 

tsfile/src/main/java/org/apache/iotdb/tsfile/write/writer/TsFileIOWriter.java

@@ -326,6 +326,7 @@ public class TsFileIOWriter {
    * close the inputstream or file channel in force. This is just used for Testing.
    */
   void forceClose() throws IOException {
+    canWrite = false;
     out.close();
   }
 

tsfile/src/test/java/org/apache/iotdb/tsfile/write/writer/NativeRestorableIOWriterTest.java

@@ -25,54 +25,62 @@ import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 
 import java.io.File;
+import java.io.FileWriter;
 import java.io.IOException;
 import java.nio.ByteBuffer;
-import java.util.List;
-import java.util.stream.Collectors;
-import org.apache.iotdb.tsfile.common.conf.TSFileConfig;
-import org.apache.iotdb.tsfile.common.conf.TSFileDescriptor;
-import org.apache.iotdb.tsfile.encoding.decoder.Decoder;
 import org.apache.iotdb.tsfile.file.MetaMarker;
-import org.apache.iotdb.tsfile.file.footer.ChunkGroupFooter;
 import org.apache.iotdb.tsfile.file.header.ChunkHeader;
-import org.apache.iotdb.tsfile.file.header.PageHeader;
-import org.apache.iotdb.tsfile.file.metadata.ChunkGroupMetaData;
-import org.apache.iotdb.tsfile.file.metadata.ChunkMetaData;
-import org.apache.iotdb.tsfile.file.metadata.TsDeviceMetadata;
 import org.apache.iotdb.tsfile.file.metadata.TsDeviceMetadataIndex;
-import org.apache.iotdb.tsfile.file.metadata.TsFileMetaData;
 import org.apache.iotdb.tsfile.file.metadata.enums.CompressionType;
 import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
 import org.apache.iotdb.tsfile.file.metadata.enums.TSEncoding;
 import org.apache.iotdb.tsfile.file.metadata.statistics.FloatStatistics;
 import org.apache.iotdb.tsfile.read.TsFileSequenceReader;
-import org.apache.iotdb.tsfile.read.common.BatchData;
-import org.apache.iotdb.tsfile.read.reader.page.PageReader;
 import org.apache.iotdb.tsfile.write.TsFileWriter;
 import org.apache.iotdb.tsfile.write.record.TSRecord;
 import org.apache.iotdb.tsfile.write.record.datapoint.FloatDataPoint;
 import org.apache.iotdb.tsfile.write.schema.MeasurementSchema;
-import org.junit.After;
-import org.junit.Before;
 import org.junit.Test;
 
 public class NativeRestorableIOWriterTest {
 
   private static final String FILE_NAME = "test.ts";
 
+  @Test
+  public void testBadHeadMagic() throws Exception {
+    File file = new File(FILE_NAME);
+    FileWriter fWriter = new FileWriter(file);
+    fWriter.write("Tsfile");
+    fWriter.close();
+    try {
+      NativeRestorableIOWriter rWriter = new NativeRestorableIOWriter(file);
+    } catch (IOException e) {
+      assertTrue(e.getMessage().contains("is not using TsFile format"));
+    }
+    assertTrue(file.delete());
+  }
+
   @Test
   public void testOnlyHeadMagic() throws Exception {
     File file = new File(FILE_NAME);
     TsFileWriter writer = new TsFileWriter(file);
     writer.getIOWriter().forceClose();
+    try {
+      NativeRestorableIOWriter rWriter = new NativeRestorableIOWriter(file, false);
+      writer = new TsFileWriter(rWriter);
+      writer.close();
+    } catch (IOException e) {
+      assertTrue(e.getMessage().contains("only has header"));
+    }
+
     NativeRestorableIOWriter rWriter = new NativeRestorableIOWriter(file);
+    assertEquals(TsFileIOWriter.magicStringBytes.length, rWriter.getTruncatedPosition());
     writer = new TsFileWriter(rWriter);
     writer.close();
-    assertEquals(TsFileIOWriter.magicStringBytes.length, rWriter.getTruncatedPosition());
     assertTrue(file.delete());
-
   }
 
+
   @Test
   public void testOnlyFirstMask() throws Exception {
     File file = new File(FILE_NAME);