[apache#1658] fix(server) FileSegmentManagedBuffer#nioByteBuffer read…

… result cacheable
xumanbu · Apr 18, 2024 · a36ed14 · a36ed14
1 parent ca0f0ac
commit a36ed14
Showing 1 changed file with 18 additions and 8 deletions.
diff --git a/common/src/main/java/org/apache/uniffle/common/netty/buffer/FileSegmentManagedBuffer.java b/common/src/main/java/org/apache/uniffle/common/netty/buffer/FileSegmentManagedBuffer.java
@@ -39,6 +39,8 @@ public class FileSegmentManagedBuffer extends ManagedBuffer {
   private final File file;
   private final long offset;
   private final int length;
+  private boolean isFilled;
+  private ByteBuffer readByteBuffer;
 
   public FileSegmentManagedBuffer(File file, long offset, int length) {
     this.file = file;
@@ -58,27 +60,32 @@ public ByteBuf byteBuf() {
 
   @Override
   public ByteBuffer nioByteBuffer() {
+    if (isFilled) {
+      return readByteBuffer;
+    }
     FileChannel channel = null;
     try {
       channel = new RandomAccessFile(file, "r").getChannel();
-      ByteBuffer buf = ByteBuffer.allocate(length);
+      readByteBuffer = ByteBuffer.allocate(length);
       channel.position(offset);
-      while (buf.remaining() != 0) {
-        if (channel.read(buf) == -1) {
+      while (readByteBuffer.remaining() != 0) {
+        if (channel.read(readByteBuffer) == -1) {
           throw new IOException(
               String.format(
                   "Reached EOF before filling buffer.offset=%s,file=%s,buf.remaining=%s",
-                  offset, file.getAbsoluteFile(), buf.remaining()));
+                  offset, file.getAbsoluteFile(), readByteBuffer.remaining()));
         }
       }
-      buf.flip();
-      return buf;
+      readByteBuffer.flip();
+      isFilled = true;
+      return readByteBuffer;
     } catch (IOException e) {
-      String errorMessage = "Error in reading " + this;
+      String errorMessage = "Error in reading file " + file + " " + this;
       try {
         if (channel != null) {
           long size = channel.size();
-          errorMessage = "Error in reading " + this + " (actual file length " + size + ")";
+          errorMessage =
+              "Error in reading file " + file + " " + this + " (actual file length " + size + ")";
         }
       } catch (IOException ignored) {
         // ignore
@@ -98,6 +105,9 @@ public ManagedBuffer retain() {
 
   @Override
   public ManagedBuffer release() {
+    readByteBuffer.clear();
+    readByteBuffer = null;
+    isFilled = false;
     return this;
   }