feat: add makeCacheableSignalKeyStore util

This commit is contained in:
Adhiraj Singh
2022-09-20 12:16:05 +05:30
parent 28afa96cb9
commit 6735263696
3 changed files with 81 additions and 3 deletions

View File

@@ -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

View File

@@ -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 & {

View File

@@ -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