fix: optimize LID migration performance and introduce cache (#1782)

* chore: lid-map cache, bulk migration, fixes * fix: type error * fix: signalrepository with lid type, remove check from assertsessions, lid-mapping store level deduplication * fix: migrations logs * fix: remove migrationsops lenght check * fix: skip deletesession if cached or existing migrations * chore: remove check migrated session * fix: remove return null loadsession * fix: getLIDforPN call and jidsrequiringfetch * fix: longer lid map cache TTL * fix: lint * fix: lid type * fix: lid socket type
2025-09-14 15:58:59 -03:00
parent 42c980778e
commit 194b557b0f
10 changed files with 241 additions and 220 deletions
@@ -1,9 +1,9 @@
+/* @ts-ignore */
 import * as libsignal from 'libsignal'
-import { LRUCache } from 'lru-cache'
 import type { SignalAuthState, SignalKeyStoreWithTransaction } from '../Types'
-import type { SignalRepository } from '../Types/Signal'
+import type { SignalRepositoryWithLIDStore } from '../Types/Signal'
 import { generateSignalPubKey } from '../Utils'
-import { jidDecode } from '../WABinary'
+import { jidDecode, transferDevice } from '../WABinary'
 import type { SenderKeyStore } from './Group/group_cipher'
 import { SenderKeyName } from './Group/sender-key-name'
 import { SenderKeyRecord } from './Group/sender-key-record'
@@ -20,7 +20,7 @@ export function makeLibSignalRepository(
 		  }[]
 		| undefined
 	>
-): SignalRepository {
+): SignalRepositoryWithLIDStore {
 	const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction, onWhatsAppFunc)
 	const storage = signalStorage(auth, lidMapping)

@@ -38,13 +38,7 @@ export function makeLibSignalRepository(
 		)
 	}

