auth-utils: Make a better addTransactionCapability function
This commit is contained in:
@@ -46,6 +46,7 @@
|
|||||||
"libsignal": "git+https://github.com/whiskeysockets/libsignal-node",
|
"libsignal": "git+https://github.com/whiskeysockets/libsignal-node",
|
||||||
"lru-cache": "^11.1.0",
|
"lru-cache": "^11.1.0",
|
||||||
"music-metadata": "^11.7.0",
|
"music-metadata": "^11.7.0",
|
||||||
|
"p-queue": "^9.0.0",
|
||||||
"pino": "^9.6",
|
"pino": "^9.6",
|
||||||
"protobufjs": "^7.2.4",
|
"protobufjs": "^7.2.4",
|
||||||
"ws": "^8.13.0"
|
"ws": "^8.13.0"
|
||||||
|
|||||||
+147
-379
@@ -1,7 +1,8 @@
|
|||||||
import NodeCache from '@cacheable/node-cache'
|
import NodeCache from '@cacheable/node-cache'
|
||||||
|
import { AsyncLocalStorage } from 'async_hooks'
|
||||||
import { Mutex } from 'async-mutex'
|
import { Mutex } from 'async-mutex'
|
||||||
import { randomBytes } from 'crypto'
|
import { randomBytes } from 'crypto'
|
||||||
import { LRUCache } from 'lru-cache'
|
import PQueue from 'p-queue'
|
||||||
import { DEFAULT_CACHE_TTLS } from '../Defaults'
|
import { DEFAULT_CACHE_TTLS } from '../Defaults'
|
||||||
import type {
|
import type {
|
||||||
AuthenticationCreds,
|
AuthenticationCreds,
|
||||||
@@ -15,6 +16,16 @@ import type {
|
|||||||
import { Curve, signedKeyPair } from './crypto'
|
import { Curve, signedKeyPair } from './crypto'
|
||||||
import { delay, generateRegistrationId } from './generics'
|
import { delay, generateRegistrationId } from './generics'
|
||||||
import type { ILogger } from './logger'
|
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
|
* Adds caching capability to a SignalKeyStore
|
||||||
@@ -47,6 +58,7 @@ export function makeCacheableSignalKeyStore(
|
|||||||
return cacheMutex.runExclusive(async () => {
|
return cacheMutex.runExclusive(async () => {
|
||||||
const data: { [_: string]: SignalDataTypeMap[typeof type] } = {}
|
const data: { [_: string]: SignalDataTypeMap[typeof type] } = {}
|
||||||
const idsToFetch: string[] = []
|
const idsToFetch: string[] = []
|
||||||
|
|
||||||
for (const id of ids) {
|
for (const id of ids) {
|
||||||
const item = (await cache.get<SignalDataTypeMap[typeof type]>(getUniqueId(type, id))) as any
|
const item = (await cache.get<SignalDataTypeMap[typeof type]>(getUniqueId(type, id))) as any
|
||||||
if (typeof item !== 'undefined') {
|
if (typeof item !== 'undefined') {
|
||||||
@@ -82,7 +94,6 @@ export function makeCacheableSignalKeyStore(
|
|||||||
}
|
}
|
||||||
|
|
||||||
logger?.trace({ keys }, 'updated cache')
|
logger?.trace({ keys }, 'updated cache')
|
||||||
|
|
||||||
await store.set(data)
|
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
|
* Adds DB-like transaction capability to the SignalKeyStore
|
||||||
*/
|
* Uses AsyncLocalStorage for automatic context management
|
||||||
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
|
|
||||||
* @param state the key store to apply this capability to
|
* @param state the key store to apply this capability to
|
||||||
* @param logger logger to log events
|
* @param logger logger to log events
|
||||||
* @returns SignalKeyStore with transaction capability
|
* @returns SignalKeyStore with transaction capability
|
||||||
@@ -280,261 +116,193 @@ export const addTransactionCapability = (
|
|||||||
logger: ILogger,
|
logger: ILogger,
|
||||||
{ maxCommitRetries, delayBetweenTriesMs }: TransactionCapabilityOptions
|
{ maxCommitRetries, delayBetweenTriesMs }: TransactionCapabilityOptions
|
||||||
): SignalKeyStoreWithTransaction => {
|
): SignalKeyStoreWithTransaction => {
|
||||||
// number of queries made to the DB during the transaction
|
// AsyncLocalStorage for transaction context
|
||||||
// only there for logging purposes
|
const txStorage = new AsyncLocalStorage<TransactionContext>()
|
||||||
let dbQueriesInTransaction = 0
|
|
||||||
let transactionCache: SignalDataSet = {}
|
|
||||||
let mutations: SignalDataSet = {}
|
|
||||||
|
|
||||||
// LRU Cache to hold mutexes for different key types
|
// Queues for concurrency control
|
||||||
const mutexCache = new LRUCache<string, Mutex>({
|
const keyQueues = new Map<string, PQueue>()
|
||||||
ttl: 60 * 60 * 1000, // 1 hour
|
const txMutexes = new Map<string, Mutex>()
|
||||||
ttlAutopurge: true,
|
|
||||||
updateAgeOnGet: true
|
|
||||||
})
|
|
||||||
|
|
||||||
let transactionsInProgress = 0
|
// Pre-key manager for specialized operations
|
||||||
|
const preKeyManager = new PreKeyManager(state, logger)
|
||||||
|
|
||||||
function getKeyTypeMutex(type: string): Mutex {
|
/**
|
||||||
return getMutex(`keytype:${type}`)
|
* 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 }))
|
||||||
}
|
}
|
||||||
|
|
||||||
function getSenderKeyMutex(senderKeyName: string): Mutex {
|
return keyQueues.get(key)!
|
||||||
return getMutex(`senderkey:${senderKeyName}`)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
function getTransactionMutex(key: string): Mutex {
|
/**
|
||||||
return getMutex(`transaction:${key}`)
|
* Get or create a transaction mutex
|
||||||
|
*/
|
||||||
|
function getTxMutex(key: string): Mutex {
|
||||||
|
if (!txMutexes.has(key)) {
|
||||||
|
txMutexes.set(key, new Mutex())
|
||||||
}
|
}
|
||||||
|
|
||||||
// Get or create a mutex for a specific key name
|
return txMutexes.get(key)!
|
||||||
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')
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return mutex
|
/**
|
||||||
|
* Check if currently in a transaction
|
||||||
|
*/
|
||||||
|
function isInTransaction(): boolean {
|
||||||
|
return !!txStorage.getStore()
|
||||||
}
|
}
|
||||||
|
|
||||||
// Sender key operations with proper mutex sequencing
|
/**
|
||||||
function queueSenderKeyOperation<T>(senderKeyName: string, operation: () => Promise<T>): Promise<T> {
|
* Commit transaction with retries
|
||||||
return getSenderKeyMutex(senderKeyName).runExclusive(operation)
|
*/
|
||||||
}
|
async function commitWithRetry(mutations: SignalDataSet): Promise<void> {
|
||||||
|
if (Object.keys(mutations).length === 0) {
|
||||||
// Check if we are currently in a transaction
|
|
||||||
function isInTransaction() {
|
|
||||||
return transactionsInProgress > 0
|
|
||||||
}
|
|
||||||
|
|
||||||
// Helper function to handle transaction commit with retries
|
|
||||||
async function commitTransaction(): Promise<void> {
|
|
||||||
if (!Object.keys(mutations).length) {
|
|
||||||
logger.trace('no mutations in transaction')
|
logger.trace('no mutations in transaction')
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
logger.trace('committing transaction')
|
logger.trace('committing transaction')
|
||||||
let tries = maxCommitRetries
|
|
||||||
|
|
||||||
while (tries > 0) {
|
for (let attempt = 0; attempt < maxCommitRetries; attempt++) {
|
||||||
tries -= 1
|
|
||||||
try {
|
try {
|
||||||
await state.set(mutations)
|
await state.set(mutations)
|
||||||
logger.trace({ dbQueriesInTransaction }, 'committed transaction')
|
logger.trace({ mutationCount: Object.keys(mutations).length }, 'committed transaction')
|
||||||
return
|
return
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
logger.warn(`failed to commit ${Object.keys(mutations).length} mutations, tries left=${tries}`)
|
const retriesLeft = maxCommitRetries - attempt - 1
|
||||||
if (tries > 0) {
|
logger.warn(`failed to commit mutations, retries left=${retriesLeft}`)
|
||||||
|
|
||||||
|
if (retriesLeft === 0) {
|
||||||
|
throw error
|
||||||
|
}
|
||||||
|
|
||||||
await delay(delayBetweenTriesMs)
|
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 {
|
return {
|
||||||
get: async (type, ids) => {
|
get: async (type, ids) => {
|
||||||
if (isInTransaction()) {
|
const ctx = txStorage.getStore()
|
||||||
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
|
|
||||||
|
|
||||||
// Use per-sender-key queue for sender-key operations when possible
|
if (!ctx) {
|
||||||
if (type === 'sender-key') {
|
// No transaction - direct read with queue protection
|
||||||
logger.info({ idsRequiringFetch }, 'processing sender keys in transaction')
|
return getQueue(type).add(async () => state.get(type, ids))
|
||||||
// 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)
|
|
||||||
|
|
||||||
// Update transaction cache
|
// In transaction - check cache first
|
||||||
transactionCache[type] ||= {}
|
const cached = ctx.cache[type] || {}
|
||||||
Object.assign(transactionCache[type]!, result)
|
const missing = ids.filter(id => !(id in cached))
|
||||||
})
|
|
||||||
|
if (missing.length > 0) {
|
||||||
|
ctx.dbQueries++
|
||||||
|
logger.trace({ type, count: missing.length }, 'fetching missing keys in transaction')
|
||||||
|
|
||||||
|
const fetched = await getQueue(type).add(async () => state.get(type, missing))
|
||||||
|
|
||||||
|
// 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 ids.reduce((dict: { [T in string]: any }, id) => {
|
return result
|
||||||
const value = transactionCache[type]?.[id]
|
|
||||||
if (value) {
|
|
||||||
dict[id] = value
|
|
||||||
}
|
|
||||||
|
|
||||||
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)
|
|
||||||
}
|
|
||||||
|
|
||||||
return results
|
|
||||||
} else {
|
|
||||||
return await getKeyTypeMutex(type as string).runExclusive(() => state.get(type, ids))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
},
|
},
|
||||||
|
|
||||||
set: async data => {
|
set: async data => {
|
||||||
if (isInTransaction()) {
|
const ctx = txStorage.getStore()
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
|
|
||||||
|
// In transaction - update cache and mutations
|
||||||
logger.trace({ types: Object.keys(data) }, 'caching in transaction')
|
logger.trace({ types: Object.keys(data) }, 'caching in transaction')
|
||||||
|
|
||||||
for (const key_ in data) {
|
for (const key_ in data) {
|
||||||
const key = key_ as keyof SignalDataTypeMap
|
const key = key_ as keyof SignalDataTypeMap
|
||||||
transactionCache[key] = transactionCache[key] || ({} as any)
|
|
||||||
|
|
||||||
// Special handling for pre-keys and signed-pre-keys
|
// 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') {
|
if (key === 'pre-key') {
|
||||||
await handlePreKeyOperations(data, key, transactionCache, mutations, logger, true)
|
await preKeyManager.processOperations(data, key, ctx.cache, ctx.mutations, true)
|
||||||
} else {
|
} else {
|
||||||
// Normal handling for other key types
|
// Normal key types
|
||||||
handleNormalKeyOperations(data, key, transactionCache, mutations)
|
Object.assign(ctx.cache[key]!, data[key])
|
||||||
}
|
Object.assign(ctx.mutations[key]!, data[key])
|
||||||
}
|
|
||||||
} 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')
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
// Handle any non-sender-key data with regular mutexes
|
|
||||||
const nonSenderKeyData = { ...data }
|
|
||||||
delete nonSenderKeyData['sender-key']
|
|
||||||
|
|
||||||
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)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Apply changes to the store
|
|
||||||
await state.set(nonSenderKeyData)
|
|
||||||
})
|
|
||||||
}
|
|
||||||
} 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)
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
|
||||||
isInTransaction,
|
isInTransaction,
|
||||||
async transaction(work, key) {
|
|
||||||
const releaseTxMutex = await getTransactionMutex(key).acquire()
|
|
||||||
|
|
||||||
try {
|
transaction: async (work, key) => {
|
||||||
transactionsInProgress += 1
|
const existing = txStorage.getStore()
|
||||||
if (transactionsInProgress === 1) {
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
|
|
||||||
logger.trace('entering transaction')
|
logger.trace('entering transaction')
|
||||||
}
|
|
||||||
|
|
||||||
// Release the transaction mutex now that we've updated the counter
|
|
||||||
// This allows other transactions to start preparing
|
|
||||||
releaseTxMutex()
|
|
||||||
|
|
||||||
try {
|
try {
|
||||||
return await executeTransactionWork(work)
|
const result = await txStorage.run(ctx, work)
|
||||||
} finally {
|
|
||||||
cleanupTransactionState()
|
// Commit mutations
|
||||||
}
|
await commitWithRetry(ctx.mutations)
|
||||||
|
|
||||||
|
logger.trace({ dbQueries: ctx.dbQueries }, 'transaction completed')
|
||||||
|
|
||||||
|
return result
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
releaseTxMutex()
|
logger.error({ error }, 'transaction failed, rolling back')
|
||||||
throw error
|
throw error
|
||||||
}
|
}
|
||||||
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -3004,6 +3004,7 @@ __metadata:
|
|||||||
lru-cache: "npm:^11.1.0"
|
lru-cache: "npm:^11.1.0"
|
||||||
music-metadata: "npm:^11.7.0"
|
music-metadata: "npm:^11.7.0"
|
||||||
open: "npm:^8.4.2"
|
open: "npm:^8.4.2"
|
||||||
|
p-queue: "npm:^9.0.0"
|
||||||
pino: "npm:^9.6"
|
pino: "npm:^9.6"
|
||||||
pino-pretty: "npm:^13.1.1"
|
pino-pretty: "npm:^13.1.1"
|
||||||
prettier: "npm:^3.5.3"
|
prettier: "npm:^3.5.3"
|
||||||
@@ -4785,6 +4786,13 @@ __metadata:
|
|||||||
languageName: node
|
languageName: node
|
||||||
linkType: hard
|
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":
|
"events@npm:^3.3.0":
|
||||||
version: 3.3.0
|
version: 3.3.0
|
||||||
resolution: "events@npm:3.3.0"
|
resolution: "events@npm:3.3.0"
|
||||||
@@ -8343,6 +8351,23 @@ __metadata:
|
|||||||
languageName: node
|
languageName: node
|
||||||
linkType: hard
|
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":
|
"p-try@npm:^1.0.0":
|
||||||
version: 1.0.0
|
version: 1.0.0
|
||||||
resolution: "p-try@npm:1.0.0"
|
resolution: "p-try@npm:1.0.0"
|
||||||
|
|||||||
Reference in New Issue
Block a user