RDBC-859 OOM on bulk insert - node.js client

Author: ml054

src/Documents/BulkInsertOperation.ts

@@ -31,6 +31,9 @@ import { BulkInsertOnProgressEventArgs } from "./Session/SessionEvents";
 import * as semaphore from "semaphore";
 import { acquireSemaphore } from "../Utility/SemaphoreUtil";
 import { Buffer } from "node:buffer";
+import { createGzip, Gzip } from "node:zlib";
+import { pipeline } from "readable-stream";
+import { promisify } from "node:util";
 
 class BulkInsertStream {
 
@@ -105,7 +108,7 @@ export class BulkInsertOperation {
     private _bulkInsertExecuteTask: Promise<any>;
     private _bulkInsertExecuteTaskErrored = false;
 
-    private _stream: RequestBodyStream;
+    private _stream: stream.Readable & { flush: (callback?: () => void) => void; };
 
     private _first: boolean = true;
     private _inProgressCommand: CommandType;
@@ -129,8 +132,8 @@ export class BulkInsertOperation {
     private readonly _streamLock: semaphore.Semaphore;
     private _heartbeatCheckInterval = 40_000;
 
-    //TODO:  private GZipStream _compressedStream;
-    private _requestBodyStream: RequestBodyStream; //TODO: raw reableable or wrapped with compressed
+    private _compressedStream: Gzip;
+    private _requestBodyStream: RequestBodyStream;
     private _requestBodyStreamFinished: boolean = false;
     private _currentWriter: BulkInsertStream;
     private _backgroundWriter: BulkInsertStream;
@@ -460,6 +463,12 @@ export class BulkInsertOperation {
                 this._isInitialWrite = false;
 
                 await this._requestBodyStream.flush();
+
+                if (this._compressedStream) {
+                    const flush = promisify(this._compressedStream.flush);
+                    await flush.call(this._compressedStream);
+                }
+
                 this._lastWriteToStream = new Date();
             }
         } finally {
@@ -535,29 +544,23 @@ export class BulkInsertOperation {
     }
 
     private async _ensureStream() {
-        //TODO: sync with c#
-        //TODO if (CompressionLevel != CompressionLevel.NoCompression)
-        //                 _streamExposerContent.Headers.ContentEncoding.Add("gzip");
-
         try {
             this._requestBodyStream = new RequestBodyStream();
-            this._stream = this._requestBodyStream; //TODO:
+            this._stream = this._requestBodyStream;
+
+            if (this.useCompression) {
+                this._compressedStream = createGzip();
+                pipeline(this._requestBodyStream, this._compressedStream);
+            }
+
             const bulkCommand =
-                new BulkInsertCommand(this._operationId, this._requestBodyStream, this._nodeTag, this._options.skipOverwriteIfUnchanged);
+                new BulkInsertCommand(this._operationId, this._compressedStream ?? this._requestBodyStream, this._nodeTag, this._options.skipOverwriteIfUnchanged);
             bulkCommand.useCompression = this._useCompression;
 
             this._bulkInsertExecuteTask = this._requestExecutor.execute(bulkCommand);
 
             this._currentWriter.push("[");
 
-            /* TODO
-              if (CompressionLevel != CompressionLevel.NoCompression)
-            {
-                _compressedStream = new GZipStream(_stream, CompressionLevel, leaveOpen: true);
-                _requestBodyStream = _compressedStream;
-            }
-             */
-
             this._bulkInsertExecuteTask
                 .catch(() => this._bulkInsertExecuteTaskErrored = true);
 
@@ -598,8 +601,7 @@ export class BulkInsertOperation {
                 try {
                     this._currentWriter.push("]");
                     await this._asyncWrite;
-                    this._requestBodyStream.write(this._currentWriter.toBuffer());
-                    //TODO: _compressedStream?.Dispose();
+                    this._requestBodyStream.push(this._currentWriter.toBuffer());
                     await this._stream.flush();
                 } finally {
                     context.dispose();
@@ -1017,7 +1019,7 @@ export class BulkInsertOperation {
 
                     await this._operation.flushIfNeeded();
 
-                    this._operation._currentWriter.push(bytes); //TODO: do we want to stream here?
+                    this._operation._currentWriter.push(bytes);
 
                     await this._operation.flushIfNeeded();
                 } catch (e) {
@@ -1078,8 +1080,12 @@ export class BulkInsertCommand extends RavenCommand<void> {
             + "/bulk_insert?id=" + this._id
             + "&skipOverwriteIfUnchanged=" + (this._skipOverwriteIfUnchanged ? "true" : "false");
 
-        const headers = this._headers().typeAppJson().build();
-        // TODO: useCompression ? new GzipCompressingEntity(_stream) : _stream);
+        const headersBuilder = this._headers().typeAppJson();
+        if (this.useCompression) {
+            headersBuilder.with("Content-Encoding", "gzip");
+        }
+
+        const headers = headersBuilder.build();
         return {
             method: "POST",
             uri,

test/Ported/BulkInsert/BulkInsertsTest.ts

@@ -8,6 +8,8 @@ import { createMetadataDictionary } from "../../../src/Mapping/MetadataAsDiction
 import { CONSTANTS } from "../../../src/Constants";
 import { DateUtil } from "../../../src/Utility/DateUtil";
 import { delay } from "../../../src/Utility/PromiseUtil";
+import { assertThat } from "../../Utils/AssertExtensions";
+import { Timer } from "../../../src/Primitives/Timer";
 
 describe("bulk insert", function () {
 
@@ -20,7 +22,7 @@ describe("bulk insert", function () {
     afterEach(async () =>
         await disposeTestDocumentStore(store));
 
-    it("simple bulk insert should work", async () => {
+    const bulkInsertTest = async (compressed: boolean) => {
         const fooBar1 = new FooBar();
         fooBar1.name = "John Doe";
 
@@ -33,7 +35,9 @@ describe("bulk insert", function () {
         const fooBar4 = new FooBar();
         fooBar4.name = "Mega Jane";
 
-        const bulkInsert = store.bulkInsert();
+        const bulkInsert = store.bulkInsert({
+            useCompression: compressed
+        });
 
         await bulkInsert.store(fooBar1);
         await bulkInsert.store(fooBar2);
@@ -65,8 +69,44 @@ describe("bulk insert", function () {
         } finally {
             session.dispose();
         }
+    }
+
+    it("simple bulk insert should work - no compressed", async () => {
+        await bulkInsertTest(false);
     });
 
+    it("simple bulk insert should work - compressed", async () => {
+        await bulkInsertTest(true);
+    });
+
+    it ("can send heartbeats", async () => {
+
+        const bulkInsert = store.bulkInsert();
+        try {
+
+            // access private variable
+            const timer = (bulkInsert as any)._timer as Timer;
+            const checkInterval = (bulkInsert as any)._heartbeatCheckInterval;
+
+            assertThat(timer)
+                .isNotNull();
+            assertThat(checkInterval)
+                .isEqualTo(40_000);
+
+            timer.change(20, 20);
+            (bulkInsert as any)._heartbeatCheckInterval = 20;
+
+            await bulkInsert.store(new FooBar());
+
+            await delay(250); //it should send heartbeats
+
+            await bulkInsert.store(new FooBar());
+
+        } finally {
+            await bulkInsert.finish();
+        }
+    })
+
     it("can be killed early before making connection", async () => {
         const bulkInsert = store.bulkInsert();
         try {