-	// Simple operation-level deduplication (5 minutes)
-	const recentMigrations = new LRUCache<string, boolean>({
-		max: 500,
-		ttl: 5 * 60 * 1000
-	})
-
-	const repository: SignalRepository = {
+	const repository: SignalRepositoryWithLIDStore = {
 		decryptGroupMessage({ group, authorJid, msg }) {
 			const senderName = jidToSignalSenderKeyName(group, authorJid)
 			const cipher = new GroupCipher(storage, senderName)
@@ -136,7 +130,7 @@ export function makeLibSignalRepository(

 						if (pnSession) {
 							// Migrate PN to LID
-							await repository.migrateSession(jid, lidForPN)
+							await repository.migrateSession([jid], lidForPN)
 							encryptionJid = lidForPN
 						}
 					}
@@ -185,13 +179,8 @@ export function makeLibSignalRepository(
 			return jidToSignalProtocolAddress(jid).toString()
 		},

-		async storeLIDPNMapping(lid: string, pn: string) {
-			await lidMapping.storeLIDPNMapping(lid, pn)
-		},
-
-		getLIDMappingStore() {
-			return lidMapping
-		},
+		// Optimized direct access to LID mapping store
+		lidMapping,

 		async validateSession(jid: string) {
 			try {
@@ -212,71 +201,97 @@ export function makeLibSignalRepository(
 			}
 		},

-		async deleteSession(jid: string) {
-			const addr = jidToSignalProtocolAddress(jid)
+		async deleteSession(jids: string[]) {
+			if (!jids.length) return

+			// Convert JIDs to signal addresses and prepare for bulk deletion
+			const sessionUpdates: { [key: string]: null } = {}
+			jids.forEach(jid => {
+				const addr = jidToSignalProtocolAddress(jid)
+				sessionUpdates[addr.toString()] = null
+			})
+
+			// Single transaction for all deletions
 			return parsedKeys.transaction(async () => {
-				await auth.keys.set({ session: { [addr.toString()]: null } })
-			}, jid)
+				await auth.keys.set({ session: sessionUpdates })
+			}, `delete-${jids.length}-sessions`)
 		},

-		async migrateSession(fromJid: string, toJid: string) {
-			// Only migrate PN → LID
-			if (!fromJid.includes('@s.whatsapp.net') || !toJid.includes('@lid')) {
-				return
-			}
+		async migrateSession(
+			fromJids: string[],
+			toJid: string
+		): Promise<{ migrated: number; skipped: number; total: number }> {
+			if (!fromJids.length || !toJid.includes('@lid')) return { migrated: 0, skipped: 0, total: 0 }

-			const fromDecoded = jidDecode(fromJid)
-			const toDecoded = jidDecode(toJid)
-			if (!fromDecoded || !toDecoded) return
+			// Filter valid PN JIDs
+			const validJids = fromJids.filter(jid => jid.includes('@s.whatsapp.net'))
+			if (!validJids.length) return { migrated: 0, skipped: 0, total: fromJids.length }

-			const deviceId = fromDecoded.device || 0
-			const migrationKey = `${fromDecoded.user}.${deviceId}→${toDecoded.user}.${deviceId}`
+			// Single optimized transaction for all migrations
+			return parsedKeys.transaction(
+				async (): Promise<{ migrated: number; skipped: number; total: number }> => {
+					// 1. Batch store all LID mappings
+					const mappings = validJids.map(jid => ({
+						lid: transferDevice(jid, toJid),
+						pn: jid
+					}))
+					await lidMapping.storeLIDPNMappings(mappings)

-			// Check if recently migrated (5 min window)
-			if (recentMigrations.has(migrationKey)) {
-				return
-			}
+					// 2. Prepare migration operations
+					const migrationOps = validJids.map(jid => {
+						const lidWithDevice = transferDevice(jid, toJid)
+						const fromDecoded = jidDecode(jid)!
+						const toDecoded = jidDecode(lidWithDevice)!

-			// Check if LID session already exists
-			const lidAddr = jidToSignalProtocolAddress(toJid)
-			const { [lidAddr.toString()]: lidExists } = await auth.keys.get('session', [lidAddr.toString()])
-			if (lidExists) {
-				recentMigrations.set(migrationKey, true)
-				return
-			}
+						return {
+							fromJid: jid,
+							toJid: lidWithDevice,
+							pnUser: fromDecoded.user,
+							lidUser: toDecoded.user,
+							deviceId: fromDecoded.device || 0,
+							fromAddr: jidToSignalProtocolAddress(jid),
+							toAddr: jidToSignalProtocolAddress(lidWithDevice)
+						}
+					})

-			return parsedKeys.transaction(async () => {
-				// Store mapping
-				await lidMapping.storeLIDPNMapping(toJid, fromJid)
+					// 3. Batch check which LID sessions already exist
+					const lidAddrs = migrationOps.map(op => op.toAddr.toString())
+					const existingSessions = await auth.keys.get('session', lidAddrs)

-				// Load and copy session
-				const fromAddr = jidToSignalProtocolAddress(fromJid)
-				const fromSession = await storage.loadSession(fromAddr.toString())
+					// 4. Filter out sessions that already have LID sessions
+					const opsToMigrate = migrationOps.filter(op => !existingSessions[op.toAddr.toString()])
+					const skippedCount = migrationOps.length - opsToMigrate.length

-				if (fromSession?.haveOpenSession()) {
-					// Deep copy session to prevent reference issues
-					const sessionBytes = fromSession.serialize()
-					const copiedSession = libsignal.SessionRecord.deserialize(sessionBytes)
+					if (!opsToMigrate.length) {
+						return { migrated: 0, skipped: skippedCount, total: validJids.length }
+					}

-					// Store at LID address
-					await storage.storeSession(lidAddr.toString(), copiedSession)
+					// 5. Execute all migrations in parallel
+					await Promise.all(
+						opsToMigrate.map(async op => {
+							const fromSession = await storage.loadSession(op.fromAddr.toString())

-					// Delete PN session - maintain single encryption layer
-					await auth.keys.set({ session: { [fromAddr.toString()]: null } })
-				}
+							if (fromSession?.haveOpenSession()) {
+								// Copy session to LID address
+								const sessionBytes = fromSession.serialize()
+								const copiedSession = libsignal.SessionRecord.deserialize(sessionBytes)
+								await storage.storeSession(op.toAddr.toString(), copiedSession)

-				recentMigrations.set(migrationKey, true)
-			}, fromJid)
+								// Delete PN session
+								await auth.keys.set({ session: { [op.fromAddr.toString()]: null } })
+							}
+						})
+					)
+
+					return { migrated: opsToMigrate.length, skipped: skippedCount, total: validJids.length }
+				},
+				`migrate-${validJids.length}-sessions-${jidDecode(toJid)?.user}`
+			)
 		},

 		async encryptMessageWithWire({ encryptionJid, wireJid, data }) {
 			const result = await repository.encryptMessage({ jid: encryptionJid, data })
 			return { ...result, wireJid }
-		},
-
-		destroy() {
-			recentMigrations.clear()
 		}
 	}

@@ -287,6 +302,12 @@ const jidToSignalProtocolAddress = (jid: string): libsignal.ProtocolAddress => {
 	const decoded = jidDecode(jid)!
 	const { user, device, server } = decoded

+	if (!user) {
+		throw new Error(
+			`JID decoded but user is empty: "${jid}" -> user: "${user}", server: "${server}", device: ${device}`
+		)
+	}
+
 	// LID addresses get _1 suffix for Signal protocol
 	const signalUser = server === 'lid' ? `${user}_1` : user
 	const finalDevice = device || 0
@@ -1,9 +1,14 @@
+import { LRUCache } from 'lru-cache'
 import type { SignalKeyStoreWithTransaction } from '../Types'
 import logger from '../Utils/logger'
 import { isLidUser, isPnUser, jidDecode } from '../WABinary'

-//TODO: Caching
 export class LIDMappingStore {
+	private readonly mappingCache = new LRUCache<string, string>({
+		ttl: 7 * 24 * 60 * 60 * 1000, // 7 days
+		ttlAutopurge: true,
+		updateAgeOnGet: true
+	})
 	private readonly keys: SignalKeyStoreWithTransaction
 	private onWhatsAppFunc?: (...jids: string[]) => Promise<
 		| {
@@ -29,13 +34,6 @@ export class LIDMappingStore {
 		this.onWhatsAppFunc = onWhatsAppFunc // needed to get LID from PN if not found
 	}

-	/**
-	 * Store LID-PN mapping - USER LEVEL
-	 */
-	async storeLIDPNMapping(lid: string, pn: string): Promise<void> {
-		return this.storeLIDPNMappings([{ lid, pn }])
-	}
-
 	/**
 	 * Store LID-PN mapping - USER LEVEL
 	 */
@@ -57,6 +55,25 @@ export class LIDMappingStore {

 			const pnUser = pnDecoded.user
 			const lidUser = lidDecoded.user
+
+			// Check if mapping already exists (cache first, then database)
+			let existingLidUser = this.mappingCache.get(`pn:${pnUser}`)
+			if (!existingLidUser) {
+				// Cache miss - check database
+				const stored = await this.keys.get('lid-mapping', [pnUser])
+				existingLidUser = stored[pnUser]
+				if (existingLidUser) {
+					// Update cache with database value
+					this.mappingCache.set(`pn:${pnUser}`, existingLidUser)
+					this.mappingCache.set(`lid:${existingLidUser}`, pnUser)
+				}
+			}
+
+			if (existingLidUser === lidUser) {
+				logger.debug({ pnUser, lidUser }, 'LID mapping already exists, skipping')
+				continue
+			}
+
 			pairMap[pnUser] = lidUser
 		}

@@ -70,6 +87,10 @@ export class LIDMappingStore {
 						[`${lidUser}_reverse`]: pnUser // "102765716062358_reverse" -> "554396160286"
 					}
 				})
+
+				// Update cache with both directions
+				this.mappingCache.set(`pn:${pnUser}`, lidUser)
+				this.mappingCache.set(`lid:${lidUser}`, pnUser)
 			}
 		}, 'lid-mapping')
 	}
