From 2be91cbd3a24f58c5d9ce6b57de3b22004b68a3e Mon Sep 17 00:00:00 2001
From: Hornwitser <github@hornwitser.no>
Date: Wed, 2 Jun 2021 16:19:18 +0200
Subject: [PATCH] Add slave update subscriptions

Allow the web interface to listen for updates to the slave list and
apply it live to the places where slaves are shown in the web interface.
---
 packages/ctl/ctl.js                           |  2 +
 packages/lib/link/messages.js                 | 36 ++++++++---
 packages/lib/users.js                         |  6 ++
 packages/master/src/ControlConnection.js      | 18 ++++++
 packages/master/src/Master.js                 | 18 ++++++
 packages/master/src/WsServer.js               |  2 +
 .../src/components/AssignInstanceModal.jsx    | 12 ++--
 packages/web_ui/src/components/SlavesPage.jsx | 40 ++++++++----
 packages/web_ui/src/model/slave.jsx           | 45 ++++++++++++++
 packages/web_ui/src/util/websocket.jsx        | 61 +++++++++++++++++++
 test/integration/clusterio.js                 | 55 ++++++++++++++++-
 test/integration/index.js                     |  9 +++
 12 files changed, 275 insertions(+), 29 deletions(-)

diff --git a/packages/ctl/ctl.js b/packages/ctl/ctl.js
index f691fe08..04c50c73 100755
--- a/packages/ctl/ctl.js
+++ b/packages/ctl/ctl.js
@@ -918,6 +918,8 @@ class Control extends libLink.Link {
 		this.keepOpen = false;
 	}
 
+	async slaveUpdateEventHandler() { }
+
 	async instanceUpdateEventHandler() { }
 
 	async saveListUpdateEventHandler() { }
diff --git a/packages/lib/link/messages.js b/packages/lib/link/messages.js
index 4e3c17e1..72dbe3a8 100644
--- a/packages/lib/link/messages.js
+++ b/packages/lib/link/messages.js
@@ -331,6 +331,14 @@ messages.setMasterConfigProp = new Request({
 	},
 });
 
+let slaveProperties = {
+	"agent": { type: "string" },
+	"version": { type: "string" },
+	"name": { type: "string" },
+	"id": { type: "integer" },
+	"connected": { type: "boolean" },
+};
+
 messages.listSlaves = new Request({
 	type: "list_slaves",
 	links: ["control-master"],
@@ -341,18 +349,25 @@ messages.listSlaves = new Request({
 			items: {
 				additionalProperties: false,
 				required: ["agent", "version", "name", "id", "connected"],
-				properties: {
-					"agent": { type: "string" },
-					"version": { type: "string" },
-					"name": { type: "string" },
-					"id": { type: "integer" },
-					"connected": { type: "boolean" },
-				},
+				properties: slaveProperties,
 			},
 		},
 	},
 });
 
+messages.setSlaveSubscriptions = new Request({
+	type: "set_slave_subscriptions",
+	links: ["control-master"],
+	permission: "core.slave.subscribe",
+	requestProperties: {
+		"all": { type: "boolean" },
+		"slave_ids": {
+			type: "array",
+			items: { type: "integer" },
+		},
+	},
+});
+
 let instanceProperties = {
 	"name": { type: "string" },
 	"id": { type: "integer" },
@@ -1083,6 +1098,13 @@ messages.logMessage = new Event({
 	},
 });
 
