memory optimization

Signed-off-by: Hoang Pham <hoangmaths96@gmail.com>

websocket_server/app.js

@@ -5,8 +5,7 @@
 
 import dotenv from 'dotenv'
 import express from 'express'
-import { register } from 'prom-client'
-import { getSystemOverview } from './monitoring.js'
+import { register, updatePrometheusMetrics } from './prom-metrics.js'
 import { rooms } from './roomData.js'
 
 dotenv.config()
@@ -24,13 +23,10 @@ app.get('/metrics', async (req, res) => {
 	if (!METRICS_TOKEN || token !== METRICS_TOKEN) {
 		return res.status(403).send('Unauthorized')
 	}
+	updatePrometheusMetrics(rooms)
 	const metrics = await register.metrics()
 	res.set('Content-Type', register.contentType)
 	res.end(metrics)
 })
 
-app.get('/system-overview', (req, res) => {
-	res.json(getSystemOverview(rooms))
-})
-
 export default app

websocket_server/prom-metrics.js

@@ -0,0 +1,42 @@
+/**
+ * SPDX-FileCopyrightText: 2024 Nextcloud GmbH and Nextcloud contributors
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+import { register, Gauge } from 'prom-client'
+import { getSystemOverview } from './monitoring.js'
+
+const memoryUsageGauge = new Gauge({
+	name: 'whiteboard_memory_usage',
+	help: 'Memory usage of the server',
+	labelNames: ['type'],
+})
+
+const roomStatsGauge = new Gauge({
+	name: 'whiteboard_room_stats',
+	help: 'Room statistics',
+	labelNames: ['stat'],
+})
+
+const cacheInfoGauge = new Gauge({
+	name: 'whiteboard_cache_info',
+	help: 'Cache information',
+	labelNames: ['info'],
+})
+
+export function updatePrometheusMetrics(rooms) {
+	const overview = getSystemOverview(rooms)
+
+	Object.entries(overview.memoryUsage).forEach(([key, value]) => {
+		memoryUsageGauge.set({ type: key }, parseFloat(value) || 0)
+	})
+
+	roomStatsGauge.set({ stat: 'activeRooms' }, Number(overview.roomStats.activeRooms) || 0)
+	roomStatsGauge.set({ stat: 'totalUsers' }, Number(overview.roomStats.totalUsers) || 0)
+	roomStatsGauge.set({ stat: 'totalDataSize' }, parseFloat(overview.roomStats.totalDataSize) || 0)
+
+	cacheInfoGauge.set({ info: 'size' }, Number(overview.cacheInfo.size) || 0)
+	cacheInfoGauge.set({ info: 'maxSize' }, Number(overview.cacheInfo.maxSize) || 0)
+}
+
+export { register }

websocket_server/roomData.js

@@ -69,20 +69,26 @@ export const rooms = new LRUCache({
 	ttl: INACTIVE_THRESHOLD,
 	updateAgeOnGet: true,
 	dispose: async (value, key) => {
-		if (value.data) {
-			await saveRoomDataToFile(key, value.data)
+		console.log('Disposing room', key)
+
+		if (value && value.data && value.lastEditedUser) {
+			try {
+				await saveRoomDataToFile(key, value.data, value.lastEditedUser)
+			} catch (error) {
+				console.error(`Failed to save room ${key} data:`, error)
+			}
 		}
 	},
 })
 
-const fetchOptions = (method, token, body = null, roomId = null) => ({
+const fetchOptions = (method, token, body = null, roomId = null, lastEditedUser = null) => ({
 	method,
 	headers: {
 		'Content-Type': 'application/json',
 		...(method === 'GET' && { Authorization: `Bearer ${token}` }),
 		...(method === 'PUT' && {
 			'X-Whiteboard-Auth': generateSharedToken(roomId),
-			'X-Whiteboard-User': getLastEditedUser(roomId),
+			'X-Whiteboard-User': lastEditedUser || 'unknown',
 		}),
 	},
 	...(body && { body: JSON.stringify(body) }),
@@ -135,6 +141,7 @@ export const updateLastEditedUser = (roomId, userId) => {
 }
 
 export const getRoomData = (roomId) => {
+	console.log('Getting data from memory for room', roomId)
 	const room = rooms.get(roomId)
 	return room ? room.data : null
 }
@@ -155,22 +162,18 @@ export const getRoomDataFromFile = async (roomID, jwtToken) => {
 	return elements || null
 }
 
-export const saveRoomDataToFile = async (roomID) => {
-	console.log(`[${roomID}] Saving room data to file: ${roomID} with:`)
+export const saveRoomDataToFile = async (roomID, roomData, lastEditedUser) => {
+	console.log('Saving room data to file', roomID)
 	const url = `${NEXTCLOUD_URL}/index.php/apps/whiteboard/${roomID}`
-	const roomData = getRoomData(roomID)
 	const body = { data: { elements: roomData } }
-	const options = fetchOptions('PUT', null, body, roomID)
-	console.log(options)
+	const options = fetchOptions('PUT', null, body, roomID, lastEditedUser)
 	await fetchData(url, options)
 }
 
 export const handleEmptyRoom = async (roomID) => {
 	const roomData = getRoomData(roomID)
 	if (roomData) {
-		await saveRoomDataToFile(roomID)
-		console.log('Removing data for room', roomID)
-		rooms.delete(roomID)
+		rooms.delete(roomID) // This will trigger the dispose function so that the data is saved to file
 	}
 }
 
@@ -181,8 +184,3 @@ export const saveAllRoomsData = () =>
 export const removeAllRoomData = () => {
 	rooms.clear()
 }
-
-export const getLastEditedUser = (roomId) => {
-	const room = rooms.get(roomId)
-	return room ? room.lastEditedUser : null
-}

websocket_server/socket.js

@@ -19,6 +19,7 @@ import {
 } from './roomData.js'
 import { convertArrayBufferToString, convertStringToArrayBuffer } from './utils.js'
 import dotenv from 'dotenv'
+import { LRUCache } from 'lru-cache'
 
 dotenv.config()
 
@@ -27,7 +28,16 @@ const {
 	JWT_SECRET_KEY,
 } = process.env
 
-const tokenCache = new Map()
+const TOKEN_CACHE_TTL = 10 * 60 * 1000 // 10 minutes, << JWT expiration time
+const tokenCache = new LRUCache({
+	ttl: TOKEN_CACHE_TTL,
+	updateAgeOnGet: false,
+	max: 1000,
+})
+
+export const removeTokenFromCache = (token) => {
+	tokenCache.delete(token)
+}
 
 const verifyToken = async (token) => {
 	if (tokenCache.has(token)) {