Improve import onyx performance

src/components/ImportOnyxState/index.native.tsx

@@ -1,18 +1,17 @@
 import React, {useState} from 'react';
 import ReactNativeBlobUtil from 'react-native-blob-util';
-import Onyx, {useOnyx} from 'react-native-onyx';
+import {useOnyx} from 'react-native-onyx';
 import type {FileObject} from '@components/AttachmentModal';
-import {KEYS_TO_PRESERVE, setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
+import {setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
 import {setShouldForceOffline} from '@libs/actions/Network';
+import {rollbackOngoingRequest} from '@libs/actions/PersistedRequests';
+import {cleanAndTransformState, importState} from '@libs/ImportOnyxStateUtils';
 import Navigation from '@libs/Navigation/Navigation';
 import type {OnyxValues} from '@src/ONYXKEYS';
 import ONYXKEYS from '@src/ONYXKEYS';
 import ROUTES from '@src/ROUTES';
 import BaseImportOnyxState from './BaseImportOnyxState';
 import type ImportOnyxStateProps from './types';
-import {cleanAndTransformState} from './utils';
-
-const CHUNK_SIZE = 100;
 
 function readOnyxFile(fileUri: string) {
     const filePath = decodeURIComponent(fileUri.replace('file://', ''));
@@ -25,27 +24,6 @@ function readOnyxFile(fileUri: string) {
     });
 }
 
-function chunkArray<T>(array: T[], size: number): T[][] {
-    const result = [];
-    for (let i = 0; i < array.length; i += size) {
-        result.push(array.slice(i, i + size));
-    }
-    return result;
-}
-
-function applyStateInChunks(state: OnyxValues) {
-    const entries = Object.entries(state);
-    const chunks = chunkArray(entries, CHUNK_SIZE);
-
-    let promise = Promise.resolve();
-    chunks.forEach((chunk) => {
-        const partialOnyxState = Object.fromEntries(chunk) as Partial<OnyxValues>;
-        promise = promise.then(() => Onyx.multiSet(partialOnyxState));
-    });
-
-    return promise;
-}
-
 export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
     const [isErrorModalVisible, setIsErrorModalVisible] = useState(false);
     const [session] = useOnyx(ONYXKEYS.SESSION);
@@ -58,19 +36,23 @@ export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
         setIsLoading(true);
         readOnyxFile(file.uri)
             .then((fileContent: string) => {
+                rollbackOngoingRequest();
                 const transformedState = cleanAndTransformState<OnyxValues>(fileContent);
                 const currentUserSessionCopy = {...session};
                 setPreservedUserSession(currentUserSessionCopy);
                 setShouldForceOffline(true);
-                Onyx.clear(KEYS_TO_PRESERVE).then(() => {
-                    applyStateInChunks(transformedState).then(() => {
-                        setIsUsingImportedState(true);
-                        Navigation.navigate(ROUTES.HOME);
-                    });
-                });
+                return importState(transformedState);
             })
-            .catch(() => {
+            .then(() => {
+                setIsUsingImportedState(true);
+                Navigation.navigate(ROUTES.HOME);
+            })
+            .catch((error) => {
+                console.error('Error importing state:', error);
                 setIsErrorModalVisible(true);
+            })
+            .finally(() => {
+                setIsLoading(false);
             });
     };
 

src/components/ImportOnyxState/index.tsx

@@ -1,15 +1,16 @@
 import React, {useState} from 'react';
-import Onyx, {useOnyx} from 'react-native-onyx';
+import {useOnyx} from 'react-native-onyx';
 import type {FileObject} from '@components/AttachmentModal';
-import {KEYS_TO_PRESERVE, setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
+import {setIsUsingImportedState, setPreservedUserSession} from '@libs/actions/App';
 import {setShouldForceOffline} from '@libs/actions/Network';
+import {rollbackOngoingRequest} from '@libs/actions/PersistedRequests';
+import {cleanAndTransformState, importState} from '@libs/ImportOnyxStateUtils';
 import Navigation from '@libs/Navigation/Navigation';
 import type {OnyxValues} from '@src/ONYXKEYS';
 import ONYXKEYS from '@src/ONYXKEYS';
 import ROUTES from '@src/ROUTES';
 import BaseImportOnyxState from './BaseImportOnyxState';
 import type ImportOnyxStateProps from './types';
-import {cleanAndTransformState} from './utils';
 
 export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
     const [isErrorModalVisible, setIsErrorModalVisible] = useState(false);
@@ -21,26 +22,33 @@ export default function ImportOnyxState({setIsLoading}: ImportOnyxStateProps) {
         }
 
         setIsLoading(true);
+
         const blob = new Blob([file as BlobPart]);
         const response = new Response(blob);
 
         response
             .text()
             .then((text) => {
+                rollbackOngoingRequest();
                 const fileContent = text;
+
                 const transformedState = cleanAndTransformState<OnyxValues>(fileContent);
+
                 const currentUserSessionCopy = {...session};
                 setPreservedUserSession(currentUserSessionCopy);
                 setShouldForceOffline(true);
-                Onyx.clear(KEYS_TO_PRESERVE).then(() => {
-                    Onyx.multiSet(transformedState).then(() => {
-                        setIsUsingImportedState(true);
-                        Navigation.navigate(ROUTES.HOME);
-                    });
-                });
+
+                return importState(transformedState);
+            })
+            .then(() => {
+                setIsUsingImportedState(true);
+                Navigation.navigate(ROUTES.HOME);
             })
-            .catch(() => {
+            .catch((error) => {
+                console.error('Error importing state:', error);
                 setIsErrorModalVisible(true);
+            })
+            .finally(() => {
                 setIsLoading(false);
             });
     };

src/components/ImportOnyxState/utils.ts → src/libs/ImportOnyxStateUtils.ts

@@ -1,6 +1,10 @@
 import cloneDeep from 'lodash/cloneDeep';
+import type {OnyxEntry, OnyxKey} from 'react-native-onyx';
 import type {UnknownRecord} from 'type-fest';
+import type {OnyxCollectionKey, OnyxCollectionValuesMapping, OnyxValues} from '@src/ONYXKEYS';
 import ONYXKEYS from '@src/ONYXKEYS';
+import type CollectionDataSet from '@src/types/utils/CollectionDataSet';
+import {clearOnyxStateBeforeImport, importOnyxCollectionState, importOnyxRegularState} from './actions/ImportOnyxState';
 
 // List of Onyx keys from the .txt file we want to keep for the local override
 const keysToOmit = [ONYXKEYS.ACTIVE_CLIENTS, ONYXKEYS.FREQUENTLY_USED_EMOJIS, ONYXKEYS.NETWORK, ONYXKEYS.CREDENTIALS, ONYXKEYS.PREFERRED_THEME];
@@ -12,7 +16,7 @@ function isRecord(value: unknown): value is Record<string, unknown> {
 function transformNumericKeysToArray(data: UnknownRecord): UnknownRecord | unknown[] {
     const dataCopy = cloneDeep(data);
     if (!isRecord(dataCopy)) {
-        return Array.isArray(dataCopy) ? (dataCopy as UnknownRecord[]).map(transformNumericKeysToArray) : (dataCopy as UnknownRecord);
+        return Array.isArray(dataCopy) ? (dataCopy as UnknownRecord[]).map(transformNumericKeysToArray) : dataCopy;
     }
 
     const keys = Object.keys(dataCopy);
@@ -50,4 +54,35 @@ function cleanAndTransformState<T>(state: string): T {
     return transformedState;
 }
 
-export {transformNumericKeysToArray, cleanAndTransformState};
+function importState(transformedState: OnyxValues): Promise<void> {
+    const collectionKeys = [...new Set(Object.values(ONYXKEYS.COLLECTION))];
+    const collectionsMap = new Map<keyof OnyxCollectionValuesMapping, CollectionDataSet<OnyxCollectionKey>>();
+    const regularState: Partial<Record<OnyxKey, OnyxEntry<OnyxKey>>> = {};
+
+    Object.entries(transformedState).forEach(([entryKey, entryValue]) => {
+        const key = entryKey as OnyxKey;
+        const value = entryValue as NonNullable<OnyxEntry<OnyxKey>>;
+
+        const collectionKey = collectionKeys.find((cKey) => key.startsWith(cKey));
+        if (collectionKey) {
+            if (!collectionsMap.has(collectionKey)) {
+                collectionsMap.set(collectionKey, {});
+            }
+
+            const collection = collectionsMap.get(collectionKey);
+            if (!collection) {
+                return;
+            }
+
+            collection[key as OnyxCollectionKey] = value;
+        } else {
+            regularState[key] = value;
+        }
+    });
+
+    return clearOnyxStateBeforeImport()
+        .then(() => importOnyxCollectionState(collectionsMap))
+        .then(() => importOnyxRegularState(regularState));
+}
+
+export {cleanAndTransformState, importState, transformNumericKeysToArray};

src/libs/actions/App.ts

@@ -27,7 +27,7 @@ import ROUTES from '@src/ROUTES';
 import type * as OnyxTypes from '@src/types/onyx';
 import type {OnyxData} from '@src/types/onyx/Request';
 import {setShouldForceOffline} from './Network';
-import {getAll, save} from './PersistedRequests';
+import {getAll, rollbackOngoingRequest, save} from './PersistedRequests';
 import {createDraftInitialWorkspace, createWorkspace, generatePolicyID} from './Policy/Policy';
 import {resolveDuplicationConflictAction} from './RequestConflictUtils';
 import {isAnonymousUser} from './Session';
@@ -570,38 +570,42 @@ function clearOnyxAndResetApp(shouldNavigateToHomepage?: boolean) {
     const isStateImported = isUsingImportedState;
     const shouldUseStagingServer = preservedShouldUseStagingServer;
     const sequentialQueue = getAll();
-    Onyx.clear(KEYS_TO_PRESERVE).then(() => {
-        // Network key is preserved, so when using imported state, we should stop forcing offline mode so that the app can re-fetch the network
-        if (isStateImported) {
-            setShouldForceOffline(false);
-        }
-
-        if (shouldNavigateToHomepage) {
-            Navigation.navigate(ROUTES.HOME);
-        }
 
-        if (preservedUserSession) {
-            Onyx.set(ONYXKEYS.SESSION, preservedUserSession);
-            Onyx.set(ONYXKEYS.PRESERVED_USER_SESSION, null);
-        }
+    rollbackOngoingRequest();
+    Onyx.clear(KEYS_TO_PRESERVE)
+        .then(() => {
+            // Network key is preserved, so when using imported state, we should stop forcing offline mode so that the app can re-fetch the network
+            if (isStateImported) {
+                setShouldForceOffline(false);
+            }
 
-        if (shouldUseStagingServer) {
-            Onyx.set(ONYXKEYS.USER, {shouldUseStagingServer});
-        }
+            if (shouldNavigateToHomepage) {
+                Navigation.navigate(ROUTES.HOME);
+            }
 
-        // Requests in a sequential queue should be called even if the Onyx state is reset, so we do not lose any pending data.
-        // However, the OpenApp request must be called before any other request in a queue to ensure data consistency.
-        // To do that, sequential queue is cleared together with other keys, and then it's restored once the OpenApp request is resolved.
-        openApp().then(() => {
-            if (!sequentialQueue || isStateImported) {
-                return;
+            if (preservedUserSession) {
+                Onyx.set(ONYXKEYS.SESSION, preservedUserSession);
+                Onyx.set(ONYXKEYS.PRESERVED_USER_SESSION, null);
             }
 
-            sequentialQueue.forEach((request) => {
-                save(request);
+            if (shouldUseStagingServer) {
+                Onyx.set(ONYXKEYS.USER, {shouldUseStagingServer});
+            }
+        })
+        .then(() => {
+            // Requests in a sequential queue should be called even if the Onyx state is reset, so we do not lose any pending data.
+            // However, the OpenApp request must be called before any other request in a queue to ensure data consistency.
+            // To do that, sequential queue is cleared together with other keys, and then it's restored once the OpenApp request is resolved.
+            openApp().then(() => {
+                if (!sequentialQueue || isStateImported) {
+                    return;
+                }
+
+                sequentialQueue.forEach((request) => {
+                    save(request);
+                });
             });
         });
-    });
     clearSoundAssetsCache();
 }
 

src/libs/actions/ImportOnyxState.ts

@@ -0,0 +1,25 @@
+import type {OnyxEntry, OnyxKey} from 'react-native-onyx';
+import Onyx from 'react-native-onyx';
+import type {OnyxCollectionKey, OnyxCollectionValuesMapping, OnyxValues} from '@src/ONYXKEYS';
+import type CollectionDataSet from '@src/types/utils/CollectionDataSet';
+import {KEYS_TO_PRESERVE} from './App';
+
+function clearOnyxStateBeforeImport(): Promise<void> {
+    return Onyx.clear(KEYS_TO_PRESERVE);
+}
+
+function importOnyxCollectionState(collectionsMap: Map<keyof OnyxCollectionValuesMapping, CollectionDataSet<OnyxCollectionKey>>): Promise<void[]> {
+    const collectionPromises = Array.from(collectionsMap.entries()).map(([baseKey, items]) => {
+        return items ? Onyx.setCollection(baseKey, items) : Promise.resolve();
+    });
+    return Promise.all(collectionPromises);
+}
+
+function importOnyxRegularState(state: Partial<Record<OnyxKey, OnyxEntry<OnyxKey>>>): Promise<void> {
+    if (Object.keys(state).length > 0) {
+        return Onyx.multiSet(state as Partial<OnyxValues>);
+    }
+    return Promise.resolve();
+}
+
+export {clearOnyxStateBeforeImport, importOnyxCollectionState, importOnyxRegularState};

tests/unit/ImportOnyxStateTest.ts

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/naming-convention */
-import {cleanAndTransformState, transformNumericKeysToArray} from '@components/ImportOnyxState/utils';
 import ONYXKEYS from '@src/ONYXKEYS';
+import {cleanAndTransformState, transformNumericKeysToArray} from '../../src/libs/ImportOnyxStateUtils';
 
 describe('transformNumericKeysToArray', () => {
     it('converts object with numeric keys to array', () => {