zlib: grow the output buffer of large async one-shot calls

lib/zlib.js

@@ -24,6 +24,7 @@
 const {
   ArrayBuffer,
   MathMax,
+  MathMin,
   NumberIsNaN,
   ObjectDefineProperties,
   ObjectDefineProperty,
@@ -87,6 +88,7 @@ const {
 
 const kFlushFlag = Symbol('kFlushFlag');
 const kError = Symbol('kError');
+const kBufferResult = Symbol('kBufferResult');
 
 const constants = internalBinding('constants').zlib;
 const {
@@ -124,7 +126,7 @@ for (const ckey of ObjectKeys(codes)) {
   codes[codes[ckey]] = ckey;
 }
 
-function zlibBuffer(engine, buffer, callback) {
+function zlibBuffer(engine, buffer, callback, finishInOneWrite) {
   validateFunction(callback, 'callback');
   // Streams do not support non-Uint8Array ArrayBufferViews yet. Convert it to a
   // Buffer without copying.
@@ -136,44 +138,93 @@ function zlibBuffer(engine, buffer, callback) {
   engine.buffers = null;
   engine.nread = 0;
   engine.cb = callback;
-  engine.on('data', zlibBufferOnData);
   engine.on('error', zlibBufferOnError);
+
+  // The whole input is known up front. Unless the engine is handed to the
+  // caller (`info`) or a custom flush is requested, process it with the finish
+  // flush directly, like the synchronous methods do. This takes one threadpool
+  // round trip instead of a write followed by a separate final flush, and skips
+  // the stream machinery.
+  if (finishInOneWrite && !engine._info &&
+      engine._defaultFlushFlag === FLUSH_BOUND[engine._flushBoundIdx][0]) {
+    if (typeof buffer === 'string') {
+      // Encode strings like Writable#write() does.
+      buffer = Buffer.from(buffer, engine._writableState.defaultEncoding);
+    }
+    if (isUint8Array(buffer)) {
+      engine[kBufferResult] = true;
+      processChunk(engine, buffer, engine._finishFlushFlag,
+                   zlibBufferOnWriteDone);
+      return;
+    }
+  }
+
+  engine.on('data', zlibBufferOnData);
   engine.on('end', zlibBufferOnEnd);
   engine.end(buffer);
 }
 
-function zlibBufferOnData(chunk) {
-  if (!this.buffers) {
-    this.buffers = [chunk];
+function zlibBufferAddChunk(engine, chunk) {
+  if (!engine.buffers) {
+    engine.buffers = [chunk];
   } else {
-    this.buffers.push(chunk);
+    engine.buffers.push(chunk);
   }
-  this.nread += chunk.length;
-  if (this.nread > this._maxOutputLength) {
-    this.close();
-    this.removeAllListeners('end');
-    this.cb(new ERR_BUFFER_TOO_LARGE(this._maxOutputLength));
+  engine.nread += chunk.length;
+  if (engine.nread > engine._maxOutputLength) {
+    engine.close();
+    engine.removeAllListeners('end');
+    engine.cb(new ERR_BUFFER_TOO_LARGE(engine._maxOutputLength));
   }
 }
 
+function zlibBufferOnData(chunk) {
+  zlibBufferAddChunk(this, chunk);
+}
+
+// Only the concatenated output of a one-shot call is observable, so its output
+// buffer can double in size whenever it fills up, and large results take far
+// fewer threadpool round trips. It never grows past what is needed to detect
+// that maxOutputLength was exceeded.
+const kMaxBufferResultChunkSize = 1024 * 1024;
+function zlibBufferGrowChunkSize(engine) {
+  const size = MathMin(engine._chunkSize * 2, kMaxBufferResultChunkSize,
+                       engine._maxOutputLength - engine.nread + 1);
+  if (size > engine._chunkSize)
+    engine._chunkSize = size;
+}
+
 function zlibBufferOnError(err) {
   this.removeAllListeners('end');
   this.cb(err);
 }
 
-function zlibBufferOnEnd() {
+function zlibBufferFinish(engine) {
   let buf;
-  if (this.nread === 0) {
+  if (engine.nread === 0) {
     buf = new FastBuffer();
   } else {
-    const bufs = this.buffers;
-    buf = (bufs.length === 1 ? bufs[0] : Buffer.concat(bufs, this.nread));
+    const bufs = engine.buffers;
+    buf = (bufs.length === 1 ? bufs[0] : Buffer.concat(bufs, engine.nread));
   }
-  this.close();
-  if (this._info)
-    this.cb(null, { buffer: buf, engine: this });
+  engine.close();
+  if (engine._info)
+    engine.cb(null, { buffer: buf, engine });
   else
-    this.cb(null, buf);
+    engine.cb(null, buf);
+}
+
+function zlibBufferOnEnd() {
+  zlibBufferFinish(this);
+}
+
+// Called by processCallback(), with the handle as `this`, once the whole input
+// has been processed. Errors have already been reported to the callback, either
+// through the 'error' event or by zlibBufferAddChunk().
+function zlibBufferOnWriteDone(err) {
+  const engine = this[owner_symbol];
+  if (err === undefined && !engine.destroyed)
+    zlibBufferFinish(engine);
 }
 
 function zlibBufferSync(engine, buffer) {
@@ -552,7 +603,10 @@ function processCallback() {
   if (have > 0) {
     const out = self._outBuffer.slice(self._outOffset, self._outOffset + have);
     self._outOffset += have;
-    streamBufferIsFull = !self.push(out);
+    if (self[kBufferResult] === true)
+      zlibBufferAddChunk(self, out);
+    else
+      streamBufferIsFull = !self.push(out);
   } else {
     assert(have === 0, 'have should not go down');
   }
@@ -564,6 +618,8 @@ function processCallback() {
 
   // Exhausted the output buffer, or used all the input create a new one.
   if (availOutAfter === 0 || self._outOffset >= self._chunkSize) {
+    if (availOutAfter === 0 && self[kBufferResult] === true)
+      zlibBufferGrowChunkSize(self);
     handle.availOutBefore = self._chunkSize;
     self._outOffset = 0;
     self._outBuffer = Buffer.allocUnsafe(self._chunkSize);
@@ -618,7 +674,8 @@ function processCallback() {
       return;
     }
 
-    self.push(null);
+    if (self[kBufferResult] !== true)
+      self.push(null);
   }
 
   // Finished with the chunk.
@@ -791,7 +848,8 @@ class Unzip extends Zlib {
   }
 }
 
-function createConvenienceMethod(ctor, sync, prepareOpts) {
+function createConvenienceMethod(ctor, sync, prepareOpts,
+                                 finishInOneWrite = true) {
   if (sync) {
     return function syncBufferWrapper(buffer, opts) {
       return zlibBufferSync(new ctor(opts), buffer);
@@ -805,7 +863,7 @@ function createConvenienceMethod(ctor, sync, prepareOpts) {
     if (prepareOpts !== undefined) {
       opts = prepareOpts(buffer, opts);
     }
-    return zlibBuffer(new ctor(opts), buffer, callback);
+    return zlibBuffer(new ctor(opts), buffer, callback, finishInOneWrite);
   };
 }
 
@@ -1118,7 +1176,11 @@ module.exports = {
   brotliCompressSync: createConvenienceMethod(BrotliCompress, true),
   brotliDecompress: createConvenienceMethod(BrotliDecompress, false),
   brotliDecompressSync: createConvenienceMethod(BrotliDecompress, true),
-  zstdCompress: createConvenienceMethod(ZstdCompress, false, withPledgedSrcSize),
+  // A zstd frame written by a single ZSTD_e_end call can be a few bytes larger
+  // than one written by ZSTD_e_continue followed by ZSTD_e_end, so
+  // zstdCompress() keeps writing the input and ending the frame separately.
+  zstdCompress: createConvenienceMethod(ZstdCompress, false, withPledgedSrcSize,
+                                        false),
   zstdCompressSync: createConvenienceMethod(ZstdCompress, true),
   zstdDecompress: createConvenienceMethod(ZstdDecompress, false),
   zstdDecompressSync: createConvenienceMethod(ZstdDecompress, true),

test/parallel/test-zlib-convenience-async-large.js

@@ -0,0 +1,82 @@
+'use strict';
+// The asynchronous convenience methods grow their output buffer while
+// processing large inputs. Check that large results, custom chunkSize values
+// and maxOutputLength behave exactly like the synchronous methods.
+
+const common = require('../common');
+const assert = require('assert');
+const fs = require('fs');
+const zlib = require('zlib');
+
+const { Z_MIN_CHUNK, BROTLI_PARAM_QUALITY } = zlib.constants;
+
+// Text plus pseudo-random bytes, so that both the compressed and the
+// decompressed output span many output buffers.
+const random = Buffer.alloc(256 * 1024);
+{
+  let seed = 0x2545f491;
+  for (let i = 0; i < random.length; i++) {
+    seed ^= seed << 13;
+    seed ^= seed >>> 17;
+    seed ^= seed << 5;
+    random[i] = seed & 0xff;
+  }
+}
+const data = Buffer.concat([
+  Buffer.alloc(2 * 1024 * 1024, fs.readFileSync(__filename)),
+  random,
+]);
+
+// Keep brotli fast enough for a test.
+const brotliOptions = { params: { [BROTLI_PARAM_QUALITY]: 4 } };
+const pairs = [
+  ['deflate', 'inflate'],
+  ['gzip', 'gunzip'],
+  ['gzip', 'unzip'],
+  ['deflateRaw', 'inflateRaw'],
+  ['brotliCompress', 'brotliDecompress', brotliOptions],
+  ['zstdCompress', 'zstdDecompress'],
+];
+
+function callAsync(method, input, options) {
+  return new Promise((resolve, reject) => {
+    zlib[method](input, options, common.mustCall((err, result) => {
+      if (err) reject(err);
+      else resolve(result);
+    }));
+  });
+}
+
+(async () => {
+  for (const [compress, decompress, options = {}] of pairs) {
+    const label = `${compress}/${decompress}`;
+    const compressed = await callAsync(compress, data, options);
+    const expected = zlib[`${compress}Sync`](data, options);
+    // zstdCompress() pledges the input size and ends the frame separately,
+    // which yields different (not larger) frames than zstdCompressSync().
+    if (compress !== 'zstdCompress') {
+      assert.deepStrictEqual(compressed, expected, label);
+    }
+    assert.deepStrictEqual(await callAsync(decompress, compressed, {}), data,
+                           label);
+
+    for (const chunkSize of [Z_MIN_CHUNK, data.length + 1]) {
+      assert.deepStrictEqual(
+        await callAsync(decompress, compressed, { chunkSize }), data,
+        `${label} chunkSize=${chunkSize}`);
+    }
+
+    for (const maxOutputLength of [data.length - 1, 100000, 1]) {
+      await assert.rejects(
+        callAsync(decompress, compressed, { maxOutputLength }), {
+          name: 'RangeError',
+          code: 'ERR_BUFFER_TOO_LARGE',
+          message: `Cannot create a Buffer larger than ${maxOutputLength} bytes`,
+        }, `${label} maxOutputLength=${maxOutputLength}`);
+    }
+    assert.deepStrictEqual(
+      await callAsync(decompress, compressed,
+                      { maxOutputLength: data.length }),
+      data, label);
+  }
+})().then(common.mustCall());

test/parallel/test-zlib-convenience-async-round-trip.js

@@ -0,0 +1,88 @@
+'use strict';
+// The asynchronous convenience methods know the whole input up front, so they
+// process it with the finish flush in a single write. Check that a call whose
+// output fits in one output buffer takes exactly one threadpool round trip,
+// and that larger results need few round trips.
+
+const common = require('../common');
+const assert = require('assert');
+const async_hooks = require('async_hooks');
+const zlib = require('zlib');
+
+const input = Buffer.from('hello world '.repeat(64));
+// [method, argument, function that turns the result back into `input`]
+const cases = [
+  ['deflate', input, zlib.inflateSync],
+  ['inflate', zlib.deflateSync(input)],
+  ['gzip', input, zlib.gunzipSync],
+  ['gunzip', zlib.gzipSync(input)],
+  ['unzip', zlib.gzipSync(input)],
+  ['deflateRaw', input, zlib.inflateRawSync],
+  ['inflateRaw', zlib.deflateRawSync(input)],
+  ['brotliCompress', input, zlib.brotliDecompressSync],
+  ['brotliDecompress', zlib.brotliCompressSync(input)],
+  ['zstdDecompress', zlib.zstdCompressSync(input)],
+];
+
+// Counts the callbacks into JavaScript (`before` hooks) of each ZLIB resource
+// created while `fn` runs.
+async function countRoundTrips(fn) {
+  const callbacks = new Map();
+  const hook = async_hooks.createHook({
+    init(asyncId, type) {
+      if (type === 'ZLIB') callbacks.set(asyncId, 0);
+    },
+    before(asyncId) {
+      if (callbacks.has(asyncId))
+        callbacks.set(asyncId, callbacks.get(asyncId) + 1);
+    },
+  }).enable();
+  try {
+    await fn();
+  } finally {
+    hook.disable();
+  }
+  return [...callbacks.values()];
+}
+
+function callAsync(method, buffer, options = {}) {
+  return new Promise((resolve, reject) => {
+    zlib[method](buffer, options, common.mustCall((err, result) => {
+      if (err) reject(err);
+      else resolve(result);
+    }));
+  });
+}
+
+(async () => {
+  for (const [method, buffer, undo = (result) => result] of cases) {
+    let result;
+    const roundTrips = await countRoundTrips(async () => {
+      result = await callAsync(method, buffer);
+    });
+    assert.deepStrictEqual(roundTrips, [1], method);
+    assert.deepStrictEqual(undo(result), input, method);
+  }
+
+  // Errors are also reported after a single round trip.
+  let error;
+  const roundTrips = await countRoundTrips(async () => {
+    error = await callAsync('inflate', Buffer.from('not zlib data'))
+      .catch((err) => err);
+  });
+  assert.deepStrictEqual(roundTrips, [1]);
+  assert.strictEqual(error.code, 'Z_DATA_ERROR');
+
+  // The output buffer doubles from the 16 KiB default up to 1 MiB, so 4 MiB of
+  // output takes 6 + 4 round trips instead of 256.
+  {
+    const large = Buffer.alloc(4 * 1024 * 1024, input);
+    const compressed = zlib.deflateSync(large);
+    let result;
+    const roundTrips = await countRoundTrips(async () => {
+      result = await callAsync('inflate', compressed);
+    });
+    assert.deepStrictEqual(roundTrips, [10]);
+    assert.deepStrictEqual(result, large);
+  }
+})().then(common.mustCall());