From 75067e992d798173ab8db6c5a2a4b73258682000 Mon Sep 17 00:00:00 2001
From: Hornwitser <github@hornwitser.no>
Date: Sat, 4 Apr 2020 12:03:50 +0200
Subject: [PATCH] Add WebSocket usage statistics

Add Prometheus metrics for messages sent, connections made, connections
rejected, active connections, and active slaves.
---
 CHANGELOG.md |  1 +
 master.js    | 53 ++++++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index efe173b3..db27eb9a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,7 @@ Version 2.0.0
   version to use on a per instance basis.
 - Added config to specify which path the master interface is accessed under,
   allowing it to be proxied behind web-server.
+- Added WebSocket usage statistics.
 
 ### Changes
 
diff --git a/master.js b/master.js
index 977f391b..876d9bac 100644
--- a/master.js
+++ b/master.js
@@ -86,11 +86,38 @@ require("./routes/api/getPictures")(app);
 // Set folder to serve static content from (the website)
 app.use(express.static('static'));
 
-const endpointHitCounter = new prometheus.Gauge(
-	'clusterio_master_http_endpoint_hits_total', "How many requests a particular endpoint has gotten",
+const endpointHitCounter = new prometheus.Counter(
+	"clusterio_master_http_endpoint_hits_total",
+	"How many requests a particular HTTP endpoint has gotten",
 	{ labels: ['route'] }
 );
 
+const wsMessageCounter = new prometheus.Counter(
+	"clusterio_master_websocket_message_total",
+	"How many messages have been received over WebSocket on the master server",
+	{ labels: ["direction"] }
+);
+
+const wsConnectionsCounter = new prometheus.Counter(
+	"clusterio_master_websocket_connections_total",
+	"How many WebSocket connections have been initiated on the master server"
+);
+
+const wsRejectedConnectionsCounter = new prometheus.Counter(
+	"clusterio_master_websocket_rejected_connections_total",
+	"How many WebSocket connections have been rejected during the handshake on the master server"
+);
+
+const wsActiveConnectionsGauge = new prometheus.Gauge(
+	"clusterio_master_websocket_active_connections",
+	"How many WebSocket connections are currently open to the master server"
+);
+
+const wsActiveSlavesGauge = new prometheus.Gauge(
+	"clusterio_master_active_slaves",
+	"How many slaves are currently connected to the master"
+);
+
 // Prometheus polling endpoint
 async function getMetrics(req, res, next) {
 	endpointHitCounter.labels(req.route.path).inc();
@@ -133,6 +160,10 @@ async function getMetrics(req, res, next) {
 		results.push(prometheus.deserializeResult(result));
 	}
 
+	wsActiveConnectionsGauge.set(activeConnectors.size);
+	wsActiveSlavesGauge.set(slaveConnections.size);
+
+
 	let text = await prometheus.exposition(results);
 	res.set('Content-Type', prometheus.exposition.contentType);
 	res.send(text);
@@ -513,10 +544,6 @@ class SlaveConnection extends BaseConnection {
 			version: this._version,
 		});
 
-		this.connector.on('message', () => {
-			// XXX prometheusWsUsageCounter.labels('message', "other").inc();
-		});
-
 		this.connector.on("close", () => {
 			if (slaveConnections.get(this._id) === this) {
 				slaveConnections.delete(this._id);
@@ -744,6 +771,15 @@ let activeConnectors = new Map();
 wss.on("connection", function (socket, req) {
 	console.log(`SOCKET | new connection from ${req.socket.remoteAddress}`);
 
+	// Track statistics
+	wsConnectionsCounter.inc();
+	socket.on("message", () => { wsMessageCounter.labels("in").inc(); });
+	let originalSend = socket.send;
+	socket.send = (...args) => {
+		wsMessageCounter.labels("out").inc();
+		return originalSend.call(socket, ...args);
+	};
+
 	// Start connection handshake.
 	socket.send(JSON.stringify({ seq: null, type: "hello", data: { version }}));
 
@@ -752,6 +788,7 @@ wss.on("connection", function (socket, req) {
 
 		let timeoutId = setTimeout(() => {
 			console.log(`SOCKET | closing ${req.socket.remoteAddress} after timing out on handshake`);
+			wsRejectedConnectionsCounter.inc();
 			socket.close(1008, "Handshake timeout");
 			pendingSockets.delete(socket);
 		}, 30*1000);
@@ -769,6 +806,7 @@ wss.on("connection", function (socket, req) {
 +----------------------------------------------------------------+`
 				);
 				console.error(err);
+				wsRejectedConnectionsCounter.inc();
 				socket.close(1011, "Unexpected error");
 			});
 		});
@@ -788,12 +826,14 @@ async function handleHandshake(message, socket, req, attachHandler) {
 		message = JSON.parse(message);
 	} catch (err) {
 		console.log(`SOCKET | closing ${req.socket.remoteAddress} after receiving invalid JSON`);
+		wsRejectedConnectionsCounter.inc();
 		socket.close(1001, "Invalid JSON");
 		return;
 	}
 
 	if (!schema.clientHandshake(message)) {
 		console.log(`SOCKET | closing ${req.socket.remoteAddress} after receiving invalid handshake:`, message);
+		wsRejectedConnectionsCounter.inc();
 		socket.close(1001, "Bad handshake");
 		return;
 	}
@@ -828,6 +868,7 @@ async function handleHandshake(message, socket, req, attachHandler) {
 
 	if (!result.ok) {
 		console.error(`SOCKET | authentication failed for ${req.socket.remoteAddress}`);
+		wsRejectedConnectionsCounter.inc();
 		socket.close(4003, "Authentication failed");
 		return;
 	}