Repository: microsoft/vscode (TypeScript) — Bug Fix

Issue:
Write of large buffer in WebSocketNodeSocket can hold NodeJS thread too long
<!-- ⚠️⚠️ Do Not Delete This! bug_report_template ⚠️⚠️ -->
<!-- Please read our Rules of Conduct: https://opensource.microsoft.com/codeofconduct/ -->
<!-- 🕮 Read our guide about submitting issues: https://github.com/microsoft/vscode/wiki/Submitting-Bugs-and-Suggestions -->
<!-- 🔎 Search existing issues to avoid creating duplicates. -->
<!-- 🧪 Test using the latest Insiders build to see if your issue has already been fixed: https://code.visualstudio.com/insiders/ -->
<!-- 💡 Instead of creating your report here, use 'Report Issue' from the 'Help' menu in VS Code to pre-fill useful information. -->
<!-- 🔧 Launch with `code --disable-extensions` to check. -->
Does this issue occur when all extensions are disabled?: Yes/No

<!-- 🪓 If you answered No above, use 'Help: Start Extension Bisect' from Command Palette to try to identify the cause. -->
<!-- 📣 Issues caused by an extension need to be reported directly to the extension publisher. The 'Help > Report Issue' dialog can assist with this. -->
- VS Code Version: 
- OS Version: 

In node/ipc.net.ts, a write of a large buffer during remoting can hang the NodeJS thread since the call is synchronous.  There should be a cap on the size of a particular message, and if the buffer is larger then the buffer is sent in chunks.



Discussion:

Resolving patch:
diff --git a/src/vs/base/parts/ipc/node/ipc.net.ts b/src/vs/base/parts/ipc/node/ipc.net.ts
--- a/src/vs/base/parts/ipc/node/ipc.net.ts
+++ b/src/vs/base/parts/ipc/node/ipc.net.ts
@@ -189,7 +189,15 @@ export class NodeSocket implements ISocket {
 }
 
 const enum Constants {
-	MinHeaderByteSize = 2
+	MinHeaderByteSize = 2,
+	/**
+	 * If we need to write a large buffer, we will split it into 256KB chunks and
+	 * send each chunk as a websocket message. This is to prevent that the sending
+	 * side is stuck waiting for the entire buffer to be compressed before writing
+	 * to the underlying socket or that the receiving side is stuck waiting for the
+	 * entire message to be received before processing the bytes.
+	 */
+	MaxWebSocketMessageLength = 256 * 1024 // 256 KB
 }
 
 const enum ReadState {
@@ -272,7 +280,14 @@ export class WebSocketNodeSocket extends Disposable implements ISocket, ISocketT
 		}));
 		this._incomingData = new ChunkStream();
 		this._register(this.socket.onData(data => this._acceptChunk(data)));
