From 1b7a81e41aaf1c0afc4eeda0973b9982644dee4a Mon Sep 17 00:00:00 2001
From: Stephen Paul Weber <singpolyma@singpolyma.net>
Date: Wed, 26 Jun 2024 21:21:50 -0500
Subject: [PATCH] If we get a presence flood, fetch each caps once

Since the query may still be in flight, we won't see the result in the
DB, but we still don't want to send out another query for the same caps
hash right now.
---
 snikket/Client.hx | 37 +++++++++++++++++++++++++++----------
 snikket/Map.js.hx |  4 ++++
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/snikket/Client.hx b/snikket/Client.hx
index cff1304..fdd919d 100644
--- a/snikket/Client.hx
+++ b/snikket/Client.hx
@@ -69,6 +69,7 @@ class Client extends EventEmitter {
 	);
 	private var _displayName: String;
 	private var fastMechanism: Null<String> = null;
+	private final pendingCaps: Map<String, Array<(Null<Caps>)->Chat>> = [];
 
 	/**
 		Create a new Client to connect to a particular account
@@ -342,20 +343,36 @@ class Client extends EventEmitter {
 					persistence.storeChat(accountId(), chat);
 					if (chat.livePresence()) this.trigger("chats/update", [chat]);
 				} else {
+					final handleCaps = (caps) -> {
+						chat.setPresence(JID.parse(stanza.attr.get("from")).resource, new Presence(caps, mucUser));
+						persistence.storeChat(accountId(), chat);
+						return chat;
+					};
+
 					persistence.getCaps(c.attr.get("ver"), (caps) -> {
 						if (caps == null) {
-							final discoGet = new DiscoInfoGet(stanza.attr.get("from"), c.attr.get("node") + "#" + c.attr.get("ver"));
-							discoGet.onFinished(() -> {
-								chat.setPresence(JID.parse(stanza.attr.get("from")).resource, new Presence(discoGet.getResult(), mucUser));
-								if (discoGet.getResult() != null) persistence.storeCaps(discoGet.getResult());
-								persistence.storeChat(accountId(), chat);
+							final pending = pendingCaps.get(c.attr.get("ver"));
+							if (pending == null) {
+								pendingCaps.set(c.attr.get("ver"), [handleCaps]);
+								final discoGet = new DiscoInfoGet(stanza.attr.get("from"), c.attr.get("node") + "#" + c.attr.get("ver"));
+								discoGet.onFinished(() -> {
+									final chatsToUpdate: Map<String, Chat> = [];
+									final handlers = pendingCaps.get(c.attr.get("ver")) ?? [];
+									pendingCaps.remove(c.attr.get("ver"));
+									if (discoGet.getResult() != null) persistence.storeCaps(discoGet.getResult());
+									for (handler in handlers) {
+										final c = handler(discoGet.getResult());
+										if (c.livePresence()) chatsToUpdate.set(c.chatId, c);
+									}
+									this.trigger("chats/update", Lambda.array({ iterator: () -> chatsToUpdate.iterator() }));
+								});
+								sendQuery(discoGet);
+							} else {
+								pending.push(handleCaps);
 								if (chat.livePresence()) this.trigger("chats/update", [chat]);
-							});
-							sendQuery(discoGet);
+							}
 						} else {
-							chat.setPresence(JID.parse(stanza.attr.get("from")).resource, new Presence(caps, mucUser));
-							persistence.storeChat(accountId(), chat);
-							if (chat.livePresence()) this.trigger("chats/update", [chat]);
+							handleCaps(caps);
 						}
 					});
 				}
diff --git a/snikket/Map.js.hx b/snikket/Map.js.hx
index f67920d..6f03b83 100644
--- a/snikket/Map.js.hx
+++ b/snikket/Map.js.hx
@@ -12,6 +12,10 @@ abstract Map<K,V>(NativeMap<K,V>) {
 		this.set(k, v);
 	}
 
+	public inline function get(k:K):Null<V> {
+		return this.get(k);
+	}
+
 	public inline function exists(k:K):Bool {
 		return this.has(k);
 	}