stream.Transform limitations on backpressure handling after stream is ended

[dmurvihill]

Version

v20.2.0

Platform

Darwin MacBook-Air.local 19.6.0 Darwin Kernel Version 19.6.0: Mon Aug 31 22:12:52 PDT 2020; root:xnu-6153.141.2~1/RELEASE_X86_64 x86_64

Subsystem

stream

What steps will reproduce the bug?

Run this with ts-jest:

import { Readable, Transform, TransformCallback } from 'stream';
import { pipeline } from 'stream/promises';
import NullWritable from 'null-writable';

describe('stream.Transform', () => {
  it('should emit drain events while flushing', async () => {
    const elements = range(3);
    const transform = new DelayTransform()
    const sink = new NullWritable({ objectMode: true, highWaterMark: 1 })
    sink.cork();
    const testPipeline = pipeline(
      Readable.from(elements, { highWaterMark: 1 }),
      transform,
      sink,
    );
    await new Promise(resolve => setTimeout(resolve, 10));
    console.log('sink.uncork()');
    sink.uncork();
    await testPipeline;
  });
});

class DelayTransform extends Transform {
  private readonly storedElements: any = [];

  constructor() {
    super({ objectMode: true, highWaterMark: 1 });
    this.on('data', (d) => { console.log(`emit '${d}'`); });
    this.on('drain', () => { console.log('drain'); });
  }

 _transform(chunk: any, encoding: string | null, callback: TransformCallback) {
   this.storedElements.push(chunk);
   callback();
 }

 _flush(callback: TransformCallback) {
   this.flushStoredElements()
     .then(() => callback())
     .catch((e) => callback(e));
 }

 private async flushStoredElements(): Promise<void> {
   for (const item of this.storedElements) {
     await this.safePush(item);
   }
 }

 /** Push and resolve when it is safe to push again */
 private safePush(item: string): Promise<void> {
   console.log(`push '${item}'`);
   if (this.push(item)) {
     console.log('continue');
     return Promise.resolve();
   } else {
     console.log('wait');
     return new Promise(resolve => {
       console.log('setting drain listener');
       this.once('drain', () => {
         console.log('resuming');
         resolve();
       });
     });
   }
  }
}

function range(n: number): number[] {
  const elements = [];
  for (let i = 0; i < n; i++) {
    elements.push(i);
  }
  return elements;
}

How often does it reproduce? Is there a required condition?

100% reproducible with Jest.

What is the expected behavior? Why is that the expected behavior?

We have two expectations in conflict here:

1. When a stream emits a drain event, it should be safe for the upstream source to continue writing data.
2. When a call to push returns false, the stream implementer should be able to wait for a drain event and then resume pushing data.

These two requirements are incompatible in the _flush() method. Emitting a drain event breaks condition 1 while failing to emit one breaks condition 2.

What do you see instead?

