src: track allocations made by zstd streams

addaleax · aduh95 · commit 3c8f700fd702 · 2026-04-08T23:53:01.000+02:00
This is both valuable as a diagnostic tool and as a way to inform the JS runtime about external allocations. Currently, this is a feature only enabled in the statically linked builds of zstd, so with `--shared-zstd`, we fall back to the non-tracking variant. PR-URL: #61717 Reviewed-By: Benjamin Gruenbaum <benjamingr@gmail.com> Reviewed-By: Yagiz Nizipli <yagiz@nizipli.com> Reviewed-By: Rafael Gonzaga <rafael.nunu@hotmail.com>
diff --git a/node.gypi b/node.gypi
@@ -248,6 +248,7 @@
 
     [ 'node_shared_zstd=="false"', {
       'dependencies': [ 'deps/zstd/zstd.gyp:zstd' ],
+      'defines': [ 'NODE_BUNDLED_ZSTD' ],
     }],
 
     [ 'OS=="mac"', {
diff --git a/src/node_zlib.cc b/src/node_zlib.cc
@@ -19,6 +19,10 @@
 // OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 // USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+#ifdef NODE_BUNDLED_ZSTD
+#define ZSTD_STATIC_LINKING_ONLY
+#endif
+
 #include "memory_tracker-inl.h"
 #include "node.h"
 #include "node_buffer.h"
@@ -1549,7 +1553,16 @@ void ZstdCompressContext::Close() {
 CompressionError ZstdCompressContext::Init(uint64_t pledged_src_size,
                                            std::string_view dictionary) {
   pledged_src_size_ = pledged_src_size;
+#ifdef NODE_BUNDLED_ZSTD
+  ZSTD_customMem custom_mem = {
+      CompressionStreamMemoryOwner::AllocForBrotli,
+      CompressionStreamMemoryOwner::FreeForZlib,
+      CompressionStream<ZstdCompressContext>::AllocatorOpaquePointerForContext(
+          this)};
+  cctx_.reset(ZSTD_createCCtx_advanced(custom_mem));
+#else
   cctx_.reset(ZSTD_createCCtx());
+#endif
   if (!cctx_) {
     return CompressionError("Could not initialize zstd instance",
                             "ERR_ZLIB_INITIALIZATION_FAILED",
@@ -1604,7 +1617,16 @@ void ZstdDecompressContext::Close() {
 
 CompressionError ZstdDecompressContext::Init(uint64_t pledged_src_size,
                                              std::string_view dictionary) {
+#ifdef NODE_BUNDLED_ZSTD
+  ZSTD_customMem custom_mem = {
+      CompressionStreamMemoryOwner::AllocForBrotli,
+      CompressionStreamMemoryOwner::FreeForZlib,
+      CompressionStream<
+          ZstdDecompressContext>::AllocatorOpaquePointerForContext(this)};
+  dctx_.reset(ZSTD_createDCtx_advanced(custom_mem));
+#else
   dctx_.reset(ZSTD_createDCtx());
+#endif
   if (!dctx_) {
     return CompressionError("Could not initialize zstd instance",
                             "ERR_ZLIB_INITIALIZATION_FAILED",
diff --git a/test/pummel/test-heapdump-zstd.js b/test/pummel/test-heapdump-zstd.js
@@ -0,0 +1,49 @@
+'use strict';
+// This tests heap snapshot integration of zlib stream.
+
+const common = require('../common');
+const assert = require('assert');
+const { validateByRetainingPath, validateByRetainingPathFromNodes } = require('../common/heap');
+const zlib = require('zlib');
+
+// Before zstd stream is created, no ZstdStream should be created.
+{
+  const nodes = validateByRetainingPath('Node / ZstdStream', []);
+  assert.strictEqual(nodes.length, 0);
+}
+
+const compress = zlib.createZstdCompress();
+
+// After zstd stream is created, a ZstdStream should be created.
+{
+  const streams = validateByRetainingPath('Node / ZstdStream', []);
+  validateByRetainingPathFromNodes(streams, 'Node / ZstdStream', [
+    { node_name: 'ZstdCompress', edge_name: 'native_to_javascript' },
+  ]);
+  const withMemory = validateByRetainingPathFromNodes(streams, 'Node / ZstdStream', [
+    { node_name: 'Node / zlib_memory', edge_name: 'zlib_memory' },
+  ], true);
+  if (process.config.variables.node_shared_zstd === true) {
+    assert.strictEqual(withMemory.length, 0);
+  } else {
+    assert.strictEqual(withMemory.length, 1);
+    // Between 1KB and 1MB (measured value was around ~5kB)
+    assert.ok(withMemory[0].self_size > 1024);
+    assert.ok(withMemory[0].self_size < 1024 * 1024);
+  }
+}
+
+// After zstd stream is written, zlib_memory is significantly larger.
+compress.write('hello world', common.mustCall(() => {
+  const streams = validateByRetainingPath('Node / ZstdStream', []);
+  const withMemory = validateByRetainingPathFromNodes(streams, 'Node / ZstdStream', [
+    { node_name: 'Node / zlib_memory', edge_name: 'zlib_memory' },
+  ], true);
+  if (process.config.variables.node_shared_zstd === true) {
+    assert.strictEqual(withMemory.length, 0);
+  } else {
+    assert.strictEqual(withMemory.length, 1);
+    // More than 1MB
+    assert.ok(withMemory[0].self_size > 1024 * 1024);
+  }
+}));
