@naturalcycles/nodejs-lib 12.58.0 → 12.62.0

package/src/stream/ndjson/ndjsonMap.ts

@@ -1,6 +1,6 @@
 import { createReadStream, createWriteStream } from 'fs'
 import { createGzip, createUnzip } from 'zlib'
-import { AsyncMapper, ErrorMode } from '@naturalcycles/js-lib'
+import { AbortableAsyncMapper, ErrorMode } from '@naturalcycles/js-lib'
 import {
   requireFileToExist,
   transformJsonParse,
@@ -41,7 +41,7 @@ export interface NDJSONMapOptions<IN = any, OUT = IN>
  * Zips output file automatically, if it ends with `.gz`.
  */
 export async function ndjsonMap<IN = any, OUT = any>(
-  mapper: AsyncMapper<IN, OUT>,
+  mapper: AbortableAsyncMapper<IN, OUT>,
   opt: NDJSONMapOptions<IN, OUT>,
 ): Promise<void> {
   const { inputFilePath, outputFilePath, logEveryOutput = 100_000, limitInput, limitOutput } = opt
@@ -56,19 +56,21 @@ export async function ndjsonMap<IN = any, OUT = any>(
   const transformUnzip = inputFilePath.endsWith('.gz') ? [createUnzip()] : []
   const transformZip = outputFilePath.endsWith('.gz') ? [createGzip()] : []
 
+  const readable = createReadStream(inputFilePath)
+
   await _pipeline([
-    createReadStream(inputFilePath),
+    readable,
     ...transformUnzip,
     transformSplit(), // splits by \n
     transformJsonParse(),
-    transformLimit(limitInput),
+    transformLimit({ limit: limitInput, sourceReadable: readable }),
     transformLogProgress({ metric: 'read', ...opt }),
     transformMap(mapper, {
       flattenArrayOutput: true,
       errorMode: ErrorMode.SUPPRESS,
       ...opt,
     }),
-    transformLimit(limitOutput),
+    transformLimit({ limit: limitOutput, sourceReadable: readable }),
     transformLogProgress({ metric: 'saved', logEvery: logEveryOutput }),
     transformToNDJson(),
     ...transformZip,

package/src/stream/ndjson/ndjsonStreamForEach.ts

@@ -1,6 +1,6 @@
 import * as fs from 'fs'
 import { createUnzip } from 'zlib'
-import { AsyncMapper, ErrorMode } from '@naturalcycles/js-lib'
+import { AbortableAsyncMapper, ErrorMode } from '@naturalcycles/js-lib'
 import {
   requireFileToExist,
   transformJsonParse,
@@ -23,7 +23,7 @@ export interface NDJSONStreamForEachOptions<IN = any>
  * Convenience function to `forEach` through an ndjson file.
  */
 export async function ndjsonStreamForEach<T>(
-  mapper: AsyncMapper<T, void>,
+  mapper: AbortableAsyncMapper<T, void>,
   opt: NDJSONStreamForEachOptions<T>,
 ): Promise<void> {
   requireFileToExist(opt.inputFilePath)

package/src/stream/ndjson/transformJsonParse.ts

@@ -33,17 +33,17 @@ export function transformJsonParse<OUT = any>(
   const { strict = true, reviver } = opt
 
   return new Transform({
-    objectMode: false,
+    writableObjectMode: false,
     readableObjectMode: true,
-    transform(chunk: string, _encoding, cb) {
+    transform(chunk: string, _, cb) {
       try {
         const data = JSON.parse(chunk, reviver)
         cb(null, data)
       } catch (err) {
-        // console.error(err)
         if (strict) {
           cb(err as Error) // emit error
         } else {
+          console.error(err)
           cb() // emit no error, but no result neither
         }
       }

package/src/stream/ndjson/transformToNDJson.ts

@@ -37,9 +37,9 @@ export function transformToNDJson<IN = any>(
   const { strict = true, separator = '\n', sortObjects = false, useFlatstr = false } = opt
 
   return new Transform({
-    objectMode: true,
+    writableObjectMode: true,
     readableObjectMode: false,
-    transform(chunk: IN, _encoding, cb) {
+    transform(chunk: IN, _, cb) {
       try {
         if (sortObjects) {
           chunk = _sortObjectDeep(chunk as any)

package/src/stream/pipeline/pipeline.ts

@@ -1,16 +1,109 @@
-import { pipeline } from 'stream'
-import { promisify } from 'util'
+import { pipeline, Readable, Transform, Writable } from 'stream'
+import { _last, AnyFunction, DeferredPromise, pDefer } from '@naturalcycles/js-lib'
+import { writablePushToArray } from '../../index'
 
 type AnyStream = NodeJS.ReadableStream | NodeJS.WritableStream | NodeJS.ReadWriteStream
 
+// /**
+//  * Promisified stream.pipeline()
+//  */
+// export let _pipeline = promisify(pipeline)
+//
+// // Workaround https://github.com/nodejs/node/issues/40191
+// // todo: remove it when fix is released in 16.x and in AppEngine 16.x
+// if (process.version >= 'v16.10') {
+//   const { pipeline } = require('stream/promises')
+//   _pipeline = ((streams: AnyStream[]) => pipeline(...streams)) as any
+// }
+
+export interface PipelineOptions {
+  /**
+   * Set to true to allow ERR_STREAM_PREMATURE_CLOSE.
+   * Required to support graceful close when using transformLimit
+   */
+  allowClose?: boolean
+}
+
+/**
+ * Promisified `stream.pipeline`.
+ *
+ * Supports opt.allowClose, which allows transformLimit to work (to actually stop source Readable)
+ * without throwing an error (ERR_STREAM_PREMATURE_CLOSE).
+ */
+export async function _pipeline(streams: AnyStream[], opt: PipelineOptions = {}): Promise<void> {
+  const first = streams[0] as any
+  const rest = streams.slice(1)
+
+  if (opt.allowClose) {
+    // Do the magic of making the pipeline "abortable"
+    //
+    // How does it work:
+    // It finds `sourceReadable` (basically, it's just first item in the passed array of streams)
+    // Finds last "writable" (last item), patches the `_final` method of it to detect when the whole pipeline is "done",
+    // sets the `streamDone` DeferredPromise that resolves when the pipeline is done.
+    // Scans through all passed items, finds those that are capable of "closing" the stream
+    // (currently its `transformLimit` or `transformMap`)
+    // Patches them by attaching `sourceReadable` and `streamDone`.
+    // These items (transformLimit and transformMap), when they need to "close the stream" - call `pipelineClose`.
+    // `pipelineClose` is the result of 2 sleepless nights of googling and experimentation:)
+    // It does:
+    // 1. Stops the "downstream" by doing `this.push(null)`.
+    // 2. Pauses the `sourceReadable` by calling sourceReadable.unpipe()
+    // 3. Waits for `streamDone` to ensure that downstream chunks are fully processed (e.g written to disk).
+    // 4. Calls `sourceReadable.destroy()`, which emits ERR_STREAM_PREMATURE_CLOSE
+    // 5. _pipeline (this function) catches that specific error and suppresses it (because it's expected and
+    // inevitable in this flow). Know a better way to close the stream? Tell me!
+    const streamDone = pDefer()
+    const sourceReadable = first as Readable
+    const last = _last(streams) as Writable
+    const lastFinal = last._final?.bind(last) || ((cb: AnyFunction) => cb())
+    last._final = cb => {
+      lastFinal(() => {
+        cb()
+        streamDone.resolve()
+      })
+    }
+
+    rest.forEach(s => {
+      // console.log(s)
+      if (s instanceof AbortableTransform || s.constructor.name === 'DestroyableTransform') {
+        // console.log(`found ${s.constructor.name}, setting props`)
+        ;(s as AbortableTransform).sourceReadable = sourceReadable
+        ;(s as AbortableTransform).streamDone = streamDone
+      }
+    })
+  }
+
+  return new Promise<void>((resolve, reject) => {
+    pipeline(first, ...(rest as any[]), (err: Error) => {
+      if (err) {
+        if (opt.allowClose && (err as any)?.code === 'ERR_STREAM_PREMATURE_CLOSE') {
+          console.log('_pipeline closed (as expected)')
+          return resolve()
+        }
+        // console.log(`_pipeline error`, err)
+        return reject(err)
+      }
+
+      resolve()
+    })
+  })
+}
+
 /**
- * Promisified stream.pipeline()
+ * Convenience function to make _pipeline collect all items at the end of the stream (should be Transform, not Writeable!)
+ * and return.
  */
-export let _pipeline = promisify(pipeline)
+export async function _pipelineToArray<T>(
+  streams: AnyStream[],
+  opt: PipelineOptions = {},
+): Promise<T[]> {
+  const a: T[] = []
+  await _pipeline([...streams, writablePushToArray(a)], opt)
+  return a
+}
 
-// Workaround https://github.com/nodejs/node/issues/40191
-// todo: remove it when fix is released in 16.x and in AppEngine 16.x
-if (process.version >= 'v16.10') {
-  const { pipeline } = require('stream/promises')
-  _pipeline = ((streams: AnyStream[]) => pipeline(...streams)) as any
+export class AbortableTransform extends Transform {
+  sourceReadable?: Readable
+  streamDone?: DeferredPromise
 }

package/src/stream/readable/readableCreate.ts

@@ -6,6 +6,14 @@ import { ReadableTyped } from '../stream.model'
  * Push `null` to it to complete (similar to RxJS `.complete()`).
  *
  * Difference from Readable.from() is that this readable is not "finished" yet and allows pushing more to it.
+ *
+ * Caution!
+ * The implementation of this Readable is not fully compliant,
+ * e.g the read() method doesn't return anything, so, it will hand the Node process (or cause it to process.exit(0))
+ * if read() will be called AFTER everything was pushed and Readable is closed (by pushing `null`).
+ * Beware of it when e.g doing unit testing! Jest prefers to hang (not exit-0).
+ *
+ * @deprecated because of the caution above
  */
 export function readableCreate<T>(
   items: Iterable<T> = [],
@@ -14,7 +22,7 @@ export function readableCreate<T>(
   const readable = new Readable({
     objectMode: true,
     ...opt,
-    read() {},
+    read() {}, // Caution, if this is called and Readable has not finished yet (null wasn't pushed) - it'll hang the process!
   })
   for (const item of items) {
     readable.push(item)

package/src/stream/readable/readableForEach.ts

@@ -1,4 +1,4 @@
-import { AsyncMapper, Mapper, _passNothingPredicate } from '@naturalcycles/js-lib'
+import { Mapper, _passNothingPredicate, AbortableAsyncMapper } from '@naturalcycles/js-lib'
 import { ReadableTyped, _pipeline } from '../../index'
 import { transformMap, TransformMapOptions } from '../transform/transformMap'
 
@@ -10,7 +10,7 @@ import { transformMap, TransformMapOptions } from '../transform/transformMap'
  */
 export async function readableForEach<T>(
   readable: ReadableTyped<T>,
-  mapper: AsyncMapper<T, void>,
+  mapper: AbortableAsyncMapper<T, void>,
   opt: TransformMapOptions<T, void> = {},
 ): Promise<void> {
   await _pipeline([

package/src/stream/readable/readableFromArray.ts

@@ -1,5 +1,5 @@
 import { Readable, ReadableOptions } from 'stream'
-import { AsyncMapper, pMap, _passthroughMapper } from '@naturalcycles/js-lib'
+import { _passthroughMapper, AbortableAsyncMapper } from '@naturalcycles/js-lib'
 import { ReadableTyped } from '../stream.model'
 
 /**
@@ -10,29 +10,26 @@ import { ReadableTyped } from '../stream.model'
  */
 export function readableFromArray<IN, OUT>(
   items: IN[],
-  mapper: AsyncMapper<IN, OUT> = _passthroughMapper,
+  mapper: AbortableAsyncMapper<IN, OUT> = _passthroughMapper,
   opt?: ReadableOptions,
 ): ReadableTyped<OUT> {
-  const readable = new Readable({
+  let i = -1
+
+  return new Readable({
     objectMode: true,
     ...opt,
-    read() {},
-  })
-
-  void pMap(
-    items,
-    async (item, index) => {
-      readable.push(await mapper(item, index))
+    async read() {
+      i++
+      if (i < items.length) {
+        try {
+          this.push(await mapper(items[i]!, i))
+        } catch (err) {
+          console.error(err)
+          this.destroy(err as Error)
+        }
+      } else {
+        this.push(null) // end
+      }
     },
-    { concurrency: 1 },
-  )
-    .then(() => {
-      readable.push(null) // done
-    })
-    .catch(err => {
-      console.error(err)
-      readable.push(err)
-    })
-
-  return readable
+  })
 }

package/src/stream/readable/readableMap.ts

@@ -1,28 +1,31 @@
-import { AsyncMapper } from '@naturalcycles/js-lib'
-import { readableCreate } from '../../index'
+import { Transform } from 'stream'
+import { AbortableAsyncMapper, SKIP } from '@naturalcycles/js-lib'
 import { ReadableTyped } from '../stream.model'
 
 export function readableMap<IN, OUT>(
   readable: ReadableTyped<IN>,
-  mapper: AsyncMapper<IN, OUT>,
+  mapper: AbortableAsyncMapper<IN, OUT>,
 ): ReadableTyped<OUT> {
-  const out = readableCreate<OUT>()
+  let i = -1
 
-  void (async () => {
-    try {
-      let index = 0
-      for await (const item of readable) {
-        const v = await mapper(item, index++)
-        out.push(v)
-      }
-
-      // We're done
-      out.push(null)
-    } catch (err) {
-      console.error(err)
-      out.emit('error', err)
-    }
-  })()
-
-  return out
+  // todo: check if we need to handle errors somehow specifically
+  return readable.pipe(
+    new Transform({
+      objectMode: true,
+      async transform(chunk, _enc, cb) {
+        try {
+          const r = await mapper(chunk, ++i)
+          if (r === SKIP) {
+            cb()
+          } else {
+            // _assert(r !== END, `readableMap END not supported`)
+            cb(null, r)
+          }
+        } catch (err) {
+          console.error(err)
+          cb(err as Error)
+        }
+      },
+    }),
+  )
 }

package/src/stream/sizeStack.ts

@@ -0,0 +1,56 @@
+import { _hb, CommonLogger, NumberStack } from '@naturalcycles/js-lib'
+import { yellow } from '../colors'
+import { gzipBuffer } from '../util/zip.util'
+
+export class SizeStack extends NumberStack {
+  constructor(public name: string, size: number) {
+    super(size)
+  }
+
+  total = 0
+
+  override push(item: any): this {
+    this.total += item
+    return super.push(item)
+  }
+
+  getStats(): string {
+    // const pcs = this.percentiles([50, 90])
+
+    return [
+      '  ' + this.name,
+      'avg',
+      yellow(_hb(this.avg())),
+      // 'p50',
+      // yellow(_hb(pcs[50])),
+      // 'p90',
+      // yellow(_hb(pcs[90])),
+      'total',
+      yellow(_hb(this.total)),
+    ].join(' ')
+  }
+
+  static async countItem(
+    item: any,
+    logger: CommonLogger,
+    sizes?: SizeStack,
+    sizesZipped?: SizeStack,
+  ): Promise<void> {
+    if (!sizes) return
+
+    // try-catch, because we don't want to fail the pipeline on logProgress
+    try {
+      const buf = Buffer.from(JSON.stringify(item))
+      sizes.push(buf.byteLength)
+
+      if (sizesZipped) {
+        const { byteLength } = await gzipBuffer(buf)
+        sizesZipped.push(byteLength)
+      }
+    } catch (err) {
+      logger.warn(
+        `transformLogProgress failed to JSON.stringify the chunk: ${(err as Error).message}`,
+      )
+    }
+  }
+}

package/src/stream/stream.util.ts

@@ -0,0 +1,29 @@
+import { Readable } from 'stream'
+import { CommonLogger } from '@naturalcycles/js-lib'
+
+export function pipelineClose(
+  name: string,
+  readableDownstream: Readable,
+  sourceReadable: Readable | undefined,
+  streamDone: Promise<void> | undefined,
+  logger: CommonLogger,
+): void {
+  readableDownstream.push(null) // this closes the stream, so downstream Readable will receive `end` and won't write anything
+
+  if (!sourceReadable) {
+    logger.warn(`${name} sourceReadable is not provided, readable stream will not be stopped`)
+  } else {
+    logger.log(`${name} is calling readable.unpipe() to pause the stream`)
+    sourceReadable.unpipe() // it is expected to pause the stream
+
+    if (!streamDone) {
+      logger.log(`${name} streamDone is not provided, will do readable.destroy right away`)
+      sourceReadable.destroy()
+    } else {
+      void streamDone.then(() => {
+        logger.log(`${name} streamDone, calling readable.destroy()`)
+        sourceReadable.destroy() // this throws ERR_STREAM_PREMATURE_CLOSE
+      })
+    }
+  }
+}

package/src/stream/transform/transformBuffer.ts

@@ -18,7 +18,7 @@ export function transformBuffer<IN = any>(opt: TransformBufferOptions): Transfor
   return new Transform({
     objectMode: true,
     ...opt,
-    transform(chunk, _encoding, cb) {
+    transform(chunk, _, cb) {
       buf.push(chunk)
 
       if (buf.length >= batchSize) {

package/src/stream/transform/transformFilter.ts

@@ -1,32 +1,18 @@
 import { Transform } from 'stream'
 import { AsyncPredicate, Predicate } from '@naturalcycles/js-lib'
 import { TransformOptions, TransformTyped } from '../stream.model'
+import { transformMap, TransformMapOptions } from './transformMap'
 
 /**
- * Note, that currently it's NOT concurrent! (concurrency = 1)
- * So, it's recommended to use transformMap instead, that is both concurrent and has
- * filtering feature by default.
+ * Just a convenience wrapper around `transformMap` that has built-in predicate filtering support.
  */
 export function transformFilter<IN = any>(
   predicate: AsyncPredicate<IN>,
-  opt: TransformOptions = {},
+  opt: TransformMapOptions = {},
 ): TransformTyped<IN, IN> {
-  let index = 0
-
-  return new Transform({
-    objectMode: true,
+  return transformMap(v => v, {
+    predicate,
     ...opt,
-    async transform(chunk: IN, _encoding, cb) {
-      try {
-        if (await predicate(chunk, index++)) {
-          cb(null, chunk) // pass through
-        } else {
-          cb() // signal that we've finished processing, but emit no output here
-        }
-      } catch (err) {
-        cb(err as Error)
-      }
-    },
   })
 }
 
@@ -42,7 +28,7 @@ export function transformFilterSync<IN = any>(
   return new Transform({
     objectMode: true,
     ...opt,
-    async transform(chunk: IN, _encoding, cb) {
+    transform(chunk: IN, _, cb) {
       try {
         if (predicate(chunk, index++)) {
           cb(null, chunk) // pass through

package/src/stream/transform/transformLimit.ts

@@ -1,38 +1,84 @@
-import { Transform } from 'stream'
+import { Readable } from 'stream'
 import { CommonLogger } from '@naturalcycles/js-lib'
+import { AbortableTransform, transformNoOp } from '../../index'
 import { TransformOptions, TransformTyped } from '../stream.model'
+import { pipelineClose } from '../stream.util'
 
 export interface TransformLimitOptions extends TransformOptions {
+  /**
+   * Nullish value (e.g 0 or undefined) would mean "no limit"
+   */
+  limit?: number
+
+  /**
+   * If provided (recommended!) - it will call readable.destroy() on limit.
+   * Without it - it will only stop the downstream consumers, but won't stop
+   * the Readable ("source" of the stream).
+   * It is almost always crucial to stop the Source too, so, please provide the Readable here!
+   */
+  sourceReadable?: Readable
+
+  /**
+   * Please provide it (a Promise that resolves when the Stream is done, e.g finished consuming things)
+   * to be able to wait for Consumers before calling `readable.destroy`.
+   * Has no effect if `readable` is not provided.
+   */
+  streamDone?: Promise<void>
+
   logger?: CommonLogger
+
+  /**
+   * Set to true to enable additional debug messages, e.g it'll log
+   * when readable still emits values after the limit is reached.
+   */
+  debug?: boolean
 }
 
+/**
+ * Class only exists to be able to do `instanceof TransformLimit`
+ * and to set sourceReadable+streamDone to it in `_pipeline`.
+ */
+export class TransformLimit extends AbortableTransform {}
+
 /**
  * 0 or falsy value means "no limit"
  */
-export function transformLimit<IN>(
-  limit?: number,
-  opt: TransformLimitOptions = {},
-): TransformTyped<IN, IN> {
-  const { logger = console } = opt
-  let index = 0
-  let ended = false
+export function transformLimit<IN>(opt: TransformLimitOptions = {}): TransformTyped<IN, IN> {
+  const { logger = console, limit, debug } = opt
 
-  return new Transform({
+  if (!limit) {
+    // No limit - returning pass-through transform
+    return transformNoOp()
+  }
+
+  let i = 0 // so we start first chunk with 1
+  let ended = false
+  return new TransformLimit({
     objectMode: true,
     ...opt,
-    transform(this: Transform, chunk: IN, _encoding, cb) {
-      index++
-
-      if (!ended) {
-        cb(null, chunk) // pass through the item
-      } else {
-        cb(null) // pass-through empty
-      }
+    transform(this: TransformLimit, chunk, _, cb) {
+      i++
 
-      if (limit && index === limit) {
+      if (i === limit) {
         ended = true
-        logger.log(`transformLimit: limit of ${limit} reached`)
-        // this.emit('end') // this makes it "halt" on Node 14 lts
+        logger.log(`transformLimit of ${limit} reached`)
+        this.push(chunk)
+
+        pipelineClose(
+          'transformLimit',
+          this,
+          opt.sourceReadable || this.sourceReadable,
+          opt.streamDone || this.streamDone,
+          logger,
+        )
+
+        cb() // after pause
+      } else if (!ended) {
+        cb(null, chunk)
+      } else {
+        if (debug) logger.log(`transformLimit.transform after limit`, i)
+        // If we ever HANG (don't call cb) - Node will do process.exit(0) to us
+        cb() // ended, don't emit anything
       }
     },
   })