zlib: finish async convenience calls in one write

lib/zlib.js

@@ -87,6 +87,7 @@ const {
 
 const kFlushFlag = Symbol('kFlushFlag');
 const kError = Symbol('kError');
+const kBufferResult = Symbol('kBufferResult');
 
 const constants = internalBinding('constants').zlib;
 const {
@@ -124,7 +125,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 +137,81 @@ 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);
+}
+
 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 +590,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');
   }
@@ -618,7 +659,8 @@ function processCallback() {
       return;
     }
 
-    self.push(null);
+    if (self[kBufferResult] !== true)
+      self.push(null);
   }
 
   // Finished with the chunk.
@@ -791,7 +833,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 +848,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 +1161,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-round-trip.js

@@ -0,0 +1,74 @@
+'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.
+
+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');
+})().then(common.mustCall());