From e9f494ec5db4884ac1a47df1795a30b5b46b5cb4 Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Thu, 16 Jun 2022 14:53:23 +0530
Subject: [PATCH] feat: more efficient message send

---
 src/Socket/messages-send.ts | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts
index 7ea3fed..04ed8f0 100644
--- a/src/Socket/messages-send.ts
+++ b/src/Socket/messages-send.ts
@@ -211,6 +211,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 	}
 
 	const assertSessions = async(jids: string[], force: boolean) => {
+		let didFetchNewSession = false
 		let jidsRequiringFetch: string[] = []
 		if(force) {
 			jidsRequiringFetch = jids
@@ -248,15 +249,13 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 				]
 			})
 			await parseAndInjectE2ESessions(result, authState)
-			return true
+
+			didFetchNewSession = true
 		}
 
-		return false
-	}
-
-	const createParticipantNodes = async(jids: string[], bytes: Buffer) => {
-		await assertSessions(jids, false)
-
+		// pre-fetch all sessions to improve efficiency
+		// makes a big difference when sending to large groups
+		// or when there are 100s of devices to push to
 		if(authState.keys.isInTransaction()) {
 			await authState.keys.prefetch(
 				'session',
@@ -264,7 +263,11 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 			)
 		}
 
-		const nodes = await Promise.all(
+		return didFetchNewSession
+	}
+
+	const createParticipantNodes = (jids: string[], bytes: Buffer) => (
+		Promise.all(
 			jids.map(
 				async jid => {
 					const { type, ciphertext } = await encryptSignalProto(jid, bytes, authState)
@@ -281,8 +284,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 				}
 			)
 		)
-		return nodes
-	}
+	)
 
 	const relayMessage = async(
 		jid: string,
@@ -365,6 +367,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 							}
 						})
 
+						await assertSessions(senderKeyJids, false)
+
 						participants.push(
 							...(await createParticipantNodes(senderKeyJids, encSenderKeyMsg))
 						)
@@ -395,6 +399,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 						devices.push(...additionalDevices)
 					}
 
+					const allJids: string[] = []
 					const meJids: string[] = []
 					const otherJids: string[] = []
 					for(const { user, device } of devices) {
@@ -405,8 +410,12 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 						} else {
 							otherJids.push(jid)
 						}
+
+						allJids.push(jid)
 					}
 
+					await assertSessions(allJids, false)
+
 					const [meNodes, otherNodes] = await Promise.all([
 						createParticipantNodes(meJids, encodedMeMsg),
 						createParticipantNodes(otherJids, encodedMsg)