Close Zstd Dictionary after execution to avoid any memory leak. (#9403)

* Close Dictionary after every execution to avoid any memory leak Signed-off-by: Mohit Godwani <mgodwan@amazon.com> * Close Dictionary after every execution to avoid any memory leak Signed-off-by: Mohit Godwani <mgodwan@amazon.com> * Add changelog Signed-off-by: Mohit Godwani <mgodwan@amazon.com> --------- Signed-off-by: Mohit Godwani <mgodwan@amazon.com> (cherry picked from commit 5cc7313)
opensearch-project · Aug 17, 2023 · 18d3f5a · 18d3f5a
1 parent d555d1d
commit 18d3f5a
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 29 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -85,6 +85,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Fixed
 - Fix flaky ResourceAwareTasksTests.testBasicTaskResourceTracking test ([#8993](https://github.com/opensearch-project/OpenSearch/pull/8993))
+- Fix memory leak when using Zstd Dictionary ([#9403](https://github.com/opensearch-project/OpenSearch/pull/9403))
 
 ### Security
 

diff --git a/server/src/main/java/org/opensearch/index/codec/customcodecs/ZstdCompressionMode.java b/server/src/main/java/org/opensearch/index/codec/customcodecs/ZstdCompressionMode.java
@@ -103,11 +103,13 @@ private void compress(byte[] bytes, int offset, int length, DataOutput out) thro
 
                 // dictionary compression first
                 doCompress(bytes, offset, dictLength, cctx, out);
-                cctx.loadDict(new ZstdDictCompress(bytes, offset, dictLength, compressionLevel));
+                try (ZstdDictCompress dictCompress = new ZstdDictCompress(bytes, offset, dictLength, compressionLevel)) {
+                    cctx.loadDict(dictCompress);
 
-                for (int start = offset + dictLength; start < end; start += blockLength) {
-                    int l = Math.min(blockLength, end - start);
-                    doCompress(bytes, start, l, cctx, out);
+                    for (int start = offset + dictLength; start < end; start += blockLength) {
+                        int l = Math.min(blockLength, end - start);
+                        doCompress(bytes, start, l, cctx, out);
+                    }
                 }
             }
         }
@@ -170,32 +172,33 @@ public void decompress(DataInput in, int originalLength, int offset, int length,
 
                 // decompress dictionary first
                 doDecompress(in, dctx, bytes, dictLength);
-
-                dctx.loadDict(new ZstdDictDecompress(bytes.bytes, 0, dictLength));
-
-                int offsetInBlock = dictLength;
-                int offsetInBytesRef = offset;
-
-                // Skip unneeded blocks
-                while (offsetInBlock + blockLength < offset) {
-                    final int compressedLength = in.readVInt();
-                    in.skipBytes(compressedLength);
-                    offsetInBlock += blockLength;
-                    offsetInBytesRef -= blockLength;
+                try (ZstdDictDecompress dictDecompress = new ZstdDictDecompress(bytes.bytes, 0, dictLength)) {
+                    dctx.loadDict(dictDecompress);
+
+                    int offsetInBlock = dictLength;
+                    int offsetInBytesRef = offset;
+
+                    // Skip unneeded blocks
+                    while (offsetInBlock + blockLength < offset) {
+                        final int compressedLength = in.readVInt();
+                        in.skipBytes(compressedLength);
+                        offsetInBlock += blockLength;
+                        offsetInBytesRef -= blockLength;
+                    }
+
+                    // Read blocks that intersect with the interval we need
+                    while (offsetInBlock < offset + length) {
+                        bytes.bytes = ArrayUtil.grow(bytes.bytes, bytes.length + blockLength);
+                        int l = Math.min(blockLength, originalLength - offsetInBlock);
+                        doDecompress(in, dctx, bytes, l);
+                        offsetInBlock += blockLength;
+                    }
+
+                    bytes.offset = offsetInBytesRef;
+                    bytes.length = length;
+
+                    assert bytes.isValid() : "decompression output is corrupted";
                 }
-
-                // Read blocks that intersect with the interval we need
-                while (offsetInBlock < offset + length) {
-                    bytes.bytes = ArrayUtil.grow(bytes.bytes, bytes.length + blockLength);
-                    int l = Math.min(blockLength, originalLength - offsetInBlock);
-                    doDecompress(in, dctx, bytes, l);
-                    offsetInBlock += blockLength;
-                }
-
-                bytes.offset = offsetInBytesRef;
-                bytes.length = length;
-
-                assert bytes.isValid() : "decompression output is corrupted";
             }
         }