Merge pull request #1057 from bbassie/fix/bounded-send-buffer

Limit unacknowledged messages buffered for a connection

docs/configuration.md

@@ -165,6 +165,17 @@ Must be greater than 0.
 Defaults to 60.
 
 
+### controller.max_send_buffer_size
+
+Size in MiB of messages sent to a connection that it has not yet acknowledged before the controller ends its session.
+The controller keeps these so they can be resent if the connection drops and resumes.
+A client that stops acknowledging messages, or falls far behind, gets disconnected and starts a new session when it reconnects.
+Changes apply to new sessions.
+Set to 0 for no limit.
+
+Defaults to 64.
+
+
 ### controller.metrics_timeout
 
 Timeout in seconds before a call to gather metrics from a host times out.

packages/controller/src/WsServer.ts

@@ -226,7 +226,8 @@ ${err.stack}`
 		let sessionToken = jwt.sign({ aud: this.sessionAud, sid: sessionId }, this.controller.authSecret);
 		let sessionTimeout = this.controller.config.get("controller.session_timeout");
 		let heartbeatInterval = this.controller.config.get("controller.heartbeat_interval");
-		let connector = new WsServerConnector(dst, sessionId, sessionTimeout, heartbeatInterval);
+		let maxSendBufferSize = this.controller.config.get("controller.max_send_buffer_size") * 1024 * 1024;
+		let connector = new WsServerConnector(dst, sessionId, sessionTimeout, heartbeatInterval, maxSendBufferSize);
 		this.activeConnectors.set(sessionId, connector);
 		connector.on("close", () => {
 			this.activeConnectors.delete(sessionId);

packages/controller/src/WsServerConnector.ts

@@ -21,8 +21,10 @@ export default class WsServerConnector extends lib.WebSocketBaseConnector {
 		private _sessionId: number,
 		private _sessionTimeout: number,
 		public _heartbeatInterval: number,
+		maxSendBufferSize = 0,
 	) {
 		super(new lib.Address(lib.Address.controller, 0), dst);
+		this.maxSendBufferSize = maxSendBufferSize;
 
 		// The following states are used in the server connector
 		// closed: Connection is closed
@@ -110,7 +112,7 @@ export default class WsServerConnector extends lib.WebSocketBaseConnector {
 			`Connector | resuming session with ${this.dst}, resending ${this._sendBuffer.length} buffered messages`
 		);
 		for (let message of this._sendBuffer) {
-			this._socket.send(JSON.stringify(message));
+			this._socket.send(message.text);
 		}
 		this.emit("resume");
 	}

packages/lib/src/config/definitions.ts

@@ -25,6 +25,7 @@ export interface ControllerConfigFields {
 	"controller.auth_secret": string | null;
 	"controller.heartbeat_interval": number;
 	"controller.session_timeout": number;
+	"controller.max_send_buffer_size": number;
 	"controller.metrics_timeout": number;
 	"controller.max_subscription_filters": number;
 	"controller.system_metrics_interval": number;
@@ -202,6 +203,14 @@ export class ControllerConfig extends classes.Config<ControllerConfigFields> {
 			initialValue: 60,
 			validator: validators.greaterThanZero,
 		},
+		"controller.max_send_buffer_size": {
+			title: "Max Send Buffer Size",
+			description:
+				"MiB of unacknowledged messages to keep for a connection before ending its session, 0 for no limit.",
+			type: "number",
+			initialValue: 64,
+			validator: validators.greaterThanEqualZero,
+		},
 		"controller.metrics_timeout": {
 			title: "Metrics Timeout",
 			description: "Timeout in seconds for metrics gathering from hosts.",

packages/lib/src/link/connectors.ts

@@ -171,7 +171,11 @@ export abstract class WebSocketBaseConnector<
 	_heartbeatId: ReturnType<typeof setInterval> | null = null;
 	_heartbeatInterval: number | null = null;
 	_lastReceivedSeq = undefined;
-	_sendBuffer: (libData.MessageRoutable)[] = [];
+	/** Serialised messages sent but not yet acknowledged by the other side */
+	_sendBuffer: { seq: number, text: string }[] = [];
+	_sendBufferSize = 0;
+	/** Session is ended if the unacknowledged messages exceed this many characters, 0 for no limit */
+	maxSendBufferSize = 0;
 
 	_reset() {
 		this._state = "closed";
@@ -182,6 +186,7 @@ export abstract class WebSocketBaseConnector<
 		this._heartbeatInterval = null;
 		this._lastReceivedSeq = undefined;
 		this._sendBuffer.length = 0;
+		this._sendBufferSize = 0;
 		super._reset();
 	}
 
@@ -194,6 +199,7 @@ export abstract class WebSocketBaseConnector<
 		this._heartbeatInterval = null;
 		this._lastReceivedSeq = undefined;
 		this._sendBuffer.length = 0;
+		this._sendBufferSize = 0;
 		super._invalidate();
 	}
 
@@ -219,6 +225,7 @@ export abstract class WebSocketBaseConnector<
 		for (let index = 0; index < this._sendBuffer.length; index++) {
 			if (this._sendBuffer[index].seq <= seq) {
 				dropCount += 1;
+				this._sendBufferSize -= this._sendBuffer[index].text.length;
 			} else {
 				break;
 			}
@@ -326,9 +333,31 @@ export abstract class WebSocketBaseConnector<
 			throw new libErrors.SessionLost("No session");
 		}
 
-		this._sendBuffer.push(message);
+		const text = JSON.stringify(message);
+		this._sendBuffer.push({ seq: message.seq, text });
+		this._sendBufferSize += text.length;
+		if (this.maxSendBufferSize && this._sendBufferSize > this.maxSendBufferSize) {
+			this._sendBufferOverflow();
+			return;
+		}
 		if (this._state === "connected") {
-			this._sendInternal(message);
+			this._socket!.send(text);
+		}
+	}
+
+	_sendBufferOverflow() {
+		logger.warn(
+			`Connector | closing connection to ${this.dst} after ${this._sendBuffer.length} messages ` +
+			"were left unacknowledged"
+		);
+		// Resuming is no longer possible, free the memory right away. The
+		// other side sees a normal drop and gets a new session when it reconnects.
+		this._sendBuffer.length = 0;
+		this._sendBufferSize = 0;
+		if (!this._closing) {
+			this.close(ConnectionClosed.TryAgainLater, "Send buffer limit exceeded").catch(err => {
+				logger.error(`Connector | error closing connection:\n${err.stack}`);
+			});
 		}
 	}
 
@@ -708,7 +737,7 @@ export abstract class WebSocketClientConnector extends WebSocketBaseConnector<We
 			this._heartbeatInterval = data.heartbeatInterval;
 			this.startHeartbeat();
 			for (let bufferedMessage of this._sendBuffer) {
-				this._sendInternal(bufferedMessage);
+				this._socket!.send(bufferedMessage.text);
 			}
 			this.emit("connect", data);
 
@@ -722,7 +751,7 @@ export abstract class WebSocketClientConnector extends WebSocketBaseConnector<We
 				`Connector | resuming existing session, resending ${this._sendBuffer.length} buffered messages`
 			);
 			for (let bufferedMessage of this._sendBuffer) {
-				this._sendInternal(bufferedMessage);
+				this._socket!.send(bufferedMessage.text);
 			}
 			this._startedResumingMs = null;
 			this.emit("resume");

test/integration/link.js

@@ -3,8 +3,9 @@ import events from "node:events";
 import jwt from "jsonwebtoken";
 
 import {
-	TestControlConnector, TestHostConnector, TestControl, url, controlToken,
+	TestControlConnector, TestHostConnector, TestControl, url, controlToken, getControl,
 } from "./index.js";
+import * as lib from "@clusterio/lib";
 import { ConnectionClosed, ProtocolError, PolicyViolation, AuthenticationFailed } from "@clusterio/lib";
 import { slowTest } from "../common.js";
 
@@ -92,6 +93,33 @@ describe("Integration of lib/link", function() {
 		await events.once(controlConnector, "connect");
 	});
 
+	it("should start a new session when the send buffer limit is exceeded", async function() {
+		await getControl().send(new lib.ControllerConfigSetFieldRequest("controller.max_send_buffer_size", "0.01"));
+		try {
+			await controlConnector.connect();
+		} finally {
+			await getControl().send(new lib.ControllerConfigSetFieldRequest("controller.max_send_buffer_size", "64"));
+		}
+
+		// Never acknowledge received messages
+		controlConnector._doHeartbeat = function() {
+			this._sendInternal(new lib.MessageHeartbeat());
+		};
+		let invalidated = false;
+		controlConnector.once("invalidate", () => { invalidated = true; });
+		const connected = events.once(controlConnector, "invalidate").then(
+			() => events.once(controlConnector, "connect")
+		);
+		for (let i = 0; i < 100; i++) {
+			if (invalidated) {
+				break;
+			}
+			await control.send(new lib.ControllerConfigGetRequest()).catch(() => {});
+		}
+		assert(invalidated, "session was not invalidated");
+		await connected;
+	});
+
 	it("should properly close connector if close is called during reconnect wait", async function() {
 		await controlConnector.connect();
 		controlConnector._socket.close(ConnectionClosed.Normal, "Test");

test/lib/link/connector.js

@@ -13,10 +13,11 @@ describe("lib/link/connectors", function() {
 		describe("._dropSendBufferSeq()", function() {
 			beforeEach(function() {
 				testConnector._sendBuffer = [
-					{ seq: 1, type: "test", data: {} },
-					{ seq: 2, type: "test", data: {} },
-					{ seq: 4, type: "test", data: {} },
+					{ seq: 1, text: "a" },
+					{ seq: 2, text: "bb" },
+					{ seq: 4, text: "cccc" },
 				];
+				testConnector._sendBufferSize = 7;
 			});
 			it("should do nothing if passed null", function() {
 				testConnector._dropSendBufferSeq(null);
@@ -25,6 +26,7 @@ describe("lib/link/connectors", function() {
 			it("should drop up to sequence", function() {
 				testConnector._dropSendBufferSeq(2);
 				assert.equal(testConnector._sendBuffer.length, 1, "incorrect number of entries dropped");
+				assert.equal(testConnector._sendBufferSize, 4);
 			});
 			it("should ignore holes", function() {
 				testConnector._dropSendBufferSeq(3);
@@ -61,6 +63,37 @@ describe("lib/link/connectors", function() {
 			});
 		});
 
+		describe("send buffer limit", function() {
+			afterEach(function() {
+				testConnector.maxSendBufferSize = 0;
+				testConnector._closing = false;
+			});
+			it("should keep unacknowledged messages under the limit", function() {
+				testConnector._state = "connected";
+				testConnector._sendBuffer = [];
+				testConnector._sendBufferSize = 0;
+				testConnector._socket.closeCalled = false;
+				testConnector.maxSendBufferSize = 1000;
+				testConnector.send(new lib.MessageHeartbeat(1));
+				assert.equal(testConnector._sendBuffer.length, 1);
+				assert(!testConnector._socket.closeCalled, "Close was called on the socket");
+			});
+			it("should clear the buffer and close when the limit is exceeded", function() {
+				testConnector._state = "connected";
+				testConnector._sendBuffer = [];
+				testConnector._sendBufferSize = 0;
+				testConnector._socket.closeCalled = false;
+				testConnector._socket.sentMessages = [];
+				testConnector.maxSendBufferSize = 30;
+				testConnector.send(new lib.MessageHeartbeat(1));
+				testConnector.send(new lib.MessageHeartbeat(2));
+				assert.equal(testConnector._sendBuffer.length, 0);
+				assert.equal(testConnector._sendBufferSize, 0);
+				assert(testConnector._socket.closeCalled, "Close was not called on the socket");
+				assert.equal(testConnector._socket.sentMessages.length, 1);
+			});
+		});
+
 		describe(".close()", function() {
 			it("should call close", function() {
 				testConnector._socket.closeCalled = false;