From 7fa48f1f78cd3b6e44f5a2fa36994da4bea4e9fb Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Thu, 18 Jan 2024 14:05:41 +0100
Subject: [PATCH 01/15] ElementR | backup: call expensive `roomKeyCounts` less
 often

---
 src/rust-crypto/backup.ts | 56 ++++++++++++++++++++++++++++++++++-----
 1 file changed, 49 insertions(+), 7 deletions(-)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index b969013eb10..d372727a09a 100644
--- a/src/rust-crypto/backup.ts
+++ b/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"];
@@ -323,7 +324,14 @@ 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;
+            // `olmMachine.roomKeyCounts()` is very slow on some configurations (see more comments below),
+            // we compute it only once per `backupKeysLoop` if necessary.
+            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
@@ -342,20 +350,37 @@ 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;
+                // 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();
-                        const remaining = keyCount.total - keyCount.backedUp;
-                        this.emit(CryptoEvent.KeyBackupSessionsRemaining, remaining);
+                        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;
+                    if (remainingToUploadCount !== null) {
+                        const keysCountInBatch = this.keysCountInBatch(request);
+                        remainingToUploadCount = Math.max(remainingToUploadCount - keysCountInBatch, 0);
+                        this.emit(CryptoEvent.KeyBackupSessionsRemaining, remainingToUploadCount);
+                    }
                 } catch (err) {
                     numFailures++;
                     logger.error("Backup: Error processing backup request for rust crypto-sdk", err);
@@ -388,12 +413,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 entry of Object.entries(parsedBody.rooms)) {
+            count += Object.keys(entry[1].sessions).length;
+        }
+        return count;
+    }
     /**
      * Get information about the current key backup from the server
      *

From bf72896221eda6a1ed580cb6915c68a82968d7f5 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 13:44:53 +0100
Subject: [PATCH 02/15] review: Improve doc

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
---
 src/rust-crypto/backup.ts | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index d372727a09a..fa7dc5d5a2a 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -326,8 +326,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.

From 418074a1961958d3a38c6b97280a91238643685b Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 13:52:28 +0100
Subject: [PATCH 03/15] review: Improve loop

---
 src/rust-crypto/backup.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index fa7dc5d5a2a..4c19e0ea137 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -430,8 +430,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;
     }

From 8dc1ff717d10d31b63b9af7a85cc3e0d687a0252 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 14:02:01 +0100
Subject: [PATCH 04/15] review: Add comment regarding slightly outdated
 remaining count

---
 src/rust-crypto/backup.ts | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index 4c19e0ea137..003186290bd 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -377,6 +377,11 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                     if (this.stopped) break;
                     if (remainingToUploadCount !== null) {
                         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);
                     }

From 611c236414ea9d6004c0bfa2d00257a1bad792a2 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 14:04:37 +0100
Subject: [PATCH 05/15] Review: doc fix typo

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
---
 src/rust-crypto/backup.ts | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index 003186290bd..99e79300891 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -356,10 +356,12 @@ 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
+                // 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). 
+                // 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) {

From 3e7fa5639f609c85cd9791e24a8c275681830c9b Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 14:07:47 +0100
Subject: [PATCH 06/15] review: refactor code order, count after doing the
 request

---
 src/rust-crypto/backup.ts | 35 ++++++++++++++++++-----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index 99e79300891..789f77911d5 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -356,27 +356,28 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                     return;
                 }
 
-                // 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). 
-                // 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).
+                    // 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) {
                         const keysCountInBatch = this.keysCountInBatch(request);
                         // The `remainingToUploadCount` is computed only once for the current backupKeysLoop. But new

From e5349f9bc9e6e8e46686f36b404895291fe54be6 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 14:10:23 +0100
Subject: [PATCH 07/15] review: Missing await on sleep for limit exceeded

---
 src/rust-crypto/backup.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index 789f77911d5..41d15314f8d 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -410,7 +410,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
                         }

From ac383aabf611cf5a276cad9596405ba0cbfe372d Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 14:15:25 +0100
Subject: [PATCH 08/15] review: Comment | add a note for when performance drops

---
 src/rust-crypto/backup.ts | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index 41d15314f8d..5a511f7becf 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -364,6 +364,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                     // 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.
                     //

From 76c4194e8335732c6ca20e56c7ff52a5b371c422 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 15:29:05 +0100
Subject: [PATCH 09/15] Backup: add upload loop test for rust

---
 spec/unit/rust-crypto/BackupKeyLoop.spec.ts | 156 ++++++++++++++++++++
 src/rust-crypto/backup.ts                   |   2 +-
 2 files changed, 157 insertions(+), 1 deletion(-)
 create mode 100644 spec/unit/rust-crypto/BackupKeyLoop.spec.ts

diff --git a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
new file mode 100644
index 00000000000..64d98272dd2
--- /dev/null
+++ b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
@@ -0,0 +1,156 @@
+import { Mocked } from "jest-mock";
+import fetchMock from "fetch-mock-jest";
+
+import { RustBackupManager } from "../../../src/rust-crypto/backup";
+import * as RustSdkCryptoJs from "../../../../matrix-rust-sdk-crypto-wasm";
+import { KeysBackupRequest, OlmMachine, SignatureVerification } from "../../../../matrix-rust-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";
+
+describe("PerSessionKeyBackupDownloader", () => {
+    /** The backup manager under test */
+    let rustBackupManager: RustBackupManager;
+
+    let mockOlmMachine: Mocked<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): 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<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 SignatureVerification),
+            roomKeyCounts: jest.fn(),
+        } as unknown as Mocked<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);
+    });
+
+    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);
+    });
+});
diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index 5a511f7becf..29c84c7cedd 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -380,6 +380,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                     }
 
                     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).
