reuse the ctx

Author: Junfan Zhang

client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriteBufferManager.java

@@ -17,6 +17,8 @@
 
 package org.apache.spark.shuffle.writer;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -55,7 +57,7 @@
 import org.apache.uniffle.common.util.BlockIdLayout;
 import org.apache.uniffle.common.util.ChecksumUtils;
 
-public class WriteBufferManager extends MemoryConsumer {
+public class WriteBufferManager extends MemoryConsumer implements Closeable {
 
   private static final Logger LOG = LoggerFactory.getLogger(WriteBufferManager.class);
   private int bufferSize;
@@ -674,4 +676,11 @@ public void setPartitionAssignmentRetrieveFunc(
       Function<Integer, List<ShuffleServerInfo>> partitionAssignmentRetrieveFunc) {
     this.partitionAssignmentRetrieveFunc = partitionAssignmentRetrieveFunc;
   }
+
+  @Override
+  public void close() throws IOException {
+    if (codec.isPresent()) {
+      codec.get().close();
+    }
+  }
 }

client-spark/spark2/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java

@@ -514,6 +514,11 @@ public Option<MapStatus> stop(boolean success) {
       // free all memory & metadata, or memory leak happen in executor
       if (bufferManager != null) {
         bufferManager.freeAllMemory();
+        try {
+          bufferManager.close();
+        } catch (Exception e) {
+          LOG.warn("Errors on closing the buffer manager", e);
+        }
       }
       if (shuffleManager != null) {
         shuffleManager.clearTaskMeta(taskId);

client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java

@@ -860,6 +860,11 @@ public Option<MapStatus> stop(boolean success) {
       // free all memory & metadata, or memory leak happen in executor
       if (bufferManager != null) {
         bufferManager.freeAllMemory();
+        try {
+          bufferManager.close();
+        } catch (Exception e) {
+          LOG.warn("Errors on closing the buffer manager", e);
+        }
       }
       if (shuffleManager != null) {
         shuffleManager.clearTaskMeta(taskId);

common/src/main/java/org/apache/uniffle/common/compression/Codec.java

@@ -17,14 +17,16 @@
 
 package org.apache.uniffle.common.compression;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Optional;
 
 import org.apache.uniffle.common.config.RssConf;
 
 import static org.apache.uniffle.common.config.RssClientConf.COMPRESSION_TYPE;
 
-public abstract class Codec {
+public abstract class Codec implements Closeable {
 
   public static Optional<Codec> newInstance(RssConf rssConf) {
     Type type = rssConf.get(COMPRESSION_TYPE);
@@ -76,4 +78,9 @@ public enum Type {
     SNAPPY,
     NONE,
   }
+
+  @Override
+  public void close() throws IOException {
+    // ignore
+  }
 }

common/src/main/java/org/apache/uniffle/common/compression/ZstdCodec.java

@@ -17,6 +17,7 @@
 
 package org.apache.uniffle.common.compression;
 
+import java.io.IOException;
 import java.nio.ByteBuffer;
 
 import com.github.luben.zstd.Zstd;
@@ -35,6 +36,7 @@ public class ZstdCodec extends Codec {
 
   private int compressionLevel;
   private int workerNumber;
+  private ZstdCompressCtx zstdCompressCtx;
 
   private static class LazyHolder {
     static final ZstdCodec INSTANCE = new ZstdCodec();
@@ -76,16 +78,15 @@ public void decompress(ByteBuffer src, int uncompressedLen, ByteBuffer dst, int
 
   @Override
   public byte[] compress(byte[] src) {
-    ZstdCompressCtx ctx = new ZstdCompressCtx();
-    try {
+    if (this.zstdCompressCtx == null) {
+      ZstdCompressCtx ctx = new ZstdCompressCtx();
       ctx.setLevel(compressionLevel);
       if (workerNumber > 0) {
         ctx.setWorkers(workerNumber);
       }
-      return ctx.compress(src);
-    } finally {
-      ctx.close();
+      this.zstdCompressCtx = ctx;
     }
+    return this.zstdCompressCtx.compress(src);
   }
 
   @Override
@@ -119,4 +120,11 @@ public int compress(ByteBuffer src, ByteBuffer dest) {
   public int maxCompressedLength(int sourceLength) {
     return (int) Zstd.compressBound(sourceLength);
   }
+
+  @Override
+  public void close() throws IOException {
+    if (this.zstdCompressCtx != null) {
+      zstdCompressCtx.close();
+    }
+  }
 }