memory optimization

Signed-off-by: Hoang Pham <[email protected]>

package.json

@@ -30,6 +30,7 @@
     "express": "^4.19.2",
     "jsonwebtoken": "^9.0.2",
     "lodash": "^4.17.21",
+    "lru-cache": "^11.0.0",
     "node-fetch": "^3.3.2",
     "prom-client": "^14.2.0",
     "react": "^18.3.1",
@@ -77,4 +78,4 @@
     "node": "^20",
     "npm": "^9"
   }
-}
+}

websocket_server/app.js

@@ -6,6 +6,8 @@
 import dotenv from 'dotenv'
 import express from 'express'
 import { register } from 'prom-client'
+import { getSystemOverview } from './monitoring.js'
+import { rooms } from './roomData.js'
 
 dotenv.config()
 
@@ -27,4 +29,8 @@ app.get('/metrics', async (req, res) => {
 	res.end(metrics)
 })
 
+app.get('/system-overview', (req, res) => {
+	res.json(getSystemOverview(rooms))
+})
+
 export default app

websocket_server/monitoring.js

@@ -0,0 +1,60 @@
+/**
+ * SPDX-FileCopyrightText: 2024 Nextcloud GmbH and Nextcloud contributors
+ * SPDX-License-Identifier: AGPL-3.0-or-later
+ */
+
+export function getSystemOverview(rooms) {
+	return {
+		memoryUsage: getMemoryUsage(),
+		roomStats: getRoomStats(rooms),
+		cacheInfo: getCacheInfo(rooms),
+		roomsData: getRoomsData(rooms),
+	}
+}
+
+function getMemoryUsage() {
+	const memUsage = process.memoryUsage()
+	return {
+		rss: formatBytes(memUsage.rss),
+		heapTotal: formatBytes(memUsage.heapTotal),
+		heapUsed: formatBytes(memUsage.heapUsed),
+		external: formatBytes(memUsage.external),
+		arrayBuffers: formatBytes(memUsage.arrayBuffers),
+	}
+}
+
+function getRoomStats(rooms) {
+	return {
+		activeRooms: rooms.size,
+		totalUsers: Array.from(rooms.values()).reduce((sum, room) => sum + Object.keys(room.users).length, 0),
+		totalDataSize: formatBytes(Array.from(rooms.values()).reduce((sum, room) => sum + (room.data ? JSON.stringify(room.data).length : 0), 0)),
+	}
+}
+
+function getRoomsData(rooms) {
+	return Array.from(rooms.entries()).map(([roomId, room]) => ({
+		id: roomId,
+		users: Object.keys(room.users),
+		lastEditedUser: room.lastEditedUser,
+		lastActivity: new Date(room.lastActivity).toISOString(),
+		dataSize: formatBytes(JSON.stringify(room.data).length),
+		data: room.data, // Be cautious with this if the data is very large
+	}))
+}
+
+function getCacheInfo(rooms) {
+	return {
+		size: rooms.size,
+		maxSize: rooms.max,
+		keys: Array.from(rooms.keys()),
+		recentlyUsed: Array.from(rooms.keys()).slice(0, 10), // Show 10 most recently used
+	}
+}
+
+function formatBytes(bytes) {
+	if (bytes === 0) return '0 Bytes'
+	const k = 1024
+	const sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB']
+	const i = Math.floor(Math.log(bytes) / Math.log(k))
+	return parseFloat((bytes / Math.pow(k, i)).toFixed(2)) + ' ' + sizes[i]
+}

websocket_server/roomData.js

@@ -10,6 +10,7 @@ import dotenv from 'dotenv'
 import { generateSharedToken } from './auth.js'
 import https from 'https'
 import { parseBooleanFromEnv } from './utils.js'
+import { LRUCache } from 'lru-cache'
 
 dotenv.config()
 
