ElementR | backup: call expensive roomKeyCounts less often

spec/unit/rust-crypto/BackupKeyLoop.spec.ts

@@ -0,0 +1,160 @@
+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 { defer } from "../../../src/utils";
+import { RustBackupManager } from "../../../src/rust-crypto/backup";
+
+describe("PerSessionKeyBackupDownloader", () => {
+    /** 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();
+    });
+
+    it("Should call expensive roomKeyCounts only once per loop", async () => {
+        const lastKeysCalled = defer();
+        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))
+            .mockImplementation(async () => {
+                lastKeysCalled.resolve();
+                return null;
+            });
+
+        mockOlmMachine.roomKeyCounts.mockResolvedValue({
+            total: 602,
+            // first iteration don't call count, it will in second after 200 keys have been saved
+            backedUp: 200,
+        });
+
+        await rustBackupManager.checkKeyBackupAndEnable(false);
+        await jest.runAllTimersAsync();
+
+        await lastKeysCalled.promise;
+        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 for small chunks", async () => {
+        const lastKeysCalled = defer();
+
+        const zeroRemainingWasEmitted = new Promise<void>((resolve) => {
+            rustBackupManager.on(CryptoEvent.KeyBackupSessionsRemaining, (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(2)).mockImplementation(async () => {
+            lastKeysCalled.resolve();
+            return null;
+        });
+
+        mockOlmMachine.roomKeyCounts.mockResolvedValue({
+            total: 2,
+            backedUp: 0,
+        });
+
+        await rustBackupManager.checkKeyBackupAndEnable(false);
+        await jest.runAllTimersAsync();
+
+        await lastKeysCalled.promise;
+        await zeroRemainingWasEmitted;
+
+        expect(outgoingRequestProcessor.makeOutgoingRequest).toHaveBeenCalledTimes(1);
+        expect(mockOlmMachine.roomKeyCounts).toHaveBeenCalledTimes(0);
+    });
+});

src/rust-crypto/backup.ts

@@ -330,8 +330,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
         try {
             // number of consecutive network failures for exponential backoff
             let numFailures = 0;
-            // `olmMachine.roomKeyCounts()` is very slow on some configurations (see more comments below),
-            // we compute it only once per `backupKeysLoop` if necessary.
+            // 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.
@@ -361,29 +360,38 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                     return;
                 }
 
-                // Keys count performance (`olmMachine.roomKeyCounts()`) can be pretty bad on some configurations (big database in FF).
-                // We detected on some M1 mac that when the object store reach a threshold, the count performances stops growing in O(n) and
-                // suddenly become very slow (40s, 60s or more). 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);
-                    }
-                }
-
                 try {
                     await this.outgoingRequestProcessor.makeOutgoingRequest(request);
                     numFailures = 0;
                     if (this.stopped) break;
+
+                    // 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);
+                        // The `remainingToUploadCount` is computed 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);
-                        this.emit(CryptoEvent.KeyBackupSessionsRemaining, remainingToUploadCount);
                     }
                 } catch (err) {
                     numFailures++;
@@ -407,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
                         }
@@ -435,8 +443,8 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
     private keysCountInBatch(batch: RustSdkCryptoJs.KeysBackupRequest): number {
         const parsedBody: IKeyBackup = JSON.parse(batch.body);
         let count = 0;
-        for (const entry of Object.entries(parsedBody.rooms)) {
-            count += Object.keys(entry[1].sessions).length;
+        for (const { sessions } of Object.values(parsedBody.rooms)) {
+            count += Object.keys(sessions).length;
         }
         return count;
     }