@@ -83,22 +104,38 @@ export class LIDMappingStore {
 		const decoded = jidDecode(pn)
 		if (!decoded) return null

-		// Look up user-level mapping (whatsmeow approach)
+		// Check cache first for PN → LID mapping
 		const pnUser = decoded.user
-		const stored = await this.keys.get('lid-mapping', [pnUser])
-		let lidUser = stored[pnUser]
+		let lidUser = this.mappingCache.get(`pn:${pnUser}`)

 		if (!lidUser) {
-			logger.trace(`No LID mapping found for PN user ${pnUser}; getting from USync`)
-			const { exists, lid } = (await this.onWhatsAppFunc?.(pn))?.[0]! // this function already adds LIDs to mapping
-			if (exists) {
-				lidUser = jidDecode(lid)?.user
+			// Cache miss - check database
+			const stored = await this.keys.get('lid-mapping', [pnUser])
+			lidUser = stored[pnUser]
+
+			if (lidUser) {
+				// Cache the database result
+				this.mappingCache.set(`pn:${pnUser}`, lidUser)
 			} else {
-				return null
+				// Not in database - try USync
+				logger.trace(`No LID mapping found for PN user ${pnUser}; getting from USync`)
+				const { exists, lid } = (await this.onWhatsAppFunc?.(pn))?.[0]! // this function already adds LIDs to mapping
+				if (exists && lid) {
+					lidUser = jidDecode(lid)?.user
+					if (lidUser) {
+						// Cache the USync result
+						this.mappingCache.set(`pn:${pnUser}`, lidUser)
+					}
+				} else {
+					return null
+				}
 			}
 		}

-		if (typeof lidUser !== 'string') return null
+		if (typeof lidUser !== 'string' || !lidUser) {
+			logger.warn(`Invalid or empty LID user for PN ${pn}: lidUser = "${lidUser}"`)
+			return null
+		}

 		// Push the PN device ID to the LID to maintain device separation
 		const pnDevice = decoded.device !== undefined ? decoded.device : 0
@@ -117,15 +154,21 @@ export class LIDMappingStore {
 		const decoded = jidDecode(lid)
 		if (!decoded) return null

-		// Look up reverse user mapping
+		// Check cache first for LID → PN mapping
 		const lidUser = decoded.user
-		// TODO: remove this style and instead load all mappings somehow, and then assign them in the map
-		const stored = await this.keys.get('lid-mapping', [`${lidUser}_reverse`])
-		const pnUser = stored[`${lidUser}_reverse`]
+		let pnUser = this.mappingCache.get(`lid:${lidUser}`)

 		if (!pnUser || typeof pnUser !== 'string') {
-			logger.trace(`No reverse mapping found for LID user: ${lidUser}`)
-			return null
+			// Cache miss - check database
+			const stored = await this.keys.get('lid-mapping', [`${lidUser}_reverse`])
+			pnUser = stored[`${lidUser}_reverse`]
+
+			if (!pnUser || typeof pnUser !== 'string') {
+				logger.trace(`No reverse mapping found for LID user: ${lidUser}`)
+				return null
+			}
+
+			this.mappingCache.set(`lid:${lidUser}`, pnUser)
 		}

 		// Construct device-specific PN JID