Currently the stream fails to emit a drain event. See output from the test above:

  console.log
    drain

      at DelayTransform.<anonymous> (test/node.spec.ts:29:38)

  console.log
    drain

      at DelayTransform.<anonymous> (test/node.spec.ts:29:38)

  console.log
    push '0'

      at DelayTransform.safePush (test/node.spec.ts:51:12)

  console.log
    emit '0'

      at DelayTransform.<anonymous> (test/node.spec.ts:28:38)

  console.log
    continue

      at DelayTransform.safePush (test/node.spec.ts:53:14)

  console.log
    push '1'

      at DelayTransform.safePush (test/node.spec.ts:51:12)

  console.log
    wait

      at DelayTransform.safePush (test/node.spec.ts:56:14)

  console.log
    setting drain listener

      at test/node.spec.ts:58:16

  console.log
    sink.uncork()

      at Object.<anonymous> (test/node.spec.ts:17:13)

  console.log
    emit '1'

      at DelayTransform.<anonymous> (test/node.spec.ts:28:38)

 FAIL  test/node.spec.ts (7.834 s)
  stream.Transform
    ✕ should emit drain events while flushing (5004 ms)

  ● stream.Transform › should emit drain events while flushing

    thrown: "Exceeded timeout of 5000 ms for a test.
    Add a timeout value to this test to increase the timeout, if this is a long-running test. See https://jestjs.io/docs/api#testname-fn-timeout."

      4 |
      5 | describe('stream.Transform', () => {
    > 6 |   it('should emit drain events while flushing', async () => {
        |   ^
      7 |     const elements = range(3);
      8 |     const transform = new DelayTransform()
      9 |     const sink = new NullWritable({ objectMode: true, highWaterMark: 1 })

      at test/node.spec.ts:6:3
      at Object.<anonymous> (test/node.spec.ts:5:1)

Test Suites: 1 failed, 1 total
Tests:       1 failed, 1 total
Snapshots:   0 total
Time:        7.917 s, estimated 8 s

Additional information

I'm implementing a utility for shuffling large data sets. My stream has to be sure the source has ended before it can emit any values; any data that comes in after the first emit will bias the shuffle. Therefore, all of my calls to push occur during _flush. Would appreciate a workaround if you can think of one.

[mscdex]

FWIW I'm not sure if this is related to what you're running into or is the underlying cause of your problem here, but there is still no event for a Readable stream that gets emitted when it is "safe" to push() again after push() returns false. The 'drain' event for Writable streams is for when it's safe to write() again after write() returns false, which is a separate thing.

I raised an issue about this over 8 years ago but no progress was ever made.

[lukiano]

My stream has to be sure the source has ended before it can emit any values; any data that comes in after the first emit will bias the shuffle. Therefore, all of my calls to push occur during _flush. Would appreciate a workaround if you can think of one.

It looks like it's two streams instead of one. The first one stores data in a buffer (in memory, a file, a database, etc.). The second one executes after the first one finishes and reads data from the buffer to pump it to further destinations.

Regarding the specific point, I just looked at the current implementation of stream.Transform and see the following:

Transform.prototype._write = function(chunk, encoding, callback) {
  const rState = this._readableState;
  const wState = this._writableState;
  const length = rState.length;

  this._transform(chunk, encoding, (err, val) => {
    if (err) {
      callback(err);
      return;
    }

    if (val != null) {
      this.push(val);
    }

    if (
      wState.ended || // Backwards compat.
      length === rState.length || // Backwards compat.
      rState.length < rState.highWaterMark
    ) {
      callback();
    } else {
      this[kCallback] = callback;
    }
  });
};

Transform.prototype._read = function() {
  if (this[kCallback]) {
    const callback = this[kCallback];
    this[kCallback] = null;
    callback();
  }
};

It doesn't check the result of this.push(val); but directly checks if the readable buffer is full. So in safePush() you could have (I think, haven't tested this logic):

if (this.push(item)) {
  console.log('continue');
  return Promise.resolve();
} else {
  console.log('wait');
  return new Promise(resolve => {
    console.log('setting listener');
    // a deferred promise that will be resolved when `_read()` is called to request more data.
    this[kCallback] = resolve;
    // TODO: Maybe listen for error events to reject the promise if the transform is destroyed.
  });
}

[dmurvihill]

It looks like it's two streams instead of one. The first one stores data in a buffer (in memory, a file, a database, etc.). The second one executes after the first one finishes and reads data from the buffer to pump it to further destinations.

I don't disagree that the solution could be represented this way. But I would submit that a shuffle is very plainly a transform operation and that stream should natively support representing it with a Transform. From an API design perspective, users should be able to insert a new ShuffleTransform() into their existing pipeline and expect it to Just Work.

I did actually implement a workaround similar to the one you suggested, based on one proposed by @mscdex in the thread they linked. It's available on NPM as drain-aware-transform.

[IanAtOcucom]

My issue #49938 is for a different use-case but has the same underlying cause.

Transforms should have a "public" non-internals way to be informed that the readable's buffer is no longer full. That the default implementation works this way isn't enough.