From 67c42e938ec18aeed101ac354d3c342b0bdabbb3 Mon Sep 17 00:00:00 2001
From: tamassoltesz <tamas@supertokens.com>
Date: Fri, 27 Sep 2024 12:13:38 +0200
Subject: [PATCH] fix: fixing unreliable mutithreaded bulk import with mysql

---
 .../bulkimport/BulkImportUserUtils.java       |  2 +-
 .../ProcessBulkUsersImportWorker.java         | 44 ++++++++++++++-----
 .../ProcessBulkImportUsersCronJobTest.java    | 12 ++---
 3 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/src/main/java/io/supertokens/bulkimport/BulkImportUserUtils.java b/src/main/java/io/supertokens/bulkimport/BulkImportUserUtils.java
index 7c6aac760..f2df6892e 100644
--- a/src/main/java/io/supertokens/bulkimport/BulkImportUserUtils.java
+++ b/src/main/java/io/supertokens/bulkimport/BulkImportUserUtils.java
@@ -265,7 +265,7 @@ private String validateAndNormaliseExternalUserId(String externalUserId, List<St
         }
 
         if (externalUserId.length() > 255) {
-            errors.add("externalUserId " + externalUserId + " is too long. Max length is 128.");
+            errors.add("externalUserId " + externalUserId + " is too long. Max length is 255.");
         }
 
         if (!allExternalUserIds.add(externalUserId)) {
diff --git a/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkUsersImportWorker.java b/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkUsersImportWorker.java
index 8122553d8..ef7dc9a6e 100644
--- a/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkUsersImportWorker.java
+++ b/src/main/java/io/supertokens/cronjobs/bulkimport/ProcessBulkUsersImportWorker.java
@@ -30,6 +30,7 @@
 import io.supertokens.pluginInterface.authRecipe.sqlStorage.AuthRecipeSQLStorage;
 import io.supertokens.pluginInterface.bulkimport.BulkImportStorage;
 import io.supertokens.pluginInterface.bulkimport.BulkImportUser;
+import io.supertokens.pluginInterface.bulkimport.exceptions.BulkImportTransactionRolledBackException;
 import io.supertokens.pluginInterface.bulkimport.sqlStorage.BulkImportSQLStorage;
 import io.supertokens.pluginInterface.exceptions.DbInitException;
 import io.supertokens.pluginInterface.exceptions.InvalidConfigException;
@@ -80,12 +81,16 @@ private void processMultipleUsers(AppIdentifier appIdentifier, List<BulkImportUs
             throws TenantOrAppNotFoundException, StorageQueryException, IOException,
             DbInitException {
 
-        BulkImportUser lastStartedUser = null;
+        BulkImportUser user = null;
         try {
-            for (BulkImportUser user : users) {
-                lastStartedUser = user;
-                if (Main.isTesting && Main.isTesting_skipBulkImportUserValidationInCronJob) {
+            boolean shouldRetryImmediately = false;
+            int userIndexPointer = 0;
+            while(userIndexPointer < users.size()){
+                user = users.get(userIndexPointer);
+                if ((Main.isTesting && Main.isTesting_skipBulkImportUserValidationInCronJob) || shouldRetryImmediately) {
                     // Skip validation when the flag is enabled during testing
+                    // Skip validation if it's a retry run. This already passed validation. A revalidation triggers
+                    // an invalid external user id already exists validation error - which is not true!
                 } else {
                     // Validate the user
                     bulkImportUserUtils.createBulkImportUserFromJSON(main, appIdentifier, user.toJsonObject(), user.id);
@@ -143,10 +148,11 @@ private void processMultipleUsers(AppIdentifier appIdentifier, List<BulkImportUs
                     }
                 }
 
-                bulkImportProxyStorage.startTransaction(con -> {
+                BulkImportUser finalUser = user;
+                shouldRetryImmediately = bulkImportProxyStorage.startTransaction(con -> {
                     try {
                         Storage[] allStoragesForApp = getAllProxyStoragesForApp(main, appIdentifier);
-                        BulkImport.processUserImportSteps(main, con, appIdentifier, bulkImportProxyStorage, user,
+                        BulkImport.processUserImportSteps(main, con, appIdentifier, bulkImportProxyStorage, finalUser,
                                 primaryLM, allStoragesForApp);
 
                         // We are updating the primaryUserId in the bulkImportUser entry. This will help us handle
@@ -154,7 +160,7 @@ private void processMultipleUsers(AppIdentifier appIdentifier, List<BulkImportUs
                         // If this update statement fails then the outer transaction will fail as well and the user
                         // will simpl be processed again. No inconsistency will happen in this
                         // case.
-                        baseTenantStorage.updateBulkImportUserPrimaryUserId(appIdentifier, user.id,
+                        baseTenantStorage.updateBulkImportUserPrimaryUserId(appIdentifier, finalUser.id,
                                 primaryLM.superTokensUserId);
 
                         // We need to commit the transaction manually because we have overridden that in the proxy
@@ -171,24 +177,40 @@ private void processMultipleUsers(AppIdentifier appIdentifier, List<BulkImportUs
                         // in the database.
                         // When processing the user again, we'll check if primaryUserId exists with the same email.
                         // In this case the user will exist, and we'll simply delete the entry.
-                        baseTenantStorage.deleteBulkImportUsers(appIdentifier, new String[]{user.id});
+                        baseTenantStorage.deleteBulkImportUsers(appIdentifier, new String[]{finalUser.id});
                     } catch (StorageTransactionLogicException e) {
                         // We need to rollback the transaction manually because we have overridden that in the proxy
                         // storage
                         bulkImportProxyStorage.rollbackTransactionForBulkImportProxyStorage();
-                        handleProcessUserExceptions(app, user, e, baseTenantStorage);
+                        if(isBulkImportTransactionRolledBackIsTheRealCause(e)){
+                            return true;
+                            //@see BulkImportTransactionRolledBackException for explanation
+                        }
+                        handleProcessUserExceptions(app, finalUser, e, baseTenantStorage);
                     }
-                    return null;
+                    return false;
                 });
 
+                if(!shouldRetryImmediately){
+                    userIndexPointer++;
+                }
             }
         } catch (StorageTransactionLogicException | InvalidBulkImportDataException | InvalidConfigException e) {
-            handleProcessUserExceptions(appIdentifier, lastStartedUser, e, baseTenantStorage);
+            handleProcessUserExceptions(appIdentifier, user, e, baseTenantStorage);
         } finally {
             closeAllProxyStorages(); //closing it here to reuse the existing connection with all the users
         }
     }
 
+    private boolean isBulkImportTransactionRolledBackIsTheRealCause(Throwable exception) {
+        if(exception instanceof BulkImportTransactionRolledBackException){
+            return true;
+        } else if(exception.getCause()!=null){
+            return isBulkImportTransactionRolledBackIsTheRealCause(exception.getCause());
+        }
+        return false;
+    }
+
     private void handleProcessUserExceptions(AppIdentifier appIdentifier, BulkImportUser user, Exception e,
                                              BulkImportSQLStorage baseTenantStorage)
             throws StorageQueryException {
diff --git a/src/test/java/io/supertokens/test/bulkimport/ProcessBulkImportUsersCronJobTest.java b/src/test/java/io/supertokens/test/bulkimport/ProcessBulkImportUsersCronJobTest.java
index 3e82171a0..97f14daf4 100644
--- a/src/test/java/io/supertokens/test/bulkimport/ProcessBulkImportUsersCronJobTest.java
+++ b/src/test/java/io/supertokens/test/bulkimport/ProcessBulkImportUsersCronJobTest.java
@@ -120,7 +120,7 @@ public void shouldProcessBulkImportUsersInTheSameTenant() throws Exception {
 
     @Test
     public void shouldProcessBulkImportUsersInNotSoLargeNumbersInTheSameTenant() throws Exception {
-        Utils.setValueInConfig("bulk_migration_parallelism", "2");
+        Utils.setValueInConfig("bulk_migration_parallelism", "8");
         TestingProcess process = startCronProcess();
         Main main = process.getProcess();
 
@@ -143,7 +143,7 @@ public void shouldProcessBulkImportUsersInNotSoLargeNumbersInTheSameTenant() thr
         List<BulkImportUser> users = generateBulkImportUser(usersCount);
         BulkImport.addUsers(appIdentifier, storage, users);
 
-        Thread.sleep(6000);
+        Thread.sleep(60000);
 
         List<BulkImportUser> usersAfterProcessing = storage.getBulkImportUsers(appIdentifier, 1000, null,
                 null, null);
@@ -159,7 +159,7 @@ public void shouldProcessBulkImportUsersInNotSoLargeNumbersInTheSameTenant() thr
 
     @Test
     public void shouldProcessBulkImportUsersInLargeNumbersInTheSameTenant() throws Exception {
-        Utils.setValueInConfig("bulk_migration_parallelism", "12");
+        Utils.setValueInConfig("bulk_migration_parallelism", "8");
 
         TestingProcess process = startCronProcess();
         Main main = process.getProcess();
@@ -259,7 +259,7 @@ public void shouldProcessBulkImportUsersInMultipleTenantsWithDifferentStorages()
 
     @Test
     public void shouldProcessBulkImportUsersInLargeNumberInMultipleTenantsWithDifferentStorages() throws Exception {
-        Utils.setValueInConfig("bulk_migration_parallelism", "4");
+        Utils.setValueInConfig("bulk_migration_parallelism", "8");
 
         TestingProcess process = startCronProcess();
         Main main = process.getProcess();
@@ -356,7 +356,7 @@ public void shouldDeleteEverythingFromTheDBIfAnythingFails() throws Exception {
 
     @Test
     public void shouldDeleteEverythingFromTheDBIfAnythingFailsOnMultipleThreads() throws Exception {
-        Utils.setValueInConfig("bulk_migration_parallelism", "12");
+        Utils.setValueInConfig("bulk_migration_parallelism", "8");
         // Creating a non-existing user role will result in an error.
         // Since, user role creation happens at the last step of the bulk import process, everything should be deleted from the DB.
 
@@ -399,7 +399,7 @@ public void shouldDeleteEverythingFromTheDBIfAnythingFailsOnMultipleThreads() th
 
     @Test
     public void shouldDeleteOnlyFailedFromTheDBIfAnythingFailsOnMultipleThreads() throws Exception {
-        Utils.setValueInConfig("bulk_migration_parallelism", "12");
+        Utils.setValueInConfig("bulk_migration_parallelism", "8");
         // Creating a non-existing user role will result in an error.
         // Since, user role creation happens at the last step of the bulk import process, everything should be deleted from the DB.