Skip to content

Commit

Permalink
ElementR | backup: call expensive roomKeyCounts less often (#4015)
Browse files Browse the repository at this point in the history 
* ElementR | backup: call expensive `roomKeyCounts` less often

* review: Improve doc

Co-authored-by: Richard van der Hoff <[email protected]>

* review: Improve loop

* review: Add comment regarding slightly outdated remaining count

* Review: doc fix typo

Co-authored-by: Richard van der Hoff <[email protected]>

* review: refactor code order, count after doing the request

* review: Missing await on sleep for limit exceeded

* review: Comment | add a note for when performance drops

* Backup: add upload loop test for rust

* test: quick fix backup loop tests

* test: quick fix imports backup loop tests

* review: improve comment

Co-authored-by: Richard van der Hoff <[email protected]>

* Review improve comment

Co-authored-by: Richard van der Hoff <[email protected]>

* Review: Clean and improve tests

* fix: wrong test name

---------

Co-authored-by: Richard van der Hoff <[email protected]>
  • Loading branch information
@BillCarsonFr @richvdh
BillCarsonFr and richvdh authored Jan 22, 2024
1 parent b10a804 commit c993785
Show file tree
Hide file tree
Showing 2 changed files with 202 additions and 8 deletions.
144 changes: 144 additions & 0 deletions spec/unit/rust-crypto/backup.spec.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
import { Mocked } from "jest-mock";
import fetchMock from "fetch-mock-jest";
import * as RustSdkCryptoJs from "@matrix-org/matrix-sdk-crypto-wasm";

import { CryptoEvent, HttpApiEvent, HttpApiEventHandlerMap, MatrixHttpApi, TypedEventEmitter } from "../../../src";
import { OutgoingRequestProcessor } from "../../../src/rust-crypto/OutgoingRequestProcessor";
import * as testData from "../../test-utils/test-data";
import * as TestData from "../../test-utils/test-data";
import { IKeyBackup } from "../../../src/crypto/backup";
import { IKeyBackupSession } from "../../../src/crypto/keybackup";
import { RustBackupManager } from "../../../src/rust-crypto/backup";

describe("Upload keys to backup", () => {
/** The backup manager under test */
let rustBackupManager: RustBackupManager;

let mockOlmMachine: Mocked<RustSdkCryptoJs.OlmMachine>;

let outgoingRequestProcessor: Mocked<OutgoingRequestProcessor>;

const httpAPi = new MatrixHttpApi(new TypedEventEmitter<HttpApiEvent, HttpApiEventHandlerMap>(), {
baseUrl: "http://server/",
prefix: "",
onlyData: true,
});

let idGenerator = 0;
function mockBackupRequest(keyCount: number): RustSdkCryptoJs.KeysBackupRequest {
const requestBody: IKeyBackup = {
rooms: {
"!room1:server": {
sessions: {},
},
},
};
for (let i = 0; i < keyCount; i++) {
requestBody.rooms["!room1:server"].sessions["session" + i] = {} as IKeyBackupSession;
}
return {
id: "id" + idGenerator++,
body: JSON.stringify(requestBody),
} as unknown as Mocked<RustSdkCryptoJs.KeysBackupRequest>;
}

beforeEach(async () => {
jest.useFakeTimers();
idGenerator = 0;

mockOlmMachine = {
getBackupKeys: jest.fn().mockResolvedValue({
backupVersion: TestData.SIGNED_BACKUP_DATA.version!,
decryptionKey: RustSdkCryptoJs.BackupDecryptionKey.fromBase64(TestData.BACKUP_DECRYPTION_KEY_BASE64),
} as unknown as RustSdkCryptoJs.BackupKeys),
backupRoomKeys: jest.fn(),
isBackupEnabled: jest.fn().mockResolvedValue(true),
enableBackupV1: jest.fn(),
verifyBackup: jest.fn().mockResolvedValue({
trusted: jest.fn().mockResolvedValue(true),
} as unknown as RustSdkCryptoJs.SignatureVerification),
roomKeyCounts: jest.fn(),
} as unknown as Mocked<RustSdkCryptoJs.OlmMachine>;

outgoingRequestProcessor = {
makeOutgoingRequest: jest.fn(),
} as unknown as Mocked<OutgoingRequestProcessor>;

rustBackupManager = new RustBackupManager(mockOlmMachine, httpAPi, outgoingRequestProcessor);

fetchMock.get("path:/_matrix/client/v3/room_keys/version", testData.SIGNED_BACKUP_DATA);
});

afterEach(() => {
fetchMock.reset();
jest.useRealTimers();
jest.resetAllMocks();
});

it("Should call expensive roomKeyCounts only once per loop", async () => {
const remainingEmitted: number[] = [];

const zeroRemainingWasEmitted = new Promise<void>((resolve) => {
rustBackupManager.on(CryptoEvent.KeyBackupSessionsRemaining, (count) => {
remainingEmitted.push(count);
if (count == 0) {
resolve();
}
});
});

// We want several batch of keys to check that we don't call expensive room key count several times
mockOlmMachine.backupRoomKeys
.mockResolvedValueOnce(mockBackupRequest(100))
.mockResolvedValueOnce(mockBackupRequest(100))
.mockResolvedValueOnce(mockBackupRequest(100))
.mockResolvedValueOnce(mockBackupRequest(100))
.mockResolvedValueOnce(mockBackupRequest(100))
.mockResolvedValueOnce(mockBackupRequest(100))
.mockResolvedValueOnce(mockBackupRequest(2))
.mockResolvedValue(null);

mockOlmMachine.roomKeyCounts.mockResolvedValue({
total: 602,
// First iteration won't call roomKeyCounts(); it will be called on the second iteration after 200 keys have been saved.
backedUp: 200,
});

await rustBackupManager.checkKeyBackupAndEnable(false);
await jest.runAllTimersAsync();

await zeroRemainingWasEmitted;

expect(outgoingRequestProcessor.makeOutgoingRequest).toHaveBeenCalledTimes(7);
expect(mockOlmMachine.roomKeyCounts).toHaveBeenCalledTimes(1);

// check event emission
expect(remainingEmitted[0]).toEqual(402);
expect(remainingEmitted[1]).toEqual(302);
expect(remainingEmitted[2]).toEqual(202);
expect(remainingEmitted[3]).toEqual(102);
expect(remainingEmitted[4]).toEqual(2);
expect(remainingEmitted[5]).toEqual(0);
});

it("Should not call expensive roomKeyCounts when only one iteration is needed", async () => {
const zeroRemainingWasEmitted = new Promise<void>((resolve) => {
rustBackupManager.on(CryptoEvent.KeyBackupSessionsRemaining, (count) => {
if (count == 0) {
resolve();
}
});
});

// Only returns 2 keys on the first call, then none.
mockOlmMachine.backupRoomKeys.mockResolvedValueOnce(mockBackupRequest(2)).mockResolvedValue(null);

await rustBackupManager.checkKeyBackupAndEnable(false);
await jest.runAllTimersAsync();

await zeroRemainingWasEmitted;

expect(outgoingRequestProcessor.makeOutgoingRequest).toHaveBeenCalledTimes(1);
expect(mockOlmMachine.roomKeyCounts).toHaveBeenCalledTimes(0);
});
});
66 changes: 58 additions & 8 deletions src/rust-crypto/backup.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ import { BackupDecryptor } from "../common-crypto/CryptoBackend";
import { IEncryptedPayload } from "../crypto/aes";
import { ImportRoomKeyProgressData, ImportRoomKeysOpts } from "../crypto-api";
import { IKeyBackupInfo } from "../crypto/keybackup";
import { IKeyBackup } from "../crypto/backup";

/** Authentification of the backup info, depends on algorithm */
type AuthData = KeyBackupInfo["auth_data"];
Expand Down Expand Up @@ -327,7 +328,13 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
await sleep(delay);

try {
let numFailures = 0; // number of consecutive network failures for exponential backoff
// number of consecutive network failures for exponential backoff
let numFailures = 0;
// The number of keys left to back up. (Populated lazily: see more comments below.)
let remainingToUploadCount: number | null = null;
// To avoid computing the key when only a few keys were added (after a sync for example),
// we compute the count only when at least two iterations are needed.
let isFirstIteration = true;

while (!this.stopped) {
// Get a batch of room keys to upload
Expand All @@ -346,19 +353,45 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,

if (!request || this.stopped || !this.activeBackupVersion) {
logger.log(`Backup: Ending loop for version ${this.activeBackupVersion}.`);
if (!request) {
// nothing more to upload
this.emit(CryptoEvent.KeyBackupSessionsRemaining, 0);
}
return;
}

try {
await this.outgoingRequestProcessor.makeOutgoingRequest(request);
numFailures = 0;
if (this.stopped) break;
try {
const keyCount = await this.olmMachine.roomKeyCounts();
const remaining = keyCount.total - keyCount.backedUp;
this.emit(CryptoEvent.KeyBackupSessionsRemaining, remaining);
} catch (err) {
logger.error("Backup: Failed to get key counts from rust crypto-sdk", err);

// Key count performance (`olmMachine.roomKeyCounts()`) can be pretty bad on some configurations.
// In particular, we detected on some M1 macs that when the object store reaches a threshold, the count
// performance stops growing in O(n) and suddenly becomes very slow (40s, 60s or more).
// For reference, the performance drop occurs around 300-400k keys on the platforms where this issue is observed.
// Even on other configurations, the count can take several seconds.
// This will block other operations on the database, like sending messages.
//
// This is a workaround to avoid calling `olmMachine.roomKeyCounts()` too often, and only when necessary.
// We don't call it on the first loop because there could be only a few keys to upload, and we don't want to wait for the count.
if (!isFirstIteration && remainingToUploadCount === null) {
try {
const keyCount = await this.olmMachine.roomKeyCounts();
remainingToUploadCount = keyCount.total - keyCount.backedUp;
} catch (err) {
logger.error("Backup: Failed to get key counts from rust crypto-sdk", err);
}
}

if (remainingToUploadCount !== null) {
this.emit(CryptoEvent.KeyBackupSessionsRemaining, remainingToUploadCount);
const keysCountInBatch = this.keysCountInBatch(request);
// `OlmMachine.roomKeyCounts` is called only once for the current backupKeysLoop. But new
// keys could be added during the current loop (after a sync for example).
// So the count can get out of sync with the real number of remaining keys to upload.
// Depending on the number of new keys imported and the time to complete the loop,
// this could result in multiple events being emitted with a remaining key count of 0.
remainingToUploadCount = Math.max(remainingToUploadCount - keysCountInBatch, 0);
}
} catch (err) {
numFailures++;
Expand All @@ -382,7 +415,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
// wait for that and then continue?
const waitTime = err.data.retry_after_ms;
if (waitTime > 0) {
sleep(waitTime);
await sleep(waitTime);
continue;
} // else go to the normal backoff
}
Expand All @@ -392,12 +425,29 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
// exponential backoff if we have failures
await sleep(1000 * Math.pow(2, Math.min(numFailures - 1, 4)));
}
isFirstIteration = false;
}
} finally {
this.backupKeysLoopRunning = false;
}
}

/**
* Utility method to count the number of keys in a backup request, in order to update the remaining keys count.
* This should be the chunk size of the backup request for all requests but the last, but we don't have access to it
* (it's static in the Rust SDK).
* @param batch - The backup request to count the keys from.
*
* @returns The number of keys in the backup request.
*/
private keysCountInBatch(batch: RustSdkCryptoJs.KeysBackupRequest): number {
const parsedBody: IKeyBackup = JSON.parse(batch.body);
let count = 0;
for (const { sessions } of Object.values(parsedBody.rooms)) {
count += Object.keys(sessions).length;
}
return count;
}
/**
* Get information about the current key backup from the server
*
Expand Down

0 comments on commit c993785

Please sign in to comment.