Skip to content

Commit

Permalink
Improve performance of decryption and keysharing
Browse files Browse the repository at this point in the history 
create targeted loop that does all work on high priority streams first (previously we were prioritizing decryption, but key solicitation was just whatever came first

optimize initialization and adding key solicitations. this is a really hot loop for large towns and these small changes help us out with cpu utilization
  • Loading branch information
@texuf
texuf committed Oct 21, 2024
1 parent b32828b commit c6cfb23
Show file tree
Hide file tree
Showing 4 changed files with 143 additions and 110 deletions.
214 changes: 115 additions & 99 deletions packages/encryption/src/decryptionExtensions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,19 +41,18 @@ export type DecryptionEvents = {
decryptionExtStatusChanged: (status: DecryptionStatus) => void
}

export interface NewGroupSessionItem {
streamId: string
sessions: UserInboxPayload_GroupEncryptionSessions
}

export interface EncryptedContentItem {
streamId: string
eventId: string
kind: string // kind of encrypted data
encryptedData: EncryptedData
}

interface DecryptionRetryItem {
streamId: string
event: EncryptedContentItem
retryAt: Date
}

export interface KeySolicitationContent {
deviceKey: string
fallbackKey: string
Expand Down Expand Up @@ -124,20 +123,20 @@ export abstract class BaseDecryptionExtensions {
private _status: DecryptionStatus = DecryptionStatus.initializing
private queues = {
priorityTasks: new Array<() => Promise<void>>(),
newGroupSession: new Array<UserInboxPayload_GroupEncryptionSessions>(),
newGroupSession: new Array<NewGroupSessionItem>(),
encryptedContent: new Array<EncryptedContentItem>(),
decryptionRetries: new Array<DecryptionRetryItem>(),
missingKeys: new Array<MissingKeysItem>(),
keySolicitations: new Array<KeySolicitationItem>(),
}
private upToDateStreams = new Set<string>()
private highPriorityStreams = new Set<string>()
private highPriorityStreams: string[] = []
private decryptionFailures: Record<string, Record<string, EncryptedContentItem[]>> = {} // streamId: sessionId: EncryptedContentItem[]
private inProgressTick?: Promise<void>
private timeoutId?: NodeJS.Timeout
private delayMs: number = 15
private started: boolean = false
private emitter: TypedEmitter<DecryptionEvents>
private keySolicitationsNeedsSort = false

protected _onStopFn?: () => void
protected log: {
Expand Down Expand Up @@ -218,7 +217,8 @@ export abstract class BaseDecryptionExtensions {
_senderId: string,
): void {
this.log.info('enqueueNewGroupSessions', sessions)
this.queues.newGroupSession.push(sessions)
const streamId = bin_toHexString(sessions.streamId)
this.queues.newGroupSession.push({ streamId, sessions })
this.checkStartTicking()
}

Expand All @@ -237,6 +237,38 @@ export abstract class BaseDecryptionExtensions {
this.checkStartTicking()
}

public enqueueInitKeySolicitations(
streamId: string,
members: {
userId: string
userAddress: Uint8Array
solicitations: KeySolicitationContent[]
}[],
): void {
this.queues.keySolicitations = this.queues.keySolicitations.filter(
(x) => x.streamId !== streamId,
)
for (const member of members) {
const { userId: fromUserId, userAddress: fromUserAddress } = member
for (const keySolicitation of member.solicitations) {
if (keySolicitation.deviceKey === this.userDevice.deviceKey) {
continue
}
this.queues.keySolicitations.push({
streamId,
fromUserId,
fromUserAddress,
solicitation: keySolicitation,
respondAfter: new Date(
Date.now() + this.getRespondDelayMSForKeySolicitation(streamId, fromUserId),
),
} satisfies KeySolicitationItem)
}
}
this.keySolicitationsNeedsSort = true
this.checkStartTicking()
}

public enqueueKeySolicitation(
streamId: string,
fromUserId: string,
Expand All @@ -256,19 +288,16 @@ export abstract class BaseDecryptionExtensions {
}
if (keySolicitation.sessionIds.length > 0 || keySolicitation.isNewDevice) {
this.log.debug('new key solicitation', keySolicitation)
insertSorted(
this.queues.keySolicitations,
{
streamId,
fromUserId,
fromUserAddress,
solicitation: keySolicitation,
respondAfter: new Date(
Date.now() + this.getRespondDelayMSForKeySolicitation(streamId, fromUserId),
),
} satisfies KeySolicitationItem,
(x) => x.respondAfter,
)
this.keySolicitationsNeedsSort = true
this.queues.keySolicitations.push({
streamId,
fromUserId,
fromUserAddress,
solicitation: keySolicitation,
respondAfter: new Date(
Date.now() + this.getRespondDelayMSForKeySolicitation(streamId, fromUserId),
),
} satisfies KeySolicitationItem)
this.checkStartTicking()
} else if (index > -1) {
this.log.debug('cleared key solicitation', keySolicitation)
Expand Down Expand Up @@ -408,43 +437,62 @@ export abstract class BaseDecryptionExtensions {
return priorityTask()
}

// update any new group sessions
const session = this.queues.newGroupSession.shift()
if (session) {
this.setStatus(DecryptionStatus.processingNewGroupSessions)
return this.processNewGroupSession(session)
}
for (const streamId of [...this.highPriorityStreams, undefined]) {
//
if (streamId && !this.upToDateStreams.has(streamId)) {
continue
}
//console.log('csb:dec streamId', streamId)

const encryptedContent = dequeueHighPriority(
this.queues.encryptedContent,
this.highPriorityStreams,
)
if (encryptedContent) {
this.setStatus(DecryptionStatus.decryptingEvents)
return this.processEncryptedContentItem(encryptedContent)
}
if (!streamId) {
// respond to key solicitations from yourself
const ownKeySolicitationIndex = this.queues.keySolicitations.findIndex(
(x) => x.fromUserId === this.userId,
)
if (ownKeySolicitationIndex > -1) {
const solicitation = this.queues.keySolicitations.splice(
ownKeySolicitationIndex,
1,
)[0]
if (solicitation) {
this.log.debug(' processing own key solicitation')
this.setStatus(DecryptionStatus.respondingToKeyRequests)
return this.processKeySolicitation(solicitation)
}
}
}

const decryptionRetry = dequeueUpToDate(
this.queues.decryptionRetries,
now,
(x) => x.retryAt,
this.upToDateStreams,
)
if (decryptionRetry) {
this.setStatus(DecryptionStatus.retryingDecryption)
return this.processDecryptionRetry(decryptionRetry)
}
const encryptedContent = streamId
? dequeueItemWithStreamId(this.queues.encryptedContent, streamId)
: this.queues.encryptedContent.shift()
if (encryptedContent) {
this.setStatus(DecryptionStatus.decryptingEvents)
return this.processEncryptedContentItem(encryptedContent)
}

const missingKeys = dequeueUpToDate(
this.queues.missingKeys,
now,
(x) => x.waitUntil,
this.upToDateStreams,
)
if (missingKeys) {
this.setStatus(DecryptionStatus.requestingKeys)
return this.processMissingKeys(missingKeys)
const missingKey = streamId
? dequeueItemWithStreamId(this.queues.missingKeys, streamId)
: dequeueUpToDate(
this.queues.missingKeys,
now,
(x) => x.waitUntil,
this.upToDateStreams,
)
if (missingKey) {
this.setStatus(DecryptionStatus.requestingKeys)
return this.processMissingKeys(missingKey)
}
}

if (this.keySolicitationsNeedsSort) {
this.sortKeySolicitations()
}
const keySolicitation = dequeueUpToDate(
this.queues.keySolicitations,
now,
Expand All @@ -465,17 +513,15 @@ export abstract class BaseDecryptionExtensions {
* process new group sessions that were sent to our to device stream inbox
* re-enqueue any decryption failures with matching session id
*/
private async processNewGroupSession(
session: UserInboxPayload_GroupEncryptionSessions,
): Promise<void> {
this.log.debug('processNewGroupSession', session)
private async processNewGroupSession(sessionItem: NewGroupSessionItem): Promise<void> {
const { streamId, sessions: session } = sessionItem
// check if this message is to our device
const ciphertext = session.ciphertexts[this.userDevice.deviceKey]
if (!ciphertext) {
this.log.debug('skipping, no session for our device')
return
}
const streamId = bin_toHexString(session.streamId)
this.log.debug('processNewGroupSession', session)
// check if it contains any keys we need
const neededKeyIndexs = []
for (let i = 0; i < session.sessionIds.length; i++) {
Expand Down Expand Up @@ -535,51 +581,12 @@ export abstract class BaseDecryptionExtensions {
*/
private async processEncryptedContentItem(item: EncryptedContentItem): Promise<void> {
this.log.debug('processEncryptedContentItem', item)
try {
// do the work to decrypt the event
this.log.debug('decrypting content')
await this.decryptGroupEvent(item.streamId, item.eventId, item.kind, item.encryptedData)
} catch (err: unknown) {
const sessionNotFound = isSessionNotFoundError(err)
this.log.debug('failed to decrypt', err, 'sessionNotFound', sessionNotFound)

// If !sessionNotFound, we want to know more about this error.
if (!sessionNotFound) {
this.log.info('failed to decrypt', err, 'streamId', item.streamId)
}

this.onDecryptionError(item, {
missingSession: sessionNotFound,
kind: item.kind,
encryptedData: item.encryptedData,
error: err,
})

insertSorted(
this.queues.decryptionRetries,
{
streamId: item.streamId,
event: item,
retryAt: new Date(Date.now() + 3000), // give it 3 seconds, maybe someone will send us the key
},
(x) => x.retryAt,
)
}
}

/**
* processDecryptionRetry
* retry decryption a second time for a failed decryption, keys may have arrived
*/
private async processDecryptionRetry(retryItem: DecryptionRetryItem): Promise<void> {
const item = retryItem.event
try {
this.log.debug('retrying decryption', item)
await this.decryptGroupEvent(item.streamId, item.eventId, item.kind, item.encryptedData)
} catch (err) {
const sessionNotFound = isSessionNotFoundError(err)

this.log.info('failed to decrypt on retry', err, 'sessionNotFound', sessionNotFound)
this.onDecryptionError(item, {
missingSession: sessionNotFound,
kind: item.kind,
Expand All @@ -603,6 +610,8 @@ export abstract class BaseDecryptionExtensions {
{ streamId, waitUntil: new Date(Date.now() + 1000) },
(x) => x.waitUntil,
)
} else {
this.log.info('failed to decrypt', err, 'streamId', item.streamId)
}
}
}
Expand Down Expand Up @@ -758,7 +767,14 @@ export abstract class BaseDecryptionExtensions {
}

public setHighPriorityStreams(streamIds: string[]) {
this.highPriorityStreams = new Set(streamIds)
this.highPriorityStreams = streamIds
}

private sortKeySolicitations() {
this.queues.keySolicitations.sort(
(a, b) => a.respondAfter.getTime() - b.respondAfter.getTime(),
)
this.keySolicitationsNeedsSort = false
}
}

Expand Down Expand Up @@ -808,13 +824,13 @@ function dequeueUpToDate<T extends { streamId: string }>(
return items.splice(index, 1)[0]
}

function dequeueHighPriority<T extends { streamId: string }>(
function dequeueItemWithStreamId<T extends { streamId: string }>(
items: T[],
highPriorityIds: Set<string>,
streamId: string,
): T | undefined {
const index = items.findIndex((x) => highPriorityIds.has(x.streamId))
const index = items.findIndex((x) => x.streamId === streamId)
if (index === -1) {
return items.shift()
return undefined
}
return items.splice(index, 1)[0]
}
Expand Down
11 changes: 11 additions & 0 deletions packages/sdk/src/clientDecryptionExtensions.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,11 +70,21 @@ export class ClientDecryptionExtensions extends BaseDecryptionExtensions {
keySolicitation: KeySolicitationContent,
) => this.enqueueKeySolicitation(streamId, fromUserId, fromUserAddress, keySolicitation)

const onInitKeySolicitations = (
streamId: string,
members: {
userId: string
userAddress: Uint8Array
solicitations: KeySolicitationContent[]
}[],
) => this.enqueueInitKeySolicitations(streamId, members)

client.on('streamUpToDate', onStreamUpToDate)
client.on('newGroupSessions', onNewGroupSessions)
client.on('newEncryptedContent', onNewEncryptedContent)
client.on('newKeySolicitation', onKeySolicitation)
client.on('updatedKeySolicitation', onKeySolicitation)
client.on('initKeySolicitations', onInitKeySolicitations)
client.on('streamNewUserJoined', onMembershipChange)

this._onStopFn = () => {
Expand All @@ -83,6 +93,7 @@ export class ClientDecryptionExtensions extends BaseDecryptionExtensions {
client.off('newEncryptedContent', onNewEncryptedContent)
client.off('newKeySolicitation', onKeySolicitation)
client.off('updatedKeySolicitation', onKeySolicitation)
client.off('initKeySolicitations', onInitKeySolicitations)
client.off('streamNewUserJoined', onMembershipChange)
}
this.log.debug('new ClientDecryptionExtensions', { userDevice })
Expand Down
8 changes: 8 additions & 0 deletions packages/sdk/src/streamEvents.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,14 @@ export type StreamEncryptionEvents = {
fromUserAddress: Uint8Array,
event: KeySolicitationContent,
) => void
initKeySolicitations: (
streamId: string,
members: {
userId: string
userAddress: Uint8Array
solicitations: KeySolicitationContent[]
}[],
) => void
userDeviceKeyMessage: (streamId: string, userId: string, userDevice: UserDevice) => void
}

Expand Down
Loading

0 comments on commit c6cfb23

Please sign in to comment.