Support staged rollout of migration to Rust Crypto (#12184)

* Rust migration staged rollout

* Phased rollout unit tests

test/MatrixClientPeg-test.ts

@@ -144,6 +144,117 @@ describe("MatrixClientPeg", () => {
             expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, true);
         });
 
+        describe("Rust staged rollout", () => {
+            function mockSettingStore(
+                userIsUsingRust: boolean,
+                newLoginShouldUseRust: boolean,
+                rolloutPercent: number | null,
+            ) {
+                const originalGetValue = SettingsStore.getValue;
+                jest.spyOn(SettingsStore, "getValue").mockImplementation(
+                    (settingName: string, roomId: string | null = null, excludeDefault = false) => {
+                        if (settingName === "feature_rust_crypto") {
+                            return userIsUsingRust;
+                        }
+                        return originalGetValue(settingName, roomId, excludeDefault);
+                    },
+                );
+                const originalGetValueAt = SettingsStore.getValueAt;
+                jest.spyOn(SettingsStore, "getValueAt").mockImplementation(
+                    (level: SettingLevel, settingName: string) => {
+                        if (settingName === "feature_rust_crypto") {
+                            return newLoginShouldUseRust;
+                        }
+                        // if null we let the original implementation handle it to get the default
+                        if (settingName === "RustCrypto.staged_rollout_percent" && rolloutPercent !== null) {
+                            return rolloutPercent;
+                        }
+                        return originalGetValueAt(level, settingName);
+                    },
+                );
+            }
+
+            let mockSetValue: jest.SpyInstance;
+            let mockInitCrypto: jest.SpyInstance;
+            let mockInitRustCrypto: jest.SpyInstance;
+
+            beforeEach(() => {
+                mockSetValue = jest.spyOn(SettingsStore, "setValue").mockResolvedValue(undefined);
+                mockInitCrypto = jest.spyOn(testPeg.safeGet(), "initCrypto").mockResolvedValue(undefined);
+                mockInitRustCrypto = jest.spyOn(testPeg.safeGet(), "initRustCrypto").mockResolvedValue(undefined);
+            });
+
+            it("Should not migrate existing login if rollout is 0", async () => {
+                mockSettingStore(false, true, 0);
+
+                await testPeg.start();
+                expect(mockInitCrypto).toHaveBeenCalled();
+                expect(mockInitRustCrypto).not.toHaveBeenCalledTimes(1);
+
+                // we should have stashed the setting in the settings store
+                expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, false);
+            });
+
+            it("Should migrate existing login if rollout is 100", async () => {
+                mockSettingStore(false, true, 100);
+                await testPeg.start();
+                expect(mockInitCrypto).not.toHaveBeenCalled();
+                expect(mockInitRustCrypto).toHaveBeenCalledTimes(1);
+
+                // we should have stashed the setting in the settings store
+                expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, true);
+            });
+
+            it("Should migrate existing login if user is in rollout bucket", async () => {
+                mockSettingStore(false, true, 30);
+
+                // Use a device id that is known to be in the 30% bucket (hash modulo 100 < 30)
+                const spy = jest.spyOn(testPeg.get()!, "getDeviceId").mockReturnValue("AAA");
+
+                await testPeg.start();
+                expect(mockInitCrypto).not.toHaveBeenCalled();
+                expect(mockInitRustCrypto).toHaveBeenCalledTimes(1);
+
+                // we should have stashed the setting in the settings store
+                expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, true);
+
+                spy.mockReset();
+            });
+
+            it("Should not migrate existing login if rollout is malformed", async () => {
+                mockSettingStore(false, true, 100.1);
+
+                await testPeg.start();
+                expect(mockInitCrypto).toHaveBeenCalled();
+                expect(mockInitRustCrypto).not.toHaveBeenCalledTimes(1);
+
+                // we should have stashed the setting in the settings store
+                expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, false);
+            });
+
+            it("Default is to not migrate", async () => {
+                mockSettingStore(false, true, null);
+
+                await testPeg.start();
+                expect(mockInitCrypto).toHaveBeenCalled();
+                expect(mockInitRustCrypto).not.toHaveBeenCalledTimes(1);
+
+                // we should have stashed the setting in the settings store
+                expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, false);
+            });
+
+            it("Should not migrate if feature_rust_crypto is false", async () => {
+                mockSettingStore(false, false, 100);
+
+                await testPeg.start();
+                expect(mockInitCrypto).toHaveBeenCalled();
+                expect(mockInitRustCrypto).not.toHaveBeenCalledTimes(1);
+
+                // we should have stashed the setting in the settings store
+                expect(mockSetValue).toHaveBeenCalledWith("feature_rust_crypto", null, SettingLevel.DEVICE, false);
+            });
+        });
+
         it("should reload when store database closes for a guest user", async () => {
             testPeg.safeGet().isGuest = () => true;
             const emitter = new EventEmitter();

test/utils/PhasedRolloutFeature-test.ts

@@ -0,0 +1,89 @@
+/*
+Copyright 2024 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import { PhasedRolloutFeature } from "../../src/utils/PhasedRolloutFeature";
+
+describe("Test PhasedRolloutFeature", () => {
+    function randomUserId() {
+        const characters = "abcdefghijklmnopqrstuvwxyz0123456789.=_-/+";
+        let result = "";
+        const charactersLength = characters.length;
+        const idLength = Math.floor(Math.random() * 15) + 6; // Random number between 6 and 20
+        for (let i = 0; i < idLength; i++) {
+            result += characters.charAt(Math.floor(Math.random() * charactersLength));
+        }
+        return "@" + result + ":matrix.org";
+    }
+
+    function randomDeviceId() {
+        const characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
+        let result = "";
+        const charactersLength = characters.length;
+        for (let i = 0; i < 10; i++) {
+            result += characters.charAt(Math.floor(Math.random() * charactersLength));
+        }
+        return result;
+    }
+
+    it("should only accept valid percentage", () => {
+        expect(() => new PhasedRolloutFeature("test", 0.8)).toThrow();
+        expect(() => new PhasedRolloutFeature("test", -1)).toThrow();
+        expect(() => new PhasedRolloutFeature("test", 123)).toThrow();
+    });
+
+    it("should enable for all if percentage is 100", () => {
+        const phasedRolloutFeature = new PhasedRolloutFeature("test", 100);
+
+        for (let i = 0; i < 1000; i++) {
+            expect(phasedRolloutFeature.isFeatureEnabled(randomUserId())).toBeTruthy();
+        }
+    });
+
+    it("should not enable for anyone if percentage is 0", () => {
+        const phasedRolloutFeature = new PhasedRolloutFeature("test", 0);
+
+        for (let i = 0; i < 1000; i++) {
+            expect(phasedRolloutFeature.isFeatureEnabled(randomUserId())).toBeFalsy();
+        }
+    });
+
+    it("should enable for more users if percentage grows", () => {
+        let rolloutPercentage = 0;
+        let previousBatch: string[] = [];
+        const allUsers = new Array(1000).fill(0).map(() => randomDeviceId());
+
+        while (rolloutPercentage <= 90) {
+            rolloutPercentage += 10;
+            const nextRollout = new PhasedRolloutFeature("test", rolloutPercentage);
+            const nextBatch = allUsers.filter((userId) => nextRollout.isFeatureEnabled(userId));
+            expect(previousBatch.length).toBeLessThan(nextBatch.length);
+            expect(previousBatch.every((user) => nextBatch.includes(user))).toBeTruthy();
+            previousBatch = nextBatch;
+        }
+    });
+
+    it("should distribute differently depending on the feature name", () => {
+        const allUsers = new Array(1000).fill(0).map(() => randomUserId());
+
+        const featureARollout = new PhasedRolloutFeature("FeatureA", 50);
+        const featureBRollout = new PhasedRolloutFeature("FeatureB", 50);
+
+        const featureAUsers = allUsers.filter((userId) => featureARollout.isFeatureEnabled(userId));
+        const featureBUsers = allUsers.filter((userId) => featureBRollout.isFeatureEnabled(userId));
+
+        expect(featureAUsers).not.toEqual(featureBUsers);
+    });
+});