diff --git a/package.json b/package.json index 53328b08..d80480b3 100644 --- a/package.json +++ b/package.json @@ -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" diff --git a/src/Utils/auth-utils.ts b/src/Utils/auth-utils.ts index e6eaf279..f3d8c87b 100644 --- a/src/Utils/auth-utils.ts +++ b/src/Utils/auth-utils.ts @@ -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(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 { - 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 { - 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( - keyTypes: string[], - getKeyTypeMutex: (type: string) => Mutex, - fn: () => Promise -): Promise { - 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() - // LRU Cache to hold mutexes for different key types - const mutexCache = new LRUCache({ - ttl: 60 * 60 * 1000, // 1 hour - ttlAutopurge: true, - updateAgeOnGet: true - }) + // Queues for concurrency control + const keyQueues = new Map() + const txMutexes = new Map() - 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(senderKeyName: string, operation: () => Promise): Promise { - 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 { - if (!Object.keys(mutations).length) { + /** + * Commit transaction with retries + */ + async function commitWithRetry(mutations: SignalDataSet): Promise { + 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): Promise { - 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 - } + }) } } } diff --git a/src/Utils/pre-key-manager.ts b/src/Utils/pre-key-manager.ts new file mode 100644 index 00000000..635a8195 --- /dev/null +++ b/src/Utils/pre-key-manager.ts @@ -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() + + 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 { + 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 = {} + + 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 { + 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 { + 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] + } + } + }) + } +} diff --git a/yarn.lock b/yarn.lock index 3b474557..c32595ba 100644 --- a/yarn.lock +++ b/yarn.lock @@ -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"