ElementR | backup: call expensive roomKeyCounts less often (#4015)

* ElementR | backup: call expensive `roomKeyCounts` less often

* review: Improve doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* review: Improve loop

* review: Add comment regarding slightly outdated remaining count

* Review: doc fix typo

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* 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 <1389908+richvdh@users.noreply.github.com>

* Review improve comment

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* Review: Clean and improve tests

* fix: wrong test name

---------

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

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

@@ -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);
+    });
+});

src/rust-crypto/backup.ts

@@ -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"];
@@ -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
@@ -346,6 +353,10 @@ 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;
                 }
 
@@ -353,12 +364,34 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                     await this.outgoingRequestProcessor.makeOutgoingRequest(request);
                     numFailures = 0;
                     if (this.stopped) break;
-                    try {
+
-                        const keyCount = await this.olmMachine.roomKeyCounts();
+                    // Key count performance (`olmMachine.roomKeyCounts()`) can be pretty bad on some configurations.
-                        const remaining = keyCount.total - keyCount.backedUp;
+                    // In particular, we detected on some M1 macs that when the object store reaches a threshold, the count
-                        this.emit(CryptoEvent.KeyBackupSessionsRemaining, remaining);
+                    // performance stops growing in O(n) and suddenly becomes very slow (40s, 60s or more).
-                    } catch (err) {
+                    // For reference, the performance drop occurs around 300-400k keys on the platforms where this issue is observed.
-                        logger.error("Backup: Failed to get key counts from rust crypto-sdk", err);
+                    // 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++;
@@ -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
                         }
@@ -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
      *