ruicesar/InfiniteAPI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

auth-utils: Make a better addTransactionCapability function

Browse Source
This commit is contained in:
Rajeh Taher
2025-10-03 00:03:21 +03:00
parent 39541cc162
commit 1060f44bff
4 changed files with 299 additions and 379 deletions
Download Patch File Download Diff File Expand all files Collapse all files
package.json
+1
View File
@@ -46,6 +46,7 @@
"libsignal": "git+https://github.com/whiskeysockets/libsignal-node",
"lru-cache": "^11.1.0",
"music-metadata": "^11.7.0",
"p-queue": "^9.0.0",
"pino": "^9.6",
"protobufjs": "^7.2.4",
"ws": "^8.13.0"
src/Utils/auth-utils.ts
+147 -379
View File
@@ -1,7 +1,8 @@
import NodeCache from '@cacheable/node-cache'
import { AsyncLocalStorage } from 'async_hooks'
import { Mutex } from 'async-mutex'
import { randomBytes } from 'crypto'
import { LRUCache } from 'lru-cache'
import PQueue from 'p-queue'
import { DEFAULT_CACHE_TTLS } from '../Defaults'
import type {
AuthenticationCreds,
@@ -15,6 +16,16 @@ import type {
import { Curve, signedKeyPair } from './crypto'
import { delay, generateRegistrationId } from './generics'
import type { ILogger } from './logger'
import { PreKeyManager } from './pre-key-manager'
/**
* Transaction context stored in AsyncLocalStorage
*/
interface TransactionContext {
cache: SignalDataSet
mutations: SignalDataSet
dbQueries: number
}
/**
* Adds caching capability to a SignalKeyStore
@@ -47,6 +58,7 @@ export function makeCacheableSignalKeyStore(
return cacheMutex.runExclusive(async () => {
const data: { [_: string]: SignalDataTypeMap[typeof type] } = {}
const idsToFetch: string[] = []
for (const id of ids) {
const item = (await cache.get<SignalDataTypeMap[typeof type]>(getUniqueId(type, id))) as any
if (typeof item !== 'undefined') {
@@ -82,7 +94,6 @@ export function makeCacheableSignalKeyStore(
}
logger?.trace({ keys }, 'updated cache')
await store.set(data)
})
},
@@ -93,184 +104,9 @@ export function makeCacheableSignalKeyStore(
}
}
// Module-level specialized mutexes for pre-key operations
const preKeyMutex = new Mutex()
const signedPreKeyMutex = new Mutex()
/**
* Get the appropriate mutex for the key type
*/
const getPreKeyMutex = (keyType: string): Mutex => {
return keyType === 'signed-pre-key' ? signedPreKeyMutex : preKeyMutex
}
/**
* Handles pre-key operations with mutex protection
*/
async function handlePreKeyOperations(
data: SignalDataSet,
keyType: keyof SignalDataTypeMap,
transactionCache: SignalDataSet,
mutations: SignalDataSet,
logger: ILogger,
isInTransaction: boolean,
state?: SignalKeyStore
): Promise<void> {
const mutex = getPreKeyMutex(keyType)
await mutex.runExclusive(async () => {
const keyData = data[keyType]
if (!keyData) return
// Ensure structures exist
transactionCache[keyType] = transactionCache[keyType] || ({} as any)
mutations[keyType] = mutations[keyType] || ({} as any)
// Separate deletions from updates for batch processing
const deletionKeys: string[] = []
const updateKeys: string[] = []
for (const keyId in keyData) {
if (keyData[keyId] === null) {
deletionKeys.push(keyId)
} else {
updateKeys.push(keyId)
}
}
// Process updates first (no validation needed)
for (const keyId of updateKeys) {
if (transactionCache[keyType]) {
transactionCache[keyType][keyId] = keyData[keyId]!
}
if (mutations[keyType]) {
mutations[keyType][keyId] = keyData[keyId]!
}
}
// Process deletions with validation
if (deletionKeys.length === 0) return
if (isInTransaction) {
// In transaction, only allow deletion if key exists in cache
for (const keyId of deletionKeys) {
if (transactionCache[keyType]) {
transactionCache[keyType][keyId] = null
if (mutations[keyType]) {
// Mark for deletion in mutations
mutations[keyType][keyId] = null
}
} else {
logger.warn(`Skipping deletion of non-existent ${keyType} in transaction: ${keyId}`)
}
}
return
}
// Outside transaction, batch validate all deletions
if (!state) return
const existingKeys = await state.get(keyType, deletionKeys)
for (const keyId of deletionKeys) {
if (existingKeys[keyId]) {
if (transactionCache[keyType]) transactionCache[keyType][keyId] = null
if (mutations[keyType]) mutations[keyType][keyId] = null
} else {
logger.warn(`Skipping deletion of non-existent ${keyType}: ${keyId}`)
}
}
})
}
/**
* Handles normal key operations for transactions
*/
function handleNormalKeyOperations(
data: SignalDataSet,
key: keyof SignalDataTypeMap,
transactionCache: SignalDataSet,
mutations: SignalDataSet
) {
Object.assign(transactionCache[key]!, data[key])
mutations[key] = mutations[key] || ({} as any)
Object.assign(mutations[key]!, data[key])
}
/**
* Process pre-key deletions with validation
*/
async function processPreKeyDeletions(
data: SignalDataSet,
keyType: keyof SignalDataTypeMap,
state: SignalKeyStore,
logger: ILogger
): Promise<void> {
const mutex = getPreKeyMutex(keyType)
await mutex.runExclusive(async () => {
const keyData = data[keyType]
if (!keyData) return
// Validate deletions
for (const keyId in keyData) {
if (keyData[keyId] === null) {
const existingKeys = await state.get(keyType, [keyId])
if (!existingKeys[keyId]) {
logger.warn(`Skipping deletion of non-existent ${keyType}: ${keyId}`)
if (data[keyType]) delete data[keyType][keyId]
}
}
}
})
}
/**
* Executes a function with mutexes acquired for given key types
* Uses async-mutex's runExclusive with efficient batching
*/
async function withMutexes<T>(
keyTypes: string[],
getKeyTypeMutex: (type: string) => Mutex,
fn: () => Promise<T>
): Promise<T> {
if (keyTypes.length === 0) {
return fn()
}
if (keyTypes.length === 1) {
return getKeyTypeMutex(keyTypes[0]!).runExclusive(fn)
}
// For multiple mutexes, sort by key type to prevent deadlocks
// Then acquire all mutexes in order using Promise.all for better efficiency
const sortedKeyTypes = [...keyTypes].sort()
const mutexes = sortedKeyTypes.map(getKeyTypeMutex)
// Acquire all mutexes in order to prevent deadlocks
const releases: (() => void)[] = []
try {
for (const mutex of mutexes) {
releases.push(await mutex.acquire())
}
return await fn()
} finally {
// Release in reverse order
while (releases.length > 0) {
const release = releases.pop()
if (release) release()
}
}
}
/**
* 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
* Adds DB-like transaction capability to the SignalKeyStore
* Uses AsyncLocalStorage for automatic context management
* @param state the key store to apply this capability to
* @param logger logger to log events
* @returns SignalKeyStore with transaction capability
@@ -280,261 +116,193 @@ export const addTransactionCapability = (
logger: ILogger,
{ maxCommitRetries, delayBetweenTriesMs }: TransactionCapabilityOptions
): SignalKeyStoreWithTransaction => {
// number of queries made to the DB during the transaction
// only there for logging purposes
let dbQueriesInTransaction = 0
let transactionCache: SignalDataSet = {}
let mutations: SignalDataSet = {}
// AsyncLocalStorage for transaction context
const txStorage = new AsyncLocalStorage<TransactionContext>()
// LRU Cache to hold mutexes for different key types
const mutexCache = new LRUCache<string, Mutex>({
ttl: 60 * 60 * 1000, // 1 hour
ttlAutopurge: true,
updateAgeOnGet: true
})
// Queues for concurrency control
const keyQueues = new Map<string, PQueue>()
const txMutexes = new Map<string, Mutex>()
let transactionsInProgress = 0
// Pre-key manager for specialized operations
const preKeyManager = new PreKeyManager(state, logger)
function getKeyTypeMutex(type: string): Mutex {
return getMutex(`keytype:${type}`)
}
function getSenderKeyMutex(senderKeyName: string): Mutex {
return getMutex(`senderkey:${senderKeyName}`)
}
function getTransactionMutex(key: string): Mutex {
return getMutex(`transaction:${key}`)
}
// Get or create a mutex for a specific key name
function getMutex(key: string): Mutex {
let mutex = mutexCache.get(key)
if (!mutex) {
mutex = new Mutex()
mutexCache.set(key, mutex)
logger.info({ key }, 'created new mutex')
/**
* Get or create a queue for a specific key type
*/
function getQueue(key: string): PQueue {
if (!keyQueues.has(key)) {
keyQueues.set(key, new PQueue({ concurrency: 1 }))
}
return mutex
return keyQueues.get(key)!
}
// Sender key operations with proper mutex sequencing
function queueSenderKeyOperation<T>(senderKeyName: string, operation: () => Promise<T>): Promise<T> {
return getSenderKeyMutex(senderKeyName).runExclusive(operation)
/**
* Get or create a transaction mutex
*/
function getTxMutex(key: string): Mutex {
if (!txMutexes.has(key)) {
txMutexes.set(key, new Mutex())
}
return txMutexes.get(key)!
}
// Check if we are currently in a transaction
function isInTransaction() {
return transactionsInProgress > 0
/**
* Check if currently in a transaction
*/
function isInTransaction(): boolean {
return !!txStorage.getStore()
}
// Helper function to handle transaction commit with retries
async function commitTransaction(): Promise<void> {
if (!Object.keys(mutations).length) {
/**
* Commit transaction with retries
*/
async function commitWithRetry(mutations: SignalDataSet): Promise<void> {
if (Object.keys(mutations).length === 0) {
logger.trace('no mutations in transaction')
return
}
logger.trace('committing transaction')
let tries = maxCommitRetries
while (tries > 0) {
tries -= 1
for (let attempt = 0; attempt < maxCommitRetries; attempt++) {
try {
await state.set(mutations)
logger.trace({ dbQueriesInTransaction }, 'committed transaction')
logger.trace({ mutationCount: Object.keys(mutations).length }, 'committed transaction')
return
} catch (error) {
logger.warn(`failed to commit ${Object.keys(mutations).length} mutations, tries left=${tries}`)
if (tries > 0) {
await delay(delayBetweenTriesMs)
const retriesLeft = maxCommitRetries - attempt - 1
logger.warn(`failed to commit mutations, retries left=${retriesLeft}`)
if (retriesLeft === 0) {
throw error
}
await delay(delayBetweenTriesMs)
}
}
}
// Helper function to clean up transaction state
function cleanupTransactionState(): void {
transactionsInProgress -= 1
if (transactionsInProgress === 0) {
transactionCache = {}
mutations = {}
dbQueriesInTransaction = 0
}
}
// Helper function to execute work within transaction
async function executeTransactionWork(work: () => Promise<any>): Promise<any> {
const result = await work()
// commit if this is the outermost transaction
if (transactionsInProgress === 1) {
await commitTransaction()
}
return result
}
return {
get: async (type, ids) => {
if (isInTransaction()) {
const dict = transactionCache[type]
const idsRequiringFetch = dict ? ids.filter(item => typeof dict[item] === 'undefined') : ids
// only fetch if there are any items to fetch
if (idsRequiringFetch.length) {
dbQueriesInTransaction += 1
const ctx = txStorage.getStore()
// Use per-sender-key queue for sender-key operations when possible
if (type === 'sender-key') {
logger.info({ idsRequiringFetch }, 'processing sender keys in transaction')
// For sender keys, process each one with queued operations to maintain serialization
for (const senderKeyName of idsRequiringFetch) {
await queueSenderKeyOperation(senderKeyName, async () => {
logger.info({ senderKeyName }, 'fetching sender key in transaction')
const result = await state.get(type, [senderKeyName])
// Update transaction cache
transactionCache[type] ||= {}
Object.assign(transactionCache[type]!, result)
logger.info({ senderKeyName, hasResult: !!result[senderKeyName] }, 'sender key fetch complete')
})
}
} else {
// Use runExclusive for cleaner mutex handling
await getKeyTypeMutex(type as string).runExclusive(async () => {
const result = await state.get(type, idsRequiringFetch)
if (!ctx) {
// No transaction - direct read with queue protection
return getQueue(type).add(async () => state.get(type, ids))
}
// Update transaction cache
transactionCache[type] ||= {}
Object.assign(transactionCache[type]!, result)
})
}
}
// In transaction - check cache first
const cached = ctx.cache[type] || {}
const missing = ids.filter(id => !(id in cached))
return ids.reduce((dict: { [T in string]: any }, id) => {
const value = transactionCache[type]?.[id]
if (value) {
dict[id] = value
}
if (missing.length > 0) {
ctx.dbQueries++
logger.trace({ type, count: missing.length }, 'fetching missing keys in transaction')
return dict
}, {})
} else {
// Not in transaction, fetch directly with queue protection
if (type === 'sender-key') {
// For sender keys, use individual queues to maintain per-key serialization
const results: { [key: string]: SignalDataTypeMap[typeof type] } = {}
for (const senderKeyName of ids) {
const result = await queueSenderKeyOperation(
senderKeyName,
async () => await state.get(type, [senderKeyName])
)
Object.assign(results, result)
}
const fetched = await getQueue(type).add(async () => state.get(type, missing))
return results
} else {
return await getKeyTypeMutex(type as string).runExclusive(() => state.get(type, ids))
// Update cache
ctx.cache[type] = ctx.cache[type] || ({} as any)
Object.assign(ctx.cache[type]!, fetched)
}
// Return requested ids from cache
const result: { [key: string]: any } = {}
for (const id of ids) {
const value = ctx.cache[type]?.[id]
if (value !== undefined && value !== null) {
result[id] = value
}
}
return result
},
set: async data => {
if (isInTransaction()) {
logger.trace({ types: Object.keys(data) }, 'caching in transaction')
for (const key_ in data) {
const key = key_ as keyof SignalDataTypeMap
transactionCache[key] = transactionCache[key] || ({} as any)
const ctx = txStorage.getStore()
// Special handling for pre-keys and signed-pre-keys
if (key === 'pre-key') {
await handlePreKeyOperations(data, key, transactionCache, mutations, logger, true)
} else {
// Normal handling for other key types
handleNormalKeyOperations(data, key, transactionCache, mutations)
if (!ctx) {
// No transaction - direct write with queue protection
const types = Object.keys(data)
// Process pre-keys with validation
for (const type_ of types) {
const type = type_ as keyof SignalDataTypeMap
if (type === 'pre-key') {
await preKeyManager.validateDeletions(data, type)
}
}
} else {
// Not in transaction, apply directly with mutex protection
const hasSenderKeys = 'sender-key' in data
const senderKeyNames = hasSenderKeys ? Object.keys(data['sender-key'] || {}) : []
if (hasSenderKeys) {
logger.info({ senderKeyNames }, 'processing sender key set operations')
// Handle sender key operations with per-key queues
for (const senderKeyName of senderKeyNames) {
await queueSenderKeyOperation(senderKeyName, async () => {
// Create data subset for this specific sender key
const senderKeyData = {
'sender-key': {
[senderKeyName]: data['sender-key']![senderKeyName]
}
}
logger.trace({ senderKeyName }, 'storing sender key')
// Apply changes to the store
await state.set(senderKeyData as SignalDataSet)
logger.trace({ senderKeyName }, 'sender key stored')
// Write all data in parallel
await Promise.all(
types.map(type =>
getQueue(type).add(async () => {
const typeData = { [type]: data[type as keyof SignalDataTypeMap] } as SignalDataSet
await state.set(typeData)
})
}
)
)
return
}
// Handle any non-sender-key data with regular mutexes
const nonSenderKeyData = { ...data }
delete nonSenderKeyData['sender-key']
// In transaction - update cache and mutations
logger.trace({ types: Object.keys(data) }, 'caching in transaction')
if (Object.keys(nonSenderKeyData).length > 0) {
await withMutexes(Object.keys(nonSenderKeyData), getKeyTypeMutex, async () => {
// Process pre-keys and signed-pre-keys separately with specialized mutexes
for (const key_ in nonSenderKeyData) {
const keyType = key_ as keyof SignalDataTypeMap
if (keyType === 'pre-key') {
await processPreKeyDeletions(nonSenderKeyData, keyType, state, logger)
}
}
for (const key_ in data) {
const key = key_ as keyof SignalDataTypeMap
// Apply changes to the store
await state.set(nonSenderKeyData)
})
}
// Ensure structures exist
ctx.cache[key] = ctx.cache[key] || ({} as any)
ctx.mutations[key] = ctx.mutations[key] || ({} as any)
// Special handling for pre-keys
if (key === 'pre-key') {
await preKeyManager.processOperations(data, key, ctx.cache, ctx.mutations, true)
} else {
// No sender keys - use original logic
await withMutexes(Object.keys(data), getKeyTypeMutex, async () => {
// Process pre-keys and signed-pre-keys separately with specialized mutexes
for (const key_ in data) {
const keyType = key_ as keyof SignalDataTypeMap
if (keyType === 'pre-key') {
await processPreKeyDeletions(data, keyType, state, logger)
}
}
// Apply changes to the store
await state.set(data)
})
// Normal key types
Object.assign(ctx.cache[key]!, data[key])
Object.assign(ctx.mutations[key]!, data[key])
}
}
},
isInTransaction,
async transaction(work, key) {
const releaseTxMutex = await getTransactionMutex(key).acquire()
try {
transactionsInProgress += 1
if (transactionsInProgress === 1) {
logger.trace('entering transaction')
isInTransaction,
transaction: async (work, key) => {
const existing = txStorage.getStore()
// Nested transaction - reuse existing context
if (existing) {
logger.trace('reusing existing transaction context')
return work()
}
// New transaction - acquire mutex and create context
return getTxMutex(key).runExclusive(async () => {
const ctx: TransactionContext = {
cache: {},
mutations: {},
dbQueries: 0
}
// Release the transaction mutex now that we've updated the counter
// This allows other transactions to start preparing
releaseTxMutex()
logger.trace('entering transaction')
try {
return await executeTransactionWork(work)
} finally {
cleanupTransactionState()
const result = await txStorage.run(ctx, work)
// Commit mutations
await commitWithRetry(ctx.mutations)
logger.trace({ dbQueries: ctx.dbQueries }, 'transaction completed')
return result
} catch (error) {
logger.error({ error }, 'transaction failed, rolling back')
throw error
}
} catch (error) {
releaseTxMutex()
throw error
}
})
}
}
}
src/Utils/pre-key-manager.ts
+126
View File
@@ -0,0 +1,126 @@
import PQueue from 'p-queue'
import type { SignalDataSet, SignalDataTypeMap, SignalKeyStore } from '../Types'
import type { ILogger } from './logger'
/**
* Manages pre-key operations with proper concurrency control
*/
export class PreKeyManager {
private readonly queues = new Map<string, PQueue>()
constructor(
private readonly store: SignalKeyStore,
private readonly logger: ILogger
) {}
/**
* Get or create a queue for a specific key type
*/
private getQueue(keyType: string): PQueue {
if (!this.queues.has(keyType)) {
this.queues.set(keyType, new PQueue({ concurrency: 1 }))
}
return this.queues.get(keyType)!
}
/**
* Process pre-key operations (updates and deletions)
*/
async processOperations(
data: SignalDataSet,
keyType: keyof SignalDataTypeMap,
transactionCache: SignalDataSet,
mutations: SignalDataSet,
isInTransaction: boolean
): Promise<void> {
const keyData = data[keyType]
if (!keyData) return
return this.getQueue(keyType).add(async () => {
// Ensure structures exist
transactionCache[keyType] = transactionCache[keyType] || ({} as any)
mutations[keyType] = mutations[keyType] || ({} as any)
// Separate deletions from updates
const deletions: string[] = []
const updates: Record<string, any> = {}
for (const keyId in keyData) {
if (keyData[keyId] === null) {
deletions.push(keyId)
} else {
updates[keyId] = keyData[keyId]
}
}
// Process updates (no validation needed)
if (Object.keys(updates).length > 0) {
Object.assign(transactionCache[keyType]!, updates)
Object.assign(mutations[keyType]!, updates)
}
// Process deletions with validation
if (deletions.length > 0) {
await this.processDeletions(keyType, deletions, transactionCache, mutations, isInTransaction)
}
})
}
/**
* Process deletions with validation
*/
private async processDeletions(
keyType: keyof SignalDataTypeMap,
ids: string[],
transactionCache: SignalDataSet,
mutations: SignalDataSet,
isInTransaction: boolean
): Promise<void> {
if (isInTransaction) {
// In transaction, only allow deletion if key exists in cache
for (const keyId of ids) {
if (transactionCache[keyType]?.[keyId]) {
transactionCache[keyType][keyId] = null
mutations[keyType]![keyId] = null
} else {
this.logger.warn(`Skipping deletion of non-existent ${keyType} in transaction: ${keyId}`)
}
}
} else {
// Outside transaction, validate against store
const existingKeys = await this.store.get(keyType, ids)
for (const keyId of ids) {
if (existingKeys[keyId]) {
transactionCache[keyType]![keyId] = null
mutations[keyType]![keyId] = null
} else {
this.logger.warn(`Skipping deletion of non-existent ${keyType}: ${keyId}`)
}
}
}
}
/**
* Validate and process pre-key deletions outside transactions
*/
async validateDeletions(data: SignalDataSet, keyType: keyof SignalDataTypeMap): Promise<void> {
const keyData = data[keyType]
if (!keyData) return
return this.getQueue(keyType).add(async () => {
// Find all deletion requests
const deletionIds = Object.keys(keyData).filter(id => keyData[id] === null)
if (deletionIds.length === 0) return
// Validate deletions
const existingKeys = await this.store.get(keyType, deletionIds)
for (const keyId of deletionIds) {
if (!existingKeys[keyId]) {
this.logger.warn(`Skipping deletion of non-existent ${keyType}: ${keyId}`)
delete data[keyType]![keyId]
}
}
})
}
}
yarn.lock
+25
View File
@@ -3004,6 +3004,7 @@ __metadata:
lru-cache: "npm:^11.1.0"
music-metadata: "npm:^11.7.0"
open: "npm:^8.4.2"
p-queue: "npm:^9.0.0"
pino: "npm:^9.6"
pino-pretty: "npm:^13.1.1"
prettier: "npm:^3.5.3"
@@ -4785,6 +4786,13 @@ __metadata:
languageName: node
linkType: hard
"eventemitter3@npm:^5.0.1":
version: 5.0.1
resolution: "eventemitter3@npm:5.0.1"
checksum: 10c0/4ba5c00c506e6c786b4d6262cfbce90ddc14c10d4667e5c83ae993c9de88aa856033994dd2b35b83e8dc1170e224e66a319fa80adc4c32adcd2379bbc75da814
languageName: node
linkType: hard
"events@npm:^3.3.0":
version: 3.3.0
resolution: "events@npm:3.3.0"
@@ -8343,6 +8351,23 @@ __metadata:
languageName: node
linkType: hard
"p-queue@npm:^9.0.0":
version: 9.0.0
resolution: "p-queue@npm:9.0.0"
dependencies:
eventemitter3: "npm:^5.0.1"
p-timeout: "npm:^7.0.0"
checksum: 10c0/0f27fcbec9e4e02f34cc3660f14b5746798c51a2425dc3d3c5238924c34726b1580606d5d1432fc05304e2350c00b112ded03eb43c1b49de7791a7150fbfcb4e
languageName: node
linkType: hard
"p-timeout@npm:^7.0.0":
version: 7.0.0
resolution: "p-timeout@npm:7.0.0"
checksum: 10c0/0418be5dd0269916293cfd10f738427a1a33b50d2358cf72fecee2b8a31cc2a663ae178105f1d6ef92301c50f86f78074adb32f8321856a4f332765786fedd93
languageName: node
linkType: hard
"p-try@npm:^1.0.0":
version: 1.0.0
resolution: "p-try@npm:1.0.0"