@@ -20,9 +21,59 @@ const {
 
 const agent = parseBooleanFromEnv(IS_DEV) ? new https.Agent({ rejectUnauthorized: false }) : null
 
-export const roomDataStore = {}
-export const roomUsers = new Map()
-export const lastEditedUser = new Map()
+const INACTIVE_THRESHOLD = 30 * 60 * 1000 // 30 minutes
+const MAX_ROOMS = 100
+
+class Room {
+
+	constructor(id) {
+		this.id = id
+		this.data = null
+		this.users = {}
+		this.lastEditedUser = null
+		this.lastActivity = Date.now()
+	}
+
+	addUser(userId) {
+		this.users[userId] = true
+		this.updateActivity()
+	}
+
+	removeUser(userId) {
+		delete this.users[userId]
+		this.updateActivity()
+	}
+
+	updateLastEditedUser(userId) {
+		this.lastEditedUser = userId
+		this.updateActivity()
+	}
+
+	setData(data) {
+		this.data = data
+		this.updateActivity()
+	}
+
+	isEmpty() {
+		return Object.keys(this.users).length === 0
+	}
+
+	updateActivity() {
+		this.lastActivity = Date.now()
+	}
+
+}
+
+export const rooms = new LRUCache({
+	max: MAX_ROOMS,
+	ttl: INACTIVE_THRESHOLD,
+	updateAgeOnGet: true,
+	dispose: async (value, key) => {
+		if (value.data) {
+			await saveRoomDataToFile(key, value.data)
+		}
+	},
+})
 
 const fetchOptions = (method, token, body = null, roomId = null) => ({
 	method,
@@ -52,61 +103,86 @@ const fetchData = async (url, options) => {
 	}
 }
 
+export const getOrCreateRoom = (roomId) => {
+	let room = rooms.get(roomId)
+	if (!room) {
+		room = new Room(roomId)
+		rooms.set(roomId, room)
+	}
+	return room
+}
+
+export const addUserToRoom = (roomId, userId) => {
+	const room = getOrCreateRoom(roomId)
+	room.addUser(userId)
+}
+
+export const removeUserFromRoom = (roomId, userId) => {
+	const room = rooms.get(roomId)
+	if (room) {
+		room.removeUser(userId)
+		if (room.isEmpty()) {
+			rooms.delete(roomId)
+		}
+	}
+}
+
+export const updateLastEditedUser = (roomId, userId) => {
+	const room = rooms.get(roomId)
+	if (room) {
+		room.updateLastEditedUser(userId)
+	}
+}
+
+export const getRoomData = (roomId) => {
+	const room = rooms.get(roomId)
+	return room ? room.data : null
+}
+
+export const setRoomData = (roomId, data) => {
+	const room = getOrCreateRoom(roomId)
+	room.setData(data)
+}
+
 export const getRoomDataFromFile = async (roomID, jwtToken) => {
 	const url = `${NEXTCLOUD_URL}/index.php/apps/whiteboard/${roomID}`
 	const options = fetchOptions('GET', jwtToken)
 	const result = await fetchData(url, options)
-	console.log(result)
-	return result?.data.elements
+	const elements = result?.data?.elements
+	if (elements) {
+		setRoomData(roomID, elements)
+	}
+	return elements || null
 }
 
 export const saveRoomDataToFile = async (roomID) => {
 	console.log(`[${roomID}] Saving room data to file: ${roomID} with:`)
 	const url = `${NEXTCLOUD_URL}/index.php/apps/whiteboard/${roomID}`
-	const body = { data: { elements: roomDataStore[roomID] } }
+	const roomData = getRoomData(roomID)
+	const body = { data: { elements: roomData } }
 	const options = fetchOptions('PUT', null, body, roomID)
 	console.log(options)
 	await fetchData(url, options)
 }
 
 export const handleEmptyRoom = async (roomID) => {
-	if (roomDataStore[roomID]) {
+	const roomData = getRoomData(roomID)
+	if (roomData) {
 		await saveRoomDataToFile(roomID)
 		console.log('Removing data for room', roomID)
-		delete roomDataStore[roomID]
-		lastEditedUser.delete(roomID)
+		rooms.delete(roomID)
 	}
 }
 
 export const saveAllRoomsData = () =>
-	Promise.all(Object.entries(roomDataStore).map(([roomId, roomData]) =>
-		saveRoomDataToFile(roomId, roomData)))
+	Promise.all(Array.from(rooms.entries()).map(([roomId, room]) =>
+		saveRoomDataToFile(roomId, room.data)))
 
 export const removeAllRoomData = () => {
-	Object.keys(roomDataStore).forEach(key => delete roomDataStore[key])
-	roomUsers.clear()
+	rooms.clear()
 }
 
-export const addUserToRoom = (roomId, userId) => {
-	if (!roomUsers.has(roomId)) {
-		roomUsers.set(roomId, new Set())
-	}
-	roomUsers.get(roomId).add(userId)
+export const getLastEditedUser = (roomId) => {
+	const room = rooms.get(roomId)
+	return room ? room.lastEditedUser : null
 }
-
-export const removeUserFromRoom = (roomId, userId) => {
-	const room = roomUsers.get(roomId)
-	if (room) {
-		room.delete(userId)
-		if (room.size === 0) {
-			roomUsers.delete(roomId)
-		}
-	}
-}
-
-export const updateLastEditedUser = (roomId, userId) => {
-	lastEditedUser.set(roomId, userId)
-}
-
-export const getLastEditedUser = (roomId) =>
-	lastEditedUser.get(roomId) || (roomUsers.has(roomId) ? Array.from(roomUsers.get(roomId)).pop() : null)