From 6735263696472181de7d41b0b1f674c8d60451c5 Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Tue, 20 Sep 2022 12:16:05 +0530
Subject: [PATCH] feat: add makeCacheableSignalKeyStore util

---
 Example/example.ts      |  8 +++--
 src/Types/Auth.ts       |  2 ++
 src/Utils/auth-utils.ts | 74 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/Example/example.ts b/Example/example.ts
index 1ae6bdd..89b374e 100644
--- a/Example/example.ts
+++ b/Example/example.ts
@@ -1,5 +1,5 @@
 import { Boom } from '@hapi/boom'
-import makeWASocket, { AnyMessageContent, delay, DisconnectReason, fetchLatestBaileysVersion, makeInMemoryStore, MessageRetryMap, useMultiFileAuthState } from '../src'
+import makeWASocket, { AnyMessageContent, delay, DisconnectReason, fetchLatestBaileysVersion, makeCacheableSignalKeyStore, makeInMemoryStore, MessageRetryMap, useMultiFileAuthState } from '../src'
 import MAIN_LOGGER from '../src/Utils/logger'
 
 const logger = MAIN_LOGGER.child({ })
@@ -32,7 +32,11 @@ const startSock = async() => {
 		version,
 		logger,
 		printQRInTerminal: true,
-		auth: state,
+		auth: {
+			creds: state.creds,
+			/** caching makes the store faster to send/recv messages */
+			keys: makeCacheableSignalKeyStore(state.keys, logger),
+		},
 		msgRetryCounterMap,
 		generateHighQualityLinkPreview: true,
 		// implement to handle retries
diff --git a/src/Types/Auth.ts b/src/Types/Auth.ts
index 2ccb9db..4d84348 100644
--- a/src/Types/Auth.ts
+++ b/src/Types/Auth.ts
@@ -69,6 +69,8 @@ type Awaitable<T> = T | Promise<T>
 export type SignalKeyStore = {
     get<T extends keyof SignalDataTypeMap>(type: T, ids: string[]): Awaitable<{ [id: string]: SignalDataTypeMap[T] }>
     set(data: SignalDataSet): Awaitable<void>
+    /** clear all the data in the store */
+    clear?(): Awaitable<void>
 }
 
 export type SignalKeyStoreWithTransaction = SignalKeyStore & {
diff --git a/src/Utils/auth-utils.ts b/src/Utils/auth-utils.ts
index 17f09c2..433bc15 100644
--- a/src/Utils/auth-utils.ts
+++ b/src/Utils/auth-utils.ts
@@ -1,10 +1,78 @@
 import { Boom } from '@hapi/boom'
 import { randomBytes } from 'crypto'
+import NodeCache from 'node-cache'
 import type { Logger } from 'pino'
 import type { AuthenticationCreds, SignalDataSet, SignalDataTypeMap, SignalKeyStore, SignalKeyStoreWithTransaction, TransactionCapabilityOptions } from '../Types'
 import { Curve, signedKeyPair } from './crypto'
 import { delay, generateRegistrationId } from './generics'
 
+/**
+ * Adds caching capability to a SignalKeyStore
+ * @param store the store to add caching to
+ * @param logger to log trace events
+ * @param opts NodeCache options
+ */
+export function makeCacheableSignalKeyStore(
+	store: SignalKeyStore,
+	logger: Logger,
+	opts?: NodeCache.Options
+): SignalKeyStore {
+	const cache = new NodeCache({
+		...opts || { },
+		useClones: false,
+	})
+
+	function getUniqueId(type: string, id: string) {
+		return `${type}.${id}`
+	}
+
+	return {
+		async get(type, ids) {
+			const data: { [_: string]: SignalDataTypeMap[typeof type] } = { }
+			const idsToFetch: string[] = []
+			for(const id of ids) {
+				const item = cache.get<SignalDataTypeMap[typeof type]>(getUniqueId(type, id))
+				if(typeof item !== 'undefined') {
+					data[id] = item
+				} else {
+					idsToFetch.push(id)
+				}
+			}
+
+			if(idsToFetch.length) {
+				logger.trace({ items: idsToFetch.length }, 'loading from store')
+				const fetched = await store.get(type, idsToFetch)
+				for(const id of idsToFetch) {
+					const item = fetched[id]
+					if(item) {
+						data[id] = item
+						cache.set(getUniqueId(type, id), item)
+					}
+				}
+			}
+
+			return data
+		},
+		async set(data) {
+			let keys = 0
+			for(const type in data) {
+				for(const id in data[type]) {
+					cache.set(getUniqueId(type, id), data[type][id])
+					keys += 1
+				}
+			}
+
+			logger.trace({ keys }, 'updated cache')
+
+			await store.set(data)
+		},
+		async clear() {
+			cache.flushAll()
+			await store.clear?.()
+		}
+	}
+}
+
 /**
  * Adds DB like transaction capability (https://en.wikipedia.org/wiki/Database_transaction) to the SignalKeyStore,
  * this allows batch read & write operations & improves the performance of the lib
@@ -12,7 +80,11 @@ import { delay, generateRegistrationId } from './generics'
  * @param logger logger to log events
  * @returns SignalKeyStore with transaction capability
  */
-export const addTransactionCapability = (state: SignalKeyStore, logger: Logger, { maxCommitRetries, delayBetweenTriesMs }: TransactionCapabilityOptions): SignalKeyStoreWithTransaction => {
+export const addTransactionCapability = (
+	state: SignalKeyStore,
+	logger: Logger,
+	{ maxCommitRetries, delayBetweenTriesMs }: TransactionCapabilityOptions
+): SignalKeyStoreWithTransaction => {
 	let inTransaction = false
 	// number of queries made to the DB during the transaction
 	// only there for logging purposes