-		this._register(this.socket.onClose((e) => this._onClose.fire(e)));
+		this._register(this.socket.onClose(async (e) => {
+			// Delay surfacing the close event until the async inflating is done
+			// and all data has been emitted
+			if (this._flowManager.isProcessingReadQueue()) {
+				await Event.toPromise(this._flowManager.onDidFinishProcessingReadQueue);
+			}
+			this._onClose.fire(e);
+		}));
 	}
 
 	public override dispose(): void {
@@ -300,7 +315,23 @@ export class WebSocketNodeSocket extends Disposable implements ISocket, ISocketT
 	}
 
 	public write(buffer: VSBuffer): void {
-		this._flowManager.writeMessage(buffer);
+		// If we write many logical messages (let's say 1000 messages of 100KB) during a single process tick, we do
+		// this thing where we install a process.nextTick timer and group all of them together and we then issue a
+		// single WebSocketNodeSocket.write with a 100MB buffer.
+		//
+		// The first problem is that the actual writing to the underlying node socket will only happen after all of
+		// the 100MB have been deflated (due to waiting on zlib flush). The second problem is on the reading side,
+		// where we will get a single WebSocketNodeSocket.onData event fired when all the 100MB have arrived,
+		// delaying processing the 1000 received messages until all have arrived, instead of processing them as each
+		// one arrives.
+		//
+		// We therefore split the buffer into chunks, and issue a write for each chunk.
+
+		let start = 0;
+		while (start < buffer.byteLength) {
+			this._flowManager.writeMessage(buffer.slice(start, Math.min(start + Constants.MaxWebSocketMessageLength, buffer.byteLength)));
+			start += Constants.MaxWebSocketMessageLength;
+		}
 	}
 
 	private _write(buffer: VSBuffer, compressed: boolean): void {
@@ -465,6 +496,9 @@ class WebSocketFlowManager extends Disposable {
 	private readonly _writeQueue: VSBuffer[] = [];
 	private readonly _readQueue: { data: VSBuffer; isCompressed: boolean; isLastFrameOfMessage: boolean }[] = [];
 
+	private readonly _onDidFinishProcessingReadQueue = this._register(new Emitter<void>());
+	public readonly onDidFinishProcessingReadQueue = this._onDidFinishProcessingReadQueue.event;
+
 	private readonly _onDidFinishProcessingWriteQueue = this._register(new Emitter<void>());
 	public readonly onDidFinishProcessingWriteQueue = this._onDidFinishProcessingWriteQueue.event;
 
@@ -565,6 +599,11 @@ class WebSocketFlowManager extends Disposable {
 			}
 		}
 		this._isProcessingReadQueue = false;
+		this._onDidFinishProcessingReadQueue.fire();
+	}
+
+	public isProcessingReadQueue(): boolean {
+		return (this._isProcessingReadQueue);
 	}
 
 	/**


Test patch:
diff --git a/src/vs/base/parts/ipc/test/node/ipc.net.test.ts b/src/vs/base/parts/ipc/test/node/ipc.net.test.ts
--- a/src/vs/base/parts/ipc/test/node/ipc.net.test.ts
+++ b/src/vs/base/parts/ipc/test/node/ipc.net.test.ts
@@ -5,7 +5,7 @@
 
 import * as assert from 'assert';
 import { EventEmitter } from 'events';
-import { createServer, Socket } from 'net';
+import { AddressInfo, connect, createServer, Server, Socket } from 'net';
 import { tmpdir } from 'os';
 import { Barrier, timeout } from 'vs/base/common/async';
 import { VSBuffer } from 'vs/base/common/buffer';
@@ -706,4 +706,63 @@ suite('WebSocketNodeSocket', () => {
 			assert.deepStrictEqual(actual, 'Helloworld');
 		});
 	});
+
+	test('Large buffers are split and sent in chunks', async () => {
+
+		let receivingSideOnDataCallCount = 0;
+		let receivingSideTotalBytes = 0;
+		const receivingSideSocketClosedBarrier = new Barrier();
+
+		const server = await listenOnRandomPort((socket) => {
+			// stop the server when the first connection is received
+			server.close();
+
+			const webSocketNodeSocket = new WebSocketNodeSocket(new NodeSocket(socket), true, null, false);
+			webSocketNodeSocket.onData((data) => {
+				receivingSideOnDataCallCount++;
+				receivingSideTotalBytes += data.byteLength;
+			});
+
+			webSocketNodeSocket.onClose(() => {
+				webSocketNodeSocket.dispose();
+				receivingSideSocketClosedBarrier.open();
+			});
+		});
+
+		const socket = connect({
+			host: '127.0.0.1',
+			port: (<AddressInfo>server.address()).port
+		});
+
+		const buff = generateRandomBuffer(1 * 1024 * 1024);
+
+		const webSocketNodeSocket = new WebSocketNodeSocket(new NodeSocket(socket), true, null, false);
+		webSocketNodeSocket.write(buff);
+		await webSocketNodeSocket.drain();
+		webSocketNodeSocket.dispose();
+		await receivingSideSocketClosedBarrier.wait();
+
+		assert.strictEqual(receivingSideTotalBytes, buff.byteLength);
+		assert.strictEqual(receivingSideOnDataCallCount, 4);
+	});
+
+	function generateRandomBuffer(size: number): VSBuffer {
+		const buff = VSBuffer.alloc(size);
+		for (let i = 0; i < size; i++) {
+			buff.writeUInt8(Math.floor(256 * Math.random()), i);
+		}
+		return buff;
+	}
+
+	function listenOnRandomPort(handler: (socket: Socket) => void): Promise<Server> {
+		return new Promise((resolve, reject) => {
+			const server = createServer(handler).listen(0);
+			server.on('listening', () => {
+				resolve(server);
+			});
+			server.on('error', (err) => {
+				reject(err);
+			});
+		});
+	}
 });