Support persisting battles in a Postgres database (#8442)

databases/schemas/stored-battles.sql

@@ -0,0 +1,6 @@
+CREATE TABLE stored_battles (
+	roomid TEXT NOT NULL PRIMARY KEY, -- can store both the num and the formatid
+	input_log TEXT NOT NULL,
+	players TEXT[] NOT NULL,
+	title TEXT NOT NULL
+);

lib/index.ts

@@ -7,3 +7,4 @@ export * as Utils from './utils';
 export {crashlogger} from './crashlogger';
 export * as ProcessManager from './process-manager';
 export {SQL} from './sql';
+export {PostgresDatabase} from './postgres';

lib/postgres.ts

@@ -0,0 +1,80 @@
+/**
+ * Library made to simplify accessing / connecting to postgres databases,
+ * and to cleanly handle when the pg module isn't installed.
+ * @author mia-pi-git
+ */
+
+// @ts-ignore in case module doesn't exist
+import type * as PG from 'pg';
+import type {SQLStatement} from 'sql-template-strings';
+import * as Streams from './streams';
+
+export class PostgresDatabase {
+	private pool: PG.Pool;
+	constructor(config = PostgresDatabase.getConfig()) {
+		try {
+			this.pool = new (require('pg').Pool)(config);
+		} catch (e: any) {
+			this.pool = null!;
+		}
+	}
+	async query(statement: string | SQLStatement, values?: any[]) {
+		if (!this.pool) {
+			throw new Error(`Attempting to use postgres without 'pg' installed`);
+		}
+		let result;
+		try {
+			result = await this.pool.query(statement, values);
+		} catch (e: any) {
+			// postgres won't give accurate stacks unless we do this
+			throw new Error(e.message);
+		}
+		return result?.rows || [];
+	}
+	static getConfig() {
+		let config: AnyObject = {};
+		try {
+			config = require('../config/config').usepostgres;
+			if (!config) throw new Error('Missing config for pg database');
+		} catch (e: any) {}
+		return config;
+	}
+	async transaction(callback: (conn: PG.PoolClient) => any, depth = 0): Promise<any> {
+		const conn = await this.pool.connect();
+		await conn.query(`BEGIN`);
+		let result;
+		try {
+			// eslint-disable-next-line callback-return
+			result = await callback(conn);
+		} catch (e: any) {
+			await conn.query(`ROLLBACK`);
+			// two concurrent transactions conflicted, try again
+			if (e.code === '40001' && depth <= 10) {
+				return this.transaction(callback, depth + 1);
+				// There is a bug in Postgres that causes some
+				// serialization failures to be reported as failed
+				// unique constraint checks. Only retrying once since
+				// it could be our fault (thanks chaos for this info / the first half of this comment)
+			} else if (e.code === '23505' && !depth) {
+				return this.transaction(callback, depth + 1);
+			} else {
+				throw e;
+			}
+		}
+		await conn.query(`COMMIT`);
+		return result;
+	}
+	stream<T = any>(query: string) {
+		// eslint-disable-next-line @typescript-eslint/no-this-alias
+		const db = this;
+		return new Streams.ObjectReadStream<T>({
+			async read(this: Streams.ObjectReadStream<T>) {
+				const result = await db.query(query) as T[];
+				if (!result.length) return this.pushEnd();
+				// getting one row at a time means some slower queries
+				// might help with performance
+				this.buf.push(...result);
+			},
+		});
+	}
+}

package.json

@@ -4,6 +4,7 @@
   "version": "0.11.7",
   "main": ".sim-dist/index.js",
   "dependencies": {
+    "@types/pg": "^8.6.1",
     "@swc/core": "1.2.205",
     "preact": "^10.5.15",
     "preact-render-to-string": "^5.1.19",

server/chat-commands/admin.ts

@@ -1154,17 +1154,43 @@ export const commands: Chat.ChatCommands = {
 		`/endemergency - Turns off emergency mode. Requires: &`,
 	],
 
-	kill(target, room, user) {
-		this.checkCan('lockdown');
+	async savebattles(target, room, user) {
+		this.checkCan('rangeban'); // admins can restart, so they should be able to do this if needed
+		this.sendReply(`Saving battles...`);
+		const count = await Rooms.global.saveBattles();
+		this.sendReply(`DONE.`);
+		this.sendReply(`${count} battles saved.`);
+		this.addModAction(`${user.name} used /savebattles`);
+	},
 
-		if (Rooms.global.lockdown !== true) {
-			return this.errorReply("For safety reasons, /kill can only be used during lockdown.");
+	async kill(target, room, user) {
+		this.checkCan('lockdown');
+		let noSave = toID(target) === 'nosave';
+		if (!Config.usepostgres) noSave = true;
+
+		if (Rooms.global.lockdown !== true && noSave) {
+			return this.errorReply("For safety reasons, using /kill without saving battles can only be done during lockdown.");
 		}
 
 		if (Monitor.updateServerLock) {
 			return this.errorReply("Wait for /updateserver to finish before using /kill.");
 		}
 
+		if (!noSave) {
+			this.sendReply('Saving battles...');
+			Rooms.global.lockdown = true; // we don't want more battles starting while we save
+			for (const u of Users.users.values()) {
+				u.send(
+					`|pm|&|${u.getIdentity()}|/raw <div class="broadcast-red"><b>The server is restarting soon.</b><br />` +
+					`While battles are being saved, no more can be started. If you're in a battle, it will be paused during saving.<br />` +
+					`After the restart, you will be able to resume your battles from where you left off.`
+				);
+			}
+			const count = await Rooms.global.saveBattles();
+			this.sendReply(`DONE.`);
+			this.sendReply(`${count} battles saved.`);
+		}
+
 		const logRoom = Rooms.get('staff') || Rooms.lobby || room;
 
 		if (!logRoom?.log.roomlogStream) return process.exit();
@@ -1180,7 +1206,10 @@ export const commands: Chat.ChatCommands = {
 			process.exit();
 		}, 10000);
 	},
-	killhelp: [`/kill - kills the server. Can't be done unless the server is in lockdown state. Requires: &`],
+	killhelp: [
+		`/kill - kills the server. Use the argument \`nosave\` to prevent the saving of battles.`,
+		` If this argument is used, the server must be in lockdown. Requires: &`,
+	],
 
 	loadbanlist(target, room, user, connection) {
 		this.checkCan('lockdown');

server/room-battle.ts

@@ -488,6 +488,8 @@ export interface RoomBattleOptions {
 	inputLog?: string;
 	ratedMessage?: string;
 	seed?: PRNGSeed;
+	roomid?: RoomID;
+	players?: ID[];
 }
 
 export class RoomBattle extends RoomGames.RoomGame<RoomBattlePlayer> {
@@ -532,6 +534,8 @@ export class RoomBattle extends RoomGames.RoomGame<RoomBattlePlayer> {
 	turn: number;
 	rqid: number;
 	requestCount: number;
+	options: RoomBattleOptions;
+	frozen?: boolean;
 	dataResolvers?: [((args: string[]) => void), ((error: Error) => void)][];
 	constructor(room: GameRoom, options: RoomBattleOptions) {
 		super(room);
@@ -539,6 +543,7 @@ export class RoomBattle extends RoomGames.RoomGame<RoomBattlePlayer> {
 		this.gameid = 'battle' as ID;
 		this.room = room;
 		this.title = format.name;
+		this.options = options;
 		if (!this.title.endsWith(" Battle")) this.title += " Battle";
 		this.allowRenames = options.allowRenames !== undefined ? !!options.allowRenames : (!options.rated && !options.tour);
 
@@ -633,6 +638,10 @@ export class RoomBattle extends RoomGames.RoomGame<RoomBattlePlayer> {
 		if (Rooms.global.battleCount === 0) Rooms.global.automaticKillRequest();
 	}
 	choose(user: User, data: string) {
+		if (this.frozen) {
+			user.popup(`Your battle is currently paused, so you cannot move right now.`);
+			return;
+		}
 		const player = this.playerTable[user.id];
 		const [choice, rqid] = data.split('|', 2);
 		if (!player) return;
@@ -672,6 +681,10 @@ export class RoomBattle extends RoomGames.RoomGame<RoomBattlePlayer> {
 		void this.stream.write(`>${player.slot} undo`);
 	}
 	joinGame(user: User, slot?: SideID, playerOpts?: {team?: string}) {
+		if (!this.options.players?.includes(user.id)) {
+			user.popup(`You cannot join this battle, as you were not originally playing in it.`);
+			return false;
+		}
 		if (user.id in this.playerTable) {
 			user.popup(`You have already joined this battle.`);
 			return false;

server/rooms.ts

@@ -28,7 +28,7 @@ const LAST_BATTLE_WRITE_THROTTLE = 10;
 
 const RETRY_AFTER_LOGIN = null;
 
-import {FS, Utils, Streams} from '../lib';
+import {FS, Utils, Streams, PostgresDatabase} from '../lib';
 import {RoomSection, RoomSections} from './chat-commands/room-settings';
 import {QueuedHunt} from './chat-plugins/scavengers';
 import {ScavengerGameTemplate} from './chat-plugins/scavenger-games';
@@ -1272,6 +1272,72 @@ export class GlobalRoomState {
 		} catch {}
 		this.lastBattle = Number(lastBattle) || 0;
 		this.lastWrittenBattle = this.lastBattle;
+		void this.loadBattles();
+	}
+
+	async saveBattles() {
+		let count = 0;
+		if (!Config.usepostgres) return 0;
+		const logDatabase = new PostgresDatabase();
+		for (const room of Rooms.rooms.values()) {
+			if (!room.battle || room.battle.ended) continue;
+			room.battle.frozen = true;
+			const log = await room.battle.getLog();
+			const players: ID[] = room.battle.options.players || [];
+			if (!players.length) {
+				for (const num of ['p1', 'p2', 'p3', 'p4'] as const) {
+					if (room.battle[num]?.id) {
+						players.push(room.battle[num].id);
+					}
+				}
+			}
+			if (!players.length || !log?.length) continue; // shouldn't happen???
+			await logDatabase.query(
+				`INSERT INTO stored_battles (roomid, input_log, players, title) VALUES ($1, $2, $3, $4)` +
+				` ON CONFLICT (roomid) DO UPDATE ` +
+				`SET input_log = EXCLUDED.input_log, players = EXCLUDED.players, title = EXCLUDED.title`,
+				[room.roomid, log.join('\n'), players, room.title]
+			);
+			count++;
+		}
+		return count;
+	}
+
+	async loadBattles() {
+		if (!Config.usepostgres) return;
+		const logDatabase = new PostgresDatabase();
+		const query = `DELETE FROM stored_battles WHERE roomid IN (SELECT roomid FROM stored_battles LIMIT 1) RETURNING *`;
+		for await (const battle of logDatabase.stream(query)) {
+			const {input_log, players, roomid, title} = battle;
+			const [, formatid] = roomid.split('-');
+			const room = Rooms.createBattle({
+				format: formatid,
+				inputLog: input_log.join('\n'),
+				roomid,
+				title,
+				players,
+				delayedStart: true,
+			});
+			if (!room || !room.battle) continue; // shouldn't happen???
+			room.battle.start();
+			for (const [i, p] of players.entries()) {
+				room.auth.set(p, Users.PLAYER_SYMBOL);
+				const u = Users.getExact(p);
+				if (u) {
+					room.battle.joinGame(u, `p${i + 1}` as SideID);
+				}
+			}
+		}
+	}
+
+	joinOldBattles(user: User) {
+		for (const room of Rooms.rooms.values()) {
+			const idx = room.battle?.options.players?.indexOf(user.id);
+			if (typeof idx === 'number' && idx > -1) {
+				user.joinRoom(room.roomid);
+				room.battle!.joinGame(user, `p${idx + 1}` as SideID);
+			}
+		}
 	}
 
 	modlog(entry: PartialModlogEntry, overrideID?: string) {
@@ -2000,7 +2066,7 @@ export const Rooms = {
 			options.ratedMessage = p1Special;
 		}
 
-		const roomid = Rooms.global.prepBattleRoom(options.format);
+		const roomid = options.roomid || Rooms.global.prepBattleRoom(options.format);
 		// options.rated is a number representing the lowest player rating, for searching purposes
 		// options.rated < 0 or falsy means "unrated", and will be converted to 0 here
 		// options.rated === true is converted to 1 (used in tests sometimes)
@@ -2015,6 +2081,8 @@ export const Rooms = {
 		} else if (gameType === 'freeforall') {
 			// p1 vs. p2 vs. p3 vs. p4 is too long of a title
 			roomTitle = `${p1name} and friends`;
+		} else if (options.title) {
+			roomTitle = options.title;
 		} else {
 			roomTitle = `${p1name} vs. ${p2name}`;
 		}

server/users.ts

@@ -829,6 +829,7 @@ export class User extends Chat.MessageContext {
 			Punishments.checkName(user, userid, registered);
 
 			Rooms.global.checkAutojoin(user);
+			Rooms.global.joinOldBattles(this);
 			Chat.loginfilter(user, this, userType);
 			return true;
 		}
@@ -844,6 +845,7 @@ export class User extends Chat.MessageContext {
 			return false;
 		}
 		Rooms.global.checkAutojoin(this);
+		Rooms.global.joinOldBattles(this);
 		Chat.loginfilter(this, null, userType);
 		return true;
 	}

sim/battle-stream.ts

@@ -209,6 +209,9 @@ export class BattleStream extends Streams.ObjectReadWriteStream<string> {
 		case 'requestlog':
 			this.push(`requesteddata\n${this.battle!.inputLog.join('\n')}`);
 			break;
+		case 'requestexport':
+			this.push(`requesteddata\n${this.battle!.prngSeed}\n${this.battle!.inputLog.join('\n')}`);
+			break;
 		case 'requestteam':
 			message = message.trim();
 			const slotNum = parseInt(message.slice(1)) - 1;

test/lib/postgres.js

@@ -0,0 +1,47 @@
+"use strict";
+const assert = require('assert').strict;
+const {PostgresDatabase} = require('../../lib');
+
+function testMod(mod) {
+	try {
+		require(mod);
+	} catch (e) {
+		return it.skip;
+	}
+	return it;
+}
+
+// only run these if you already have postgres configured
+describe.skip("Postgres features", () => {
+	it("Should be able to connect to a database", async () => {
+		this.database = new PostgresDatabase();
+	});
+	it("Should be able to insert data", async () => {
+		await assert.doesNotThrowAsync(async () => {
+			await this.database.query(`CREATE TABLE test (col TEXT, col2 TEXT)`);
+			await this.database.query(
+				`INSERT INTO test (col, col2) VALUES ($1, $2)`,
+				['foo', 'bar'],
+			);
+		});
+	});
+	testMod('sql-template-strings')('Should support sql-template-strings', async () => {
+		await assert.doesNotThrowAsync(async () => {
+			const SQL = require('sql-template-strings');
+			await this.database.query(SQL`INSERT INTO test (col1, col2) VALUES (${'a'}, ${'b'})`);
+		});
+	});
+	it("Should be able to run multiple statements in transaction", async () => {
+		await assert.doesNotThrowAsync(async () => {
+			await this.database.transaction(async worker => {
+				const tables = await worker.query(
+					`SELECT tablename FROM pg_catalog.pg_tables ` +
+					`WHERE tablename = 'test' LIMIT 1;`
+				);
+				for (const {tablename} of tables) {
+					await worker.query(`DROP TABLE ` + tablename);
+				}
+			});
+		});
+	});
+});