Add retry logic to copy account during personal to work profile account migration

Bug: b/329424312
Change-Id: I84661a09502de8bbce15712cb785ddbc8520a324
Test: Manual

2 files changed, 31 insertions, 2 deletions

diff --git a/core/java/android/app/admin/flags/flags.aconfig b/core/java/android/app/admin/flags/flags.aconfig
index 4fa45be57a11..be0aaff95c96 100644
--- a/core/java/android/app/admin/flags/flags.aconfig
+++ b/core/java/android/app/admin/flags/flags.aconfig
@@ -174,6 +174,17 @@ flag {
 }
 
 flag {
+  name: "copy_account_with_retry_enabled"
+  namespace: "enterprise"
+  description: "Retry copy and remove account from personal to work profile in case of failure"
+  bug: "329424312"
+  metadata {
+      purpose: PURPOSE_BUGFIX
+  }
+}
+
+
+flag {
   name: "esim_management_ux_enabled"
   namespace: "enterprise"
   description: "Enable UX changes for esim management"
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index f955b91136a3..c3175d64e2d3 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -888,6 +888,8 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
     private static final String APPLICATION_EXEMPTIONS_FLAG = "application_exemptions";
     private static final boolean DEFAULT_APPLICATION_EXEMPTIONS_FLAG = true;
 
+    private static final int RETRY_COPY_ACCOUNT_ATTEMPTS = 3;
+
     /**
      * For apps targeting U+
      * Enable multiple admins to coexist on the same device.
@@ -21514,13 +21516,26 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
             Slogf.w(LOG_TAG, "sourceUser and targetUser are the same, won't migrate account.");
             return;
         }
-        copyAccount(targetUser, sourceUser, accountToMigrate, callerPackage);
+
+        if (Flags.copyAccountWithRetryEnabled()) {
+            boolean copySucceeded = false;
+            int retryAttemptsLeft = RETRY_COPY_ACCOUNT_ATTEMPTS;
+            while (!copySucceeded && (retryAttemptsLeft > 0)) {
+                Slogf.i(LOG_TAG, "Copying account. Attempts left : " + retryAttemptsLeft);
+                copySucceeded =
+                        copyAccount(targetUser, sourceUser, accountToMigrate, callerPackage);
+                retryAttemptsLeft--;
+            }
+        } else {
+            copyAccount(targetUser, sourceUser, accountToMigrate, callerPackage);
+        }
         if (!keepAccountMigrated) {
             removeAccount(accountToMigrate, sourceUserId);
         }
+
     }
 
-    private void copyAccount(
+    private boolean copyAccount(
             UserHandle targetUser, UserHandle sourceUser, Account accountToMigrate,
             String callerPackage) {
         final long startTime = SystemClock.elapsedRealtime();
@@ -21538,6 +21553,8 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
                         DevicePolicyEnums.PLATFORM_PROVISIONING_COPY_ACCOUNT_MS,
                         startTime,
                         callerPackage);
+                Slogf.i(LOG_TAG, "Copy account successful to " + targetUser);
+                return true;
             } else {
                 logCopyAccountStatus(COPY_ACCOUNT_FAILED, callerPackage);
                 Slogf.e(LOG_TAG, "Failed to copy account to " + targetUser);
@@ -21550,6 +21567,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
             logCopyAccountStatus(COPY_ACCOUNT_EXCEPTION, callerPackage);
             Slogf.e(LOG_TAG, "Exception copying account to " + targetUser, e);
         }
+        return false;
     }
 
     private static void logCopyAccountStatus(@CopyAccountStatus int status, String callerPackage) {