[incubator-uniffle] 02/02: [Bugfix] Fix uncorrect index file (#92)

[js...@apache.org]

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

jshao pushed a commit to branch branch-0.2.0
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git

commit 75b537661f1a29291f199974c6e7fa1e39197d72
Author: roryqi <je...@gmail.com>
AuthorDate: Tue Mar 8 16:31:33 2022 +0800

    [Bugfix] Fix uncorrect index file (#92)
    
    ### What changes were proposed in this pull request?
    Modify the method that calculate the offset in the index file.
    
    ### Why are the changes needed?
    If we don't have this patch, we run 10TB tpcds, query24a will fail.
    <img width="361" alt="企业微信截图_6dc451cf-dbf4-4257-b680-e79346cd582d" src="https://user-images.githubusercontent.com/8159038/157178756-d8a39b3f-0ea6-4864-ac68-ee382a88bb0f.png">
    When we write many data to dataOutputStream, dataOutputStream.size() won't increase again. dataOutputStream.size() will
    always be Integer.MAX_VALUE.
    
    ### Does this PR introduce _any_ user-facing change?
    No
    
    ### How was this patch tested?
    Add new uts.
    
    Co-authored-by: roryqi <ro...@tencent.com>
---
 .../rss/storage/handler/impl/LocalFileWriter.java       |  6 ++----
 .../rss/storage/handler/impl/LocalFileHandlerTest.java  | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/storage/src/main/java/com/tencent/rss/storage/handler/impl/LocalFileWriter.java b/storage/src/main/java/com/tencent/rss/storage/handler/impl/LocalFileWriter.java
index 10185a4..609db7e 100644
--- a/storage/src/main/java/com/tencent/rss/storage/handler/impl/LocalFileWriter.java
+++ b/storage/src/main/java/com/tencent/rss/storage/handler/impl/LocalFileWriter.java
@@ -30,21 +30,19 @@ public class LocalFileWriter implements Closeable {
 
   private DataOutputStream dataOutputStream;
   private FileOutputStream fileOutputStream;
-  private long initSize;
   private long nextOffset;
 
   public LocalFileWriter(File file) throws IOException {
     fileOutputStream = new FileOutputStream(file, true);
     // init fsDataOutputStream
     dataOutputStream = new DataOutputStream(fileOutputStream);
-    initSize = file.length();
-    nextOffset = initSize;
+    nextOffset = file.length();
   }
 
   public void writeData(byte[] data) throws IOException {
     if (data != null && data.length > 0) {
       dataOutputStream.write(data);
-      nextOffset = initSize + dataOutputStream.size();
+      nextOffset = nextOffset + data.length;
     }
   }
 
diff --git a/storage/src/test/java/com/tencent/rss/storage/handler/impl/LocalFileHandlerTest.java b/storage/src/test/java/com/tencent/rss/storage/handler/impl/LocalFileHandlerTest.java
index 969944d..ce8915b 100644
--- a/storage/src/test/java/com/tencent/rss/storage/handler/impl/LocalFileHandlerTest.java
+++ b/storage/src/test/java/com/tencent/rss/storage/handler/impl/LocalFileHandlerTest.java
@@ -39,6 +39,7 @@ import com.tencent.rss.storage.handler.api.ServerReadHandler;
 import com.tencent.rss.storage.handler.api.ShuffleWriteHandler;
 import com.tencent.rss.storage.util.ShuffleStorageUtils;
 import java.io.File;
+import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
@@ -53,6 +54,7 @@ public class LocalFileHandlerTest {
   @Test
   public void writeTest() throws Exception {
     File tmpDir = Files.createTempDir();
+    tmpDir.deleteOnExit();
     File dataDir1 = new File(tmpDir, "data1");
     File dataDir2 = new File(tmpDir, "data2");
     String[] basePaths = new String[]{dataDir1.getAbsolutePath(),
@@ -111,6 +113,21 @@ public class LocalFileHandlerTest {
     }
   }
 
+  @Test
+  public void writeBigDataTest() throws IOException  {
+    File tmpDir = Files.createTempDir();
+    tmpDir.deleteOnExit();
+    File writeFile = new File(tmpDir, "writetest");
+    LocalFileWriter writer = new LocalFileWriter(writeFile);
+    int  size = Integer.MAX_VALUE / 100;
+    byte[] data = new byte[size];
+    for (int i = 0; i < 200; i++) {
+      writer.writeData(data);
+    }
+    long totalSize = 200L * size;
+    assertEquals(writer.nextOffset(), totalSize);
+  }
+
 
   private void writeTestData(
       ShuffleWriteHandler writeHandler,