+messages.slaveUpdate = new Event({
+	type: "slave_update",
+	links: ["master-control"],
+	eventProperties: slaveProperties,
+});
+
+
 messages.instanceInitialized = new Event({
 	type: "instance_initialized",
 	links: ["instance-slave"],
diff --git a/packages/lib/users.js b/packages/lib/users.js
index c6c0060a..3f91ab11 100644
--- a/packages/lib/users.js
+++ b/packages/lib/users.js
@@ -322,6 +322,12 @@ definePermission({
 	description: "Get the full list of slaves in the cluster.",
 	grantByDefault: true,
 });
+definePermission({
+	name: "core.slave.subscribe",
+	title: "Subscribe to slave updates",
+	description: "Subscribe to be notified on updates on the details of slaves.",
+	grantByDefault: true,
+});
 definePermission({
 	name: "core.slave.generate_token",
 	title: "Generate slave token",
diff --git a/packages/master/src/ControlConnection.js b/packages/master/src/ControlConnection.js
index 4f1c4fbd..ad53cb0b 100644
--- a/packages/master/src/ControlConnection.js
+++ b/packages/master/src/ControlConnection.js
@@ -39,6 +39,11 @@ class ControlConnection extends BaseConnection {
 		 */
 		this.user = user;
 
+		this.slaveSubscriptions = {
+			all: false,
+			slave_ids: [],
+		};
+
 		this.instanceSubscriptions = {
 			all: false,
 			instance_ids: [],
@@ -107,6 +112,19 @@ class ControlConnection extends BaseConnection {
 		return { list };
 	}
 
+	async setSlaveSubscriptionsRequestHandler(message) {
+		this.slaveSubscriptions = message.data;
+	}
+
+	slaveUpdated(slave, update) {
+		if (
+			this.slaveSubscriptions.all
+			|| this.slaveSubscriptions.slave_ids.includes(slave.id)
+		) {
+			libLink.messages.slaveUpdate.send(this, update);
+		}
+	}
+
 	generateSlaveToken(slaveId) {
 		return jwt.sign({ aud: "slave", slave: slaveId }, this._master.config.get("master.auth_secret"));
 	}
diff --git a/packages/master/src/Master.js b/packages/master/src/Master.js
index 8f97c975..f7590049 100644
--- a/packages/master/src/Master.js
+++ b/packages/master/src/Master.js
@@ -360,6 +360,24 @@ class Master {
 		}
 	}
 
+	slaveUpdated(slave) {
+		let update = {
+			agent: slave.agent,
+			version: slave.version,
+			id: slave.id,
+			name: slave.name,
+			connected: this.wsServer.slaveConnections.has(slave.id),
+		};
+
+		for (let controlConnection of this.wsServer.controlConnections) {
+			if (controlConnection.connector.closing) {
+				continue;
+			}
+
+			controlConnection.slaveUpdated(slave, update);
+		}
+	}
+
 	addInstanceHooks(instance) {
 		instance.config.on("fieldChanged", (group, field, prev) => {
 			if (group.name === "instance" && field === "name") {
diff --git a/packages/master/src/WsServer.js b/packages/master/src/WsServer.js
index 3c77824f..47342420 100644
--- a/packages/master/src/WsServer.js
+++ b/packages/master/src/WsServer.js
@@ -274,9 +274,11 @@ ${err.stack}`
 			connector.on("close", () => {
 				if (this.slaveConnections.get(data.id) === connection) {
 					this.slaveConnections.delete(data.id);
+					this.master.slaveUpdated(this.master.slaves.get(data.id));
 				}
 			});
 			this.slaveConnections.set(data.id, connection);
+			this.master.slaveUpdated(this.master.slaves.get(data.id));
 
 
 		} else if (type === "register_control") {
diff --git a/packages/web_ui/src/components/AssignInstanceModal.jsx b/packages/web_ui/src/components/AssignInstanceModal.jsx
index 8a0b2399..ad06f474 100644
--- a/packages/web_ui/src/components/AssignInstanceModal.jsx
+++ b/packages/web_ui/src/components/AssignInstanceModal.jsx
@@ -5,24 +5,20 @@ import libLink from "@clusterio/lib/link";
 
 import ControlContext from "./ControlContext";
 import { notifyErrorHandler } from "../util/notify";
+import { useSlaveList } from "../model/slave";
 
 const { Paragraph } = Typography;
 
 
 export default function AssignInstanceModal(props) {
 	let [visible, setVisible] = useState(false);
-	let [slaves, setSlaves] = useState([]);
-	let [loading, setLoading] = useState(true);
+	let [slaveList] = useSlaveList();
 	let [applying, setApplying] = useState(false);
 	let [form] = Form.useForm();
 	let control = useContext(ControlContext);
 
 	function open() {
 		setVisible(true);
-		libLink.messages.listSlaves.send(control).then((result) => {
-			setSlaves(result.list);
-			setLoading(false);
-		});
 	}
 
 	function handleAssign() {
@@ -73,8 +69,8 @@ export default function AssignInstanceModal(props) {
 			</Paragraph>
 			<Form form={form} initialValues={{ slave: props.slaveId }}>
 				<Form.Item name="slave" label="Slave">
-					<Select loading={loading}>
-						{slaves.map((slave) => <Select.Option
+					<Select>
+						{slaveList.map((slave) => <Select.Option
 							key={slave["id"]}
 							value={slave["id"]}
 							disabled={!slave["connected"]}
diff --git a/packages/web_ui/src/components/SlavesPage.jsx b/packages/web_ui/src/components/SlavesPage.jsx
index c3bd02ce..0283c8bb 100644
--- a/packages/web_ui/src/components/SlavesPage.jsx
+++ b/packages/web_ui/src/components/SlavesPage.jsx
@@ -1,28 +1,42 @@
 import React, { useContext } from "react";
+import { Table } from "antd";
 
 import libLink from "@clusterio/lib/link";
 
-import DataTable from "./data-table";
 import ControlContext from "./ControlContext";
 import PageLayout from "./PageLayout";
+import { useSlaveList } from "../model/slave";
 
 
 export default function SlavesPage() {
 	let control = useContext(ControlContext);
-
-	async function listSlaves() {
-		let result = await libLink.messages.listSlaves.send(control);
-		return result["list"].map((item) => ({
-			key: item["id"],
-			"Name": item["name"],
-			"Agent": item["agent"],
-			"Version": item["version"],
-			"Connected": item["connected"] && "Yes",
-		}));
-	}
+	let [slaveList] = useSlaveList();
 
 	return <PageLayout nav={[{ name: "Slaves" }]}>
 		<h2>Slaves</h2>
-		<DataTable DataFunction={listSlaves} />
+		<Table
+			columns={[
+				{
+					title: "Name",
+					dataIndex: "name",
+				},
+				{
+					title: "Agent",
+					dataIndex: "agent",
+				},
+				{
+					title: "Version",
+					dataIndex: "version",
+				},
+				{
+					title: "Connected",
+					key: "connected",
+					render: slave => slave["connected"] && "Yes",
+				},
+			]}
+			dataSource={slaveList}
+			rowKey={slave => slave["id"]}
+			pagination={false}
+		/>
 	</PageLayout>;
 };
diff --git a/packages/web_ui/src/model/slave.jsx b/packages/web_ui/src/model/slave.jsx
index ca6ee136..c29e7405 100644
--- a/packages/web_ui/src/model/slave.jsx
+++ b/packages/web_ui/src/model/slave.jsx
@@ -27,7 +27,52 @@ export function useSlave(id) {
 
 	useEffect(() => {
 		updateSlave();
+
+		function updateHandler(newSlave) {
+			setSlave({ ...newSlave, present: true });
+		}
+
+		control.onSlaveUpdate(id, updateHandler);
+		return () => {
+			control.offSlaveUpdate(id, updateHandler);
+		};
 	}, [id]);
 
 	return [slave, updateSlave];
 }
+
+export function useSlaveList() {
+	let control = useContext(ControlContext);
+	let [slaveList, setSlaveList] = useState([]);
+
+	function updateSlaveList() {
+		libLink.messages.listSlaves.send(control).then(result => {
+			setSlaveList(result.list);
+		});
+	}
+
+	useEffect(() => {
+		updateSlaveList();
+
+		function updateHandler(newSlave) {
+			setSlaveList(oldList => {
+				let newList = oldList.concat();
+				let index = newList.findIndex(s => s.id === newSlave.id);
+				if (index !== -1) {
+					newList[index] = newSlave;
+				} else {
+					newList.push(newSlave);
+				}
+				return newList;
+			});
+		}
+
+		control.onSlaveUpdate(null, updateHandler);
+		return () => {
+			control.offSlaveUpdate(null, updateHandler);
+		};
+	}, []);
+
+
+	return [slaveList];
+}
diff --git a/packages/web_ui/src/util/websocket.jsx b/packages/web_ui/src/util/websocket.jsx
index 63d7d689..6b07fc65 100644
--- a/packages/web_ui/src/util/websocket.jsx
+++ b/packages/web_ui/src/util/websocket.jsx
@@ -62,12 +62,16 @@ export class Control extends libLink.Link {
 		 */
 		this.accountName = null;
 
+		this.slaveUpdateHandlers = new Map();
 		this.instanceUpdateHandlers = new Map();
 		this.saveListUpdateHandlers = new Map();
 		this.instanceLogHandlers = new Map();
 
 		this.connector.on("connect", data => {
 			this.accountName = data.account.name;
+			this.updateSlaveSubscriptions().catch(err => logger.error(
+				`Unexpected error updating slave subscriptions:\n${err.stack}`
+			));
 			this.updateInstanceSubscriptions().catch(err => logger.error(
 				`Unexpected error updating instance subscriptions:\n${err.stack}`
 			));
@@ -92,6 +96,63 @@ export class Control extends libLink.Link {
 		}
 	}
 
+	async slaveUpdateEventHandler(message) {
+		let handlers = [].concat(
+			this.slaveUpdateHandlers.get(null) || [],
+			this.slaveUpdateHandlers.get(message.data.id) || [],
+		);
+		for (let handler of handlers) {
+			handler(message.data);
+		}
+	}
+
+	async onSlaveUpdate(id, handler) {
+		if (id !== null && !Number.isInteger(id)) {
+			throw new Error("Invalid slave id");
+		}
+
+		let handlers = this.slaveUpdateHandlers.get(id);
+		if (!handlers) {
+			handlers = [];
+			this.slaveUpdateHandlers.set(id, handlers);
+		}
+
+		handlers.push(handler);
+
+		if (handlers.length === 1) {
+			await this.updateSlaveSubscriptions();
+		}
+	}
+
+	async offSlaveUpdate(id, handler) {
+		let handlers = this.slaveUpdateHandlers.get(id);
+		if (!handlers || !handlers.length) {
+			throw new Error(`No handlers for slave ${id} exists`);
+		}
+
+		let index = handlers.lastIndexOf(handler);
+		if (index === -1) {
+			throw new Error(`Given handler is not registered for slave ${id}`);
+		}
+
+		handlers.slice(index, 1);
+		if (!handlers.length) {
+			this.slaveUpdateHandlers.delete(id);
+			await this.updateSlaveSubscriptions();
+		}
+	}
+
+	async updateSlaveSubscriptions() {
+		if (!this.connector.connected) {
+			return;
+		}
+
+		await libLink.messages.setSlaveSubscriptions.send(this, {
+			all: this.slaveUpdateHandlers.has(null),
+			slave_ids: [...this.slaveUpdateHandlers.keys()].filter(e => e !== null),
+		});
+	}
+
 	async instanceUpdateEventHandler(message) {
 		let handlers = this.instanceUpdateHandlers.get(message.data.id);
 		for (let handler of handlers || []) {
diff --git a/test/integration/clusterio.js b/test/integration/clusterio.js
index c73c3d32..82f7da77 100644
--- a/test/integration/clusterio.js
+++ b/test/integration/clusterio.js
@@ -10,7 +10,7 @@ const libUsers = require("@clusterio/lib/users");
 const testStrings = require("../lib/factorio/test_strings");
 const {
 	TestControl, TestControlConnector, url, controlToken, slowTest,
-	get, exec, execCtl, sendRcon, getControl, instancesDir,
+	get, exec, execCtl, sendRcon, getControl, spawn, instancesDir,
 } = require("./index");
 
 
@@ -62,6 +62,59 @@ describe("Integration of Clusterio", function() {
 		});
 	});
 
+	describe("clusterioslave", function() {
+		describe("slaveUpdateEventHandler()", function() {
+			it("should trigger when a new slave is added", async function() {
+				slowTest(this);
+				getControl().slaveUpdates = [];
+				let config = "alt-slave-config.json";
+				let configPath = path.join("temp", "test", config);
+				await fs.remove(configPath);
+				await execCtl(`slave create-config --id 5 --name alt-slave --generate-token --output ${config}`);
+				await exec(
+					`node ../../packages/slave --config ${config} config set slave.tls_ca ../../test/file/tls/cert.pem`
+				);
+
+				let slaveProcess;
+				try {
+					slaveProcess = await spawn(
+						"alt-slave:", `node ../../packages/slave run --config ${config}`, /SOCKET \| registering slave/
+					);
+				} finally {
+					if (slaveProcess) {
+						slaveProcess.kill("SIGINT");
+						await events.once(slaveProcess, "exit");
+					}
+				}
+
+				let sawUpdate = false;
+				let sawConnected = false;
+				let sawDisconnected = false;
+
+				for (let update of getControl().slaveUpdates) {
+					if (update.name !== "alt-slave") {
+						continue;
+					}
+
+					sawUpdate = true;
+					if (update.connected) {
+						sawConnected = true;
+					} else {
+						sawDisconnected = true;
+					}
+				}
+
+				assert(sawUpdate, "No slave update was sent");
+				assert(sawConnected, "No slave update with status connected was sent");
+				assert(sawDisconnected, "No slave update with status disconnected was sent");
+
+				let result = await libLink.messages.listSlaves.send(getControl());
+				let slaves = new Map(result.list.map(instance => [instance.id, instance]));
+				assert(slaves.has(5), "Slave list was not updated");
+			});
+		});
+	});
+
 	describe("clusterioctl", function() {
 		describe("master config list", function() {
 			it("runs", async function() {
diff --git a/test/integration/index.js b/test/integration/index.js
index ad2fa030..68094800 100644
--- a/test/integration/index.js
+++ b/test/integration/index.js
@@ -25,10 +25,14 @@ class TestControl extends libLink.Link {
 	constructor(connector) {
 		super("control", "master", connector);
 		libLink.attachAllMessages(this);
+		this.slaveUpdates = [];
 		this.instanceUpdates = [];
 		this.saveListUpdates = [];
 
 		this.connector.on("connect", () => {
+			libLink.messages.setSlaveSubscriptions.send(
+				this, { all: true, slave_ids: [] }
+			).catch(err => logger.error(`Error setting slave subscriptions:\n${err.stack}`));
 			libLink.messages.setInstanceSubscriptions.send(
 				this, { all: true, instance_ids: [] }
 			).catch(err => logger.error(`Error setting instance subscriptions:\n${err.stack}`));
@@ -45,6 +49,10 @@ class TestControl extends libLink.Link {
 
 	async debugWsMessageEventHandler() { }
 
+	async slaveUpdateEventHandler(message) {
+		this.slaveUpdates.push(message.data);
+	}
+
 	async instanceUpdateEventHandler(message) {
 		this.instanceUpdates.push(message.data);
 	}
@@ -223,6 +231,7 @@ module.exports = {
 	execCtl,
 	sendRcon,
 	getControl,
+	spawn,
 
 	url,
 	controlToken,