@@ -387,7 +388,6 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                         // 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++;

From a263ac722d6f7d2af4938d2ecef1e432d9cf5ac7 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 15:34:50 +0100
Subject: [PATCH 10/15] test: quick fix backup loop tests

---
 spec/unit/rust-crypto/BackupKeyLoop.spec.ts | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
index 64d98272dd2..19029d8f722 100644
--- a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
+++ b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
@@ -1,7 +1,6 @@
 import { Mocked } from "jest-mock";
 import fetchMock from "fetch-mock-jest";
 
-import { RustBackupManager } from "../../../src/rust-crypto/backup";
 import * as RustSdkCryptoJs from "../../../../matrix-rust-sdk-crypto-wasm";
 import { KeysBackupRequest, OlmMachine, SignatureVerification } from "../../../../matrix-rust-sdk-crypto-wasm";
 import { CryptoEvent, HttpApiEvent, HttpApiEventHandlerMap, MatrixHttpApi, TypedEventEmitter } from "../../../src";
@@ -11,6 +10,7 @@ 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 */
@@ -71,6 +71,11 @@ describe("PerSessionKeyBackupDownloader", () => {
         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[] = [];

From ce91e5e266a18e854f2cff27314e1feb91f6c942 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Fri, 19 Jan 2024 15:38:08 +0100
Subject: [PATCH 11/15] test: quick fix imports backup loop tests

---
 spec/unit/rust-crypto/BackupKeyLoop.spec.ts | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
index 19029d8f722..ca84b1762b3 100644
--- a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
+++ b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
@@ -1,8 +1,7 @@
 import { Mocked } from "jest-mock";
 import fetchMock from "fetch-mock-jest";
+import * as RustSdkCryptoJs from "@matrix-org/matrix-sdk-crypto-wasm";
 
-import * as RustSdkCryptoJs from "../../../../matrix-rust-sdk-crypto-wasm";
-import { KeysBackupRequest, OlmMachine, SignatureVerification } from "../../../../matrix-rust-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";
@@ -16,7 +15,7 @@ describe("PerSessionKeyBackupDownloader", () => {
     /** The backup manager under test */
     let rustBackupManager: RustBackupManager;
 
-    let mockOlmMachine: Mocked<OlmMachine>;
+    let mockOlmMachine: Mocked<RustSdkCryptoJs.OlmMachine>;
 
     let outgoingRequestProcessor: Mocked<OutgoingRequestProcessor>;
 
@@ -27,7 +26,7 @@ describe("PerSessionKeyBackupDownloader", () => {
     });
 
     let idGenerator = 0;
-    function mockBackupRequest(keyCount: number): KeysBackupRequest {
+    function mockBackupRequest(keyCount: number): RustSdkCryptoJs.KeysBackupRequest {
         const requestBody: IKeyBackup = {
             rooms: {
                 "!room1:server": {
@@ -41,7 +40,7 @@ describe("PerSessionKeyBackupDownloader", () => {
         return {
             id: "id" + idGenerator++,
             body: JSON.stringify(requestBody),
-        } as unknown as Mocked<KeysBackupRequest>;
+        } as unknown as Mocked<RustSdkCryptoJs.KeysBackupRequest>;
     }
 
     beforeEach(async () => {
@@ -58,9 +57,9 @@ describe("PerSessionKeyBackupDownloader", () => {
             enableBackupV1: jest.fn(),
             verifyBackup: jest.fn().mockResolvedValue({
                 trusted: jest.fn().mockResolvedValue(true),
-            } as unknown as SignatureVerification),
+            } as unknown as RustSdkCryptoJs.SignatureVerification),
             roomKeyCounts: jest.fn(),
-        } as unknown as Mocked<OlmMachine>;
+        } as unknown as Mocked<RustSdkCryptoJs.OlmMachine>;
 
         outgoingRequestProcessor = {
             makeOutgoingRequest: jest.fn(),

From 475e6642a9e8a9363389a8b7051d405205e7e8e2 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Mon, 22 Jan 2024 14:22:10 +0100
Subject: [PATCH 12/15] review: improve comment

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
---
 src/rust-crypto/backup.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/rust-crypto/backup.ts b/src/rust-crypto/backup.ts
index ef3ba1b1185..21964eb3afb 100644
--- a/src/rust-crypto/backup.ts
+++ b/src/rust-crypto/backup.ts
@@ -386,7 +386,7 @@ export class RustBackupManager extends TypedEventEmitter<RustBackupCryptoEvents,
                     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
+                        // `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,

From 3142b2b3d10273c78ddc0d05b11902ce31ae35b6 Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Mon, 22 Jan 2024 14:22:30 +0100
Subject: [PATCH 13/15] Review improve comment

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
---
 spec/unit/rust-crypto/BackupKeyLoop.spec.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
index ca84b1762b3..bb1c39eddcc 100644
--- a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
+++ b/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
@@ -104,7 +104,7 @@ describe("PerSessionKeyBackupDownloader", () => {
 
         mockOlmMachine.roomKeyCounts.mockResolvedValue({
             total: 602,
-            // first iteration don't call count, it will in second after 200 keys have been saved
+            // First iteration won't call roomKeyCounts(); it will be called on the second iteration after 200 keys have been saved.
             backedUp: 200,
         });
 

From 89c76eb8806b4cd6d5c7b545a9e7a698c0ffaece Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Mon, 22 Jan 2024 14:58:46 +0100
Subject: [PATCH 14/15] Review: Clean and improve tests

---
 .../{BackupKeyLoop.spec.ts => backup.spec.ts} | 28 ++++---------------
 1 file changed, 6 insertions(+), 22 deletions(-)
 rename spec/unit/rust-crypto/{BackupKeyLoop.spec.ts => backup.spec.ts} (87%)

diff --git a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts b/spec/unit/rust-crypto/backup.spec.ts
similarity index 87%
rename from spec/unit/rust-crypto/BackupKeyLoop.spec.ts
rename to spec/unit/rust-crypto/backup.spec.ts
index bb1c39eddcc..e8d15862b67 100644
--- a/spec/unit/rust-crypto/BackupKeyLoop.spec.ts
+++ b/spec/unit/rust-crypto/backup.spec.ts
@@ -8,10 +8,9 @@ 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", () => {
+describe("Import keys from backup", () => {
     /** The backup manager under test */
     let rustBackupManager: RustBackupManager;
 
@@ -73,10 +72,10 @@ describe("PerSessionKeyBackupDownloader", () => {
     afterEach(() => {
         fetchMock.reset();
         jest.useRealTimers();
+        jest.resetAllMocks();
     });
 
     it("Should call expensive roomKeyCounts only once per loop", async () => {
-        const lastKeysCalled = defer();
         const remainingEmitted: number[] = [];
 
         const zeroRemainingWasEmitted = new Promise<void>((resolve) => {
@@ -97,10 +96,7 @@ describe("PerSessionKeyBackupDownloader", () => {
             .mockResolvedValueOnce(mockBackupRequest(100))
             .mockResolvedValueOnce(mockBackupRequest(100))
             .mockResolvedValueOnce(mockBackupRequest(2))
-            .mockImplementation(async () => {
-                lastKeysCalled.resolve();
-                return null;
-            });
+            .mockResolvedValue(null);
 
         mockOlmMachine.roomKeyCounts.mockResolvedValue({
             total: 602,
@@ -111,7 +107,6 @@ describe("PerSessionKeyBackupDownloader", () => {
         await rustBackupManager.checkKeyBackupAndEnable(false);
         await jest.runAllTimersAsync();
 
-        await lastKeysCalled.promise;
         await zeroRemainingWasEmitted;
 
         expect(outgoingRequestProcessor.makeOutgoingRequest).toHaveBeenCalledTimes(7);
@@ -126,9 +121,7 @@ describe("PerSessionKeyBackupDownloader", () => {
         expect(remainingEmitted[5]).toEqual(0);
     });
 
-    it("Should not call expensive roomKeyCounts for small chunks", async () => {
-        const lastKeysCalled = defer();
-
+    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) {
@@ -137,21 +130,12 @@ describe("PerSessionKeyBackupDownloader", () => {
             });
         });
 
-        // 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,
-        });
+        // 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 lastKeysCalled.promise;
         await zeroRemainingWasEmitted;
 
         expect(outgoingRequestProcessor.makeOutgoingRequest).toHaveBeenCalledTimes(1);

From 99648ae135b9ecd7dff7e35e9e342ecb51c0682b Mon Sep 17 00:00:00 2001
From: Valere <bill.carson@valrsoft.com>
Date: Mon, 22 Jan 2024 15:07:35 +0100
Subject: [PATCH 15/15] fix: wrong test name

---
 spec/unit/rust-crypto/backup.spec.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/spec/unit/rust-crypto/backup.spec.ts b/spec/unit/rust-crypto/backup.spec.ts
index e8d15862b67..1a9c1435663 100644
--- a/spec/unit/rust-crypto/backup.spec.ts
+++ b/spec/unit/rust-crypto/backup.spec.ts
@@ -10,7 +10,7 @@ import { IKeyBackup } from "../../../src/crypto/backup";
 import { IKeyBackupSession } from "../../../src/crypto/keybackup";
 import { RustBackupManager } from "../../../src/rust-crypto/backup";
 
-describe("Import keys from backup", () => {
+describe("Upload keys to backup", () => {
     /** The backup manager under test */
     let rustBackupManager: RustBackupManager;