Merge pull request #115 from Gerrit0/master

fix: Add trash can in addMessageOwnership

Author: robertt

src/modules/playground.ts

@@ -12,14 +12,15 @@ import {
 	findCodeblockFromChannel,
 	PLAYGROUND_REGEX,
 } from '../util/findCodeblockFromChannel';
+import { LimitedSizeMap } from '../util/limitedSizeMap';
 import { addMessageOwnership, sendWithMessageOwnership } from '../util/send';
 
 export class PlaygroundModule extends Module {
 	constructor(client: CookiecordClient) {
 		super(client);
 	}
 
-	private editedLongLink = new Map<string, Message>();
+	private editedLongLink = new LimitedSizeMap<string, Message>(1000);
 
 	@command({
 		aliases: ['pg', 'playg'],
@@ -67,9 +68,10 @@ export class PlaygroundModule extends Module {
 				{ embed },
 			);
 			this.editedLongLink.set(msg.id, botMsg);
-			addMessageOwnership(botMsg, msg.author);
+			await addMessageOwnership(botMsg, msg.author);
 		}
 	}
+
 	@listener({ event: 'messageUpdate' })
 	async onLongFix(_oldMsg: Message, msg: Message) {
 		if (msg.partial) await msg.fetch();

src/util/limitedSizeMap.ts

@@ -0,0 +1,22 @@
+/**
+ * A FIFO style limited size map, used to prevent memory from growing unboundedly when tracking messages.
+ */
+export class LimitedSizeMap<K, V> extends Map<K, V> implements Map<K, V> {
+	private _maxSize: number;
+
+	constructor(maxSize: number) {
+		super();
+		this._maxSize = maxSize;
+	}
+
+	set(key: K, val: V) {
+		super.set(key, val);
+
+		if (this.size > this._maxSize) {
+			// Keys returns an iterable in insertion order, so this removes the oldest entry from the map.
+			this.delete(this.keys().next().value);
+		}
+
+		return this;
+	}
+}

src/util/send.ts

@@ -1,8 +1,7 @@
 import { Message, MessageEmbed, User } from 'discord.js';
+import { LimitedSizeMap } from './limitedSizeMap';
 
-const MAX_TRACKED_MESSAGES = 1000;
-
-const messageToUserId = new Map<string, string>();
+const messageToUserId = new LimitedSizeMap<string, string>(1000);
 
 export const DELETE_EMOJI = '🗑️';
 
@@ -11,18 +10,13 @@ export async function sendWithMessageOwnership(
 	toSend: string | { embed: MessageEmbed },
 ) {
 	const sent = await message.channel.send(toSend);
-	await sent.react(DELETE_EMOJI);
-
-	addMessageOwnership(sent, message.author);
+	await addMessageOwnership(sent, message.author);
 }
 
-export function addMessageOwnership(message: Message, user: User) {
+export async function addMessageOwnership(message: Message, user: User) {
+	await message.react(DELETE_EMOJI);
+
 	messageToUserId.set(message.id, user.id);
-	// Without this memory grows unboundedly... very slowly, but better to avoid the issue.
-	if (messageToUserId.size > MAX_TRACKED_MESSAGES) {
-		// Keys returns an iterable in insertion order, so we remove the oldest message from the map.
-		messageToUserId.delete(messageToUserId.keys().next().value);
-	}
 }
 
 export function ownsBotMessage(message: Message, userId: string) {