Pass in cancellation reason to B&R tasks

This is to give more granular information to tasks about why they got
cancelled. This will help them report why they failed.

The cancelAll parameter was only set to false for cancellations due to
timeouts, so I've removed it.

This is a no-op refactor.

Flag: EXEMPT refactor
Bug: 399645990
Test: atest CtsBackupTestCases
Change-Id: I1bd6d9eff42c5edeb0f4df453bae792cf994bb6a

12 files changed, 117 insertions, 93 deletions

diff --git a/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java b/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java
index 9a353fbc45bf..867cd51e1c2b 100644
--- a/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java
+++ b/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java
@@ -44,6 +44,7 @@ import android.util.Slog;
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.LocalServices;
+import com.android.server.backup.BackupRestoreTask.CancellationReason;
 import com.android.server.backup.internal.LifecycleOperationStorage;
 
 import java.util.Set;
@@ -298,20 +299,22 @@ public class BackupAgentConnectionManager {
             // Offload operation cancellation off the main thread as the cancellation callbacks
             // might call out to BackupTransport. Other operations started on the same package
             // before the cancellation callback has executed will also be cancelled by the callback.
-            Runnable cancellationRunnable = () -> {
-                // handleCancel() causes the PerformFullTransportBackupTask to go on to
-                // tearDownAgentAndKill: that will unbindBackupAgent in the Activity Manager, so
-                // that the package being backed up doesn't get stuck in restricted mode until the
-                // backup time-out elapses.
-                for (int token : mOperationStorage.operationTokensForPackage(packageName)) {
-                    if (DEBUG) {
-                        Slog.d(TAG,
-                                mUserIdMsg + "agentDisconnected: will handleCancel(all) for token:"
-                                        + Integer.toHexString(token));
-                    }
-                    mUserBackupManagerService.handleCancel(token, true /* cancelAll */);
-                }
-            };
+            Runnable cancellationRunnable =
+                    () -> {
+                        // On handleCancel(), the operation will call unbindAgent() which will make
+                        // sure the app doesn't get stuck in restricted mode.
+                        for (int token : mOperationStorage.operationTokensForPackage(packageName)) {
+                            if (DEBUG) {
+                                Slog.d(
+                                        TAG,
+                                        mUserIdMsg
+                                                + "agentDisconnected: cancelling for token:"
+                                                + Integer.toHexString(token));
+                            }
+                            mUserBackupManagerService.handleCancel(
+                                    token, CancellationReason.AGENT_DISCONNECTED);
+                        }
+                    };
             getThreadForCancellation(cancellationRunnable).start();
 
             mAgentConnectLock.notifyAll();
diff --git a/services/backup/java/com/android/server/backup/BackupRestoreTask.java b/services/backup/java/com/android/server/backup/BackupRestoreTask.java
index acaab0c54191..7ec5f0d786ed 100644
--- a/services/backup/java/com/android/server/backup/BackupRestoreTask.java
+++ b/services/backup/java/com/android/server/backup/BackupRestoreTask.java
@@ -16,9 +16,12 @@
 
 package com.android.server.backup;
 
-/**
- * Interface and methods used by the asynchronous-with-timeout backup/restore operations.
- */
+import android.annotation.IntDef;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/** Interface and methods used by the asynchronous-with-timeout backup/restore operations. */
 public interface BackupRestoreTask {
 
     // Execute one tick of whatever state machine the task implements
@@ -27,6 +30,24 @@ public interface BackupRestoreTask {
     // An operation that wanted a callback has completed
     void operationComplete(long result);
 
-    // An operation that wanted a callback has timed out
-    void handleCancel(boolean cancelAll);
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef({
+        CancellationReason.TIMEOUT,
+        CancellationReason.AGENT_DISCONNECTED,
+        CancellationReason.EXTERNAL,
+        CancellationReason.SCHEDULED_JOB_STOPPED,
+    })
+    @interface CancellationReason {
+        // The task timed out.
+        int TIMEOUT = 0;
+        // The agent went away before the task was able to finish (e.g. due to an app crash).
+        int AGENT_DISCONNECTED = 1;
+        // An external caller cancelled the operation (e.g. via BackupManager#cancelBackups).
+        int EXTERNAL = 2;
+        // The job scheduler has stopped an ongoing scheduled backup pass.
+        int SCHEDULED_JOB_STOPPED = 3;
+    }
+
+    /** The task is cancelled for the given {@link CancellationReason}. */
+    void handleCancel(@CancellationReason int cancellationReason);
 }
diff --git a/services/backup/java/com/android/server/backup/UserBackupManagerService.java b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
index 2143aaaa4cd6..b3af444ff9bd 100644
--- a/services/backup/java/com/android/server/backup/UserBackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
@@ -102,6 +102,7 @@ import com.android.internal.util.Preconditions;
 import com.android.server.AppWidgetBackupBridge;
 import com.android.server.EventLogTags;
 import com.android.server.LocalServices;
+import com.android.server.backup.BackupRestoreTask.CancellationReason;
 import com.android.server.backup.OperationStorage.OpState;
 import com.android.server.backup.OperationStorage.OpType;
 import com.android.server.backup.fullbackup.FullBackupEntry;
@@ -168,6 +169,7 @@ import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.IntConsumer;
 
 /** System service that performs backup/restore operations. */
 public class UserBackupManagerService {
@@ -1816,11 +1818,9 @@ public class UserBackupManagerService {
 
             for (Integer token : operationsToCancel) {
                 mOperationStorage.cancelOperation(
-                        token, /* cancelAll */
-                        true,
-                        operationType -> {
-                            /* no callback needed here */
-                        });
+                        token,
+                        operationType -> {}, // no callback needed here
+                        CancellationReason.EXTERNAL);
             }
             // We don't want the backup jobs to kick in any time soon.
             // Reschedules them to run in the distant future.
@@ -1897,19 +1897,17 @@ public class UserBackupManagerService {
     }
 
     /** Cancel the operation associated with {@code token}. */
-    public void handleCancel(int token, boolean cancelAll) {
+    public void handleCancel(int token, @CancellationReason int cancellationReason) {
         // Remove all pending timeout messages of types OpType.BACKUP_WAIT and
         // OpType.RESTORE_WAIT. On the other hand, OP_TYPE_BACKUP cannot time out and
         // doesn't require cancellation.
-        mOperationStorage.cancelOperation(
-                token,
-                cancelAll,
-                operationType -> {
-                    if (operationType == OpType.BACKUP_WAIT
-                            || operationType == OpType.RESTORE_WAIT) {
-                        mBackupHandler.removeMessages(getMessageIdForOperationType(operationType));
+        IntConsumer timeoutCallback =
+                opType -> {
+                    if (opType == OpType.BACKUP_WAIT || opType == OpType.RESTORE_WAIT) {
+                        mBackupHandler.removeMessages(getMessageIdForOperationType(opType));
                     }
-                });
+                };
+        mOperationStorage.cancelOperation(token, timeoutCallback, cancellationReason);
     }
 
     /** Returns {@code true} if a backup is currently running, else returns {@code false}. */
@@ -2219,20 +2217,17 @@ public class UserBackupManagerService {
         // offload the mRunningFullBackupTask.handleCancel() call to another thread,
         // as we might have to wait for mCancelLock
         Runnable endFullBackupRunnable =
-                new Runnable() {
-                    @Override
-                    public void run() {
-                        PerformFullTransportBackupTask pftbt = null;
-                        synchronized (mQueueLock) {
-                            if (mRunningFullBackupTask != null) {
-                                pftbt = mRunningFullBackupTask;
-                            }
-                        }
-                        if (pftbt != null) {
-                            Slog.i(TAG, mLogIdMsg + "Telling running backup to stop");
-                            pftbt.handleCancel(true);
+                () -> {
+                    PerformFullTransportBackupTask pftbt = null;
+                    synchronized (mQueueLock) {
+                        if (mRunningFullBackupTask != null) {
+                            pftbt = mRunningFullBackupTask;
                         }
                     }
+                    if (pftbt != null) {
+                        Slog.i(TAG, mLogIdMsg + "Telling running backup to stop");
+                        pftbt.handleCancel(CancellationReason.SCHEDULED_JOB_STOPPED);
+                    }
                 };
         new Thread(endFullBackupRunnable, "end-full-backup").start();
     }
diff --git a/services/backup/java/com/android/server/backup/fullbackup/PerformAdbBackupTask.java b/services/backup/java/com/android/server/backup/fullbackup/PerformAdbBackupTask.java
index 0d4364e14e03..7fc9ed3e0213 100644
--- a/services/backup/java/com/android/server/backup/fullbackup/PerformAdbBackupTask.java
+++ b/services/backup/java/com/android/server/backup/fullbackup/PerformAdbBackupTask.java
@@ -484,7 +484,7 @@ public class PerformAdbBackupTask extends FullBackupTask implements BackupRestor
     }
 
     @Override
-    public void handleCancel(boolean cancelAll) {
+    public void handleCancel(@CancellationReason int cancellationReason) {
         final PackageInfo target = mCurrentTarget;
         Slog.w(TAG, "adb backup cancel of " + target);
         if (target != null) {
diff --git a/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java b/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
index c182c2618fdf..f677c9dbf4d0 100644
--- a/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
+++ b/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
@@ -162,7 +162,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
 
     // This is true when a backup operation for some package is in progress.
     private volatile boolean mIsDoingBackup;
-    private volatile boolean mCancelAll;
+    private volatile boolean mCancelled;
     private final int mCurrentOpToken;
     private final BackupAgentTimeoutParameters mAgentTimeoutParameters;
     private final BackupEligibilityRules mBackupEligibilityRules;
@@ -199,7 +199,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
 
         if (backupManagerService.isBackupOperationInProgress()) {
             Slog.d(TAG, "Skipping full backup. A backup is already in progress.");
-            mCancelAll = true;
+            mCancelled = true;
             return;
         }
 
@@ -287,25 +287,23 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
     }
 
     @Override
-    public void handleCancel(boolean cancelAll) {
+    public void handleCancel(@CancellationReason int cancellationReason) {
         synchronized (mCancelLock) {
-            // We only support 'cancelAll = true' case for this task. Cancelling of a single package
-
-            // due to timeout is handled by SinglePackageBackupRunner and
+            // This callback is only used for cancelling the entire backup operation. Cancelling of
+            // a single package due to timeout is handled by SinglePackageBackupRunner and
             // SinglePackageBackupPreflight.
-
-            if (!cancelAll) {
-                Slog.wtf(TAG, "Expected cancelAll to be true.");
+            if (cancellationReason == CancellationReason.TIMEOUT) {
+                Slog.wtf(TAG, "This task cannot time out");
             }
 
-            if (mCancelAll) {
+            if (mCancelled) {
                 Slog.d(TAG, "Ignoring duplicate cancel call.");
                 return;
             }
 
-            mCancelAll = true;
+            mCancelled = true;
             if (mIsDoingBackup) {
-                mUserBackupManagerService.handleCancel(mBackupRunnerOpToken, cancelAll);
+                mUserBackupManagerService.handleCancel(mBackupRunnerOpToken, cancellationReason);
                 try {
                     // If we're running a backup we should be connected to a transport
                     BackupTransportClient transport =
@@ -410,7 +408,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
                 int backupPackageStatus;
                 long quota = Long.MAX_VALUE;
                 synchronized (mCancelLock) {
-                    if (mCancelAll) {
+                    if (mCancelled) {
                         break;
                     }
                     backupPackageStatus = transport.performFullBackup(currentPackage,
@@ -478,7 +476,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
                             if (nRead > 0) {
                                 out.write(buffer, 0, nRead);
                                 synchronized (mCancelLock) {
-                                    if (!mCancelAll) {
+                                    if (!mCancelled) {
                                         backupPackageStatus = transport.sendBackupData(nRead);
                                     }
                                 }
@@ -509,7 +507,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
                     synchronized (mCancelLock) {
                         mIsDoingBackup = false;
                         // If mCancelCurrent is true, we have already called cancelFullBackup().
-                        if (!mCancelAll) {
+                        if (!mCancelled) {
                             if (backupRunnerResult == BackupTransport.TRANSPORT_OK) {
                                 // If we were otherwise in a good state, now interpret the final
                                 // result based on what finishBackup() returns.  If we're in a
@@ -607,7 +605,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
                             .sendBackupOnPackageResult(mBackupObserver, packageName,
                                     BackupManager.ERROR_BACKUP_CANCELLED);
                     Slog.w(TAG, "Backup cancelled. package=" + packageName +
-                            ", cancelAll=" + mCancelAll);
+                            ", entire session cancelled=" + mCancelled);
                     EventLog.writeEvent(EventLogTags.FULL_BACKUP_CANCELLED, packageName);
                     mUserBackupManagerService.getBackupAgentConnectionManager().unbindAgent(
                             currentPackage.applicationInfo, /* allowKill= */ true);
@@ -654,7 +652,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
 
         } finally {
 
-            if (mCancelAll) {
+            if (mCancelled) {
                 backupRunStatus = BackupManager.ERROR_BACKUP_CANCELLED;
             }
 
@@ -820,7 +818,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
         }
 
         @Override
-        public void handleCancel(boolean cancelAll) {
+        public void handleCancel(@CancellationReason int cancellationReason) {
             if (DEBUG) {
                 Slog.i(TAG, "Preflight cancelled; failing");
             }
@@ -974,7 +972,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
         public void operationComplete(long result) { /* intentionally empty */ }
 
         @Override
-        public void handleCancel(boolean cancelAll) {
+        public void handleCancel(@CancellationReason int cancellationReason) {
             Slog.w(TAG, "Full backup cancel of " + mTarget.packageName);
 
             mBackupManagerMonitorEventSender.monitorEvent(
@@ -984,7 +982,7 @@ public class PerformFullTransportBackupTask extends FullBackupTask implements Ba
                     /* extras= */ null);
             mIsCancelled = true;
             // Cancel tasks spun off by this task.
-            mUserBackupManagerService.handleCancel(mEphemeralToken, cancelAll);
+            mUserBackupManagerService.handleCancel(mEphemeralToken, cancellationReason);
             mUserBackupManagerService.getBackupAgentConnectionManager().unbindAgent(
                     mTarget.applicationInfo, /* allowKill= */ true);
             // Free up everyone waiting on this task and its children.
diff --git a/services/backup/java/com/android/server/backup/internal/BackupHandler.java b/services/backup/java/com/android/server/backup/internal/BackupHandler.java
index 87cf8a313651..464dc2dfe1ec 100644
--- a/services/backup/java/com/android/server/backup/internal/BackupHandler.java
+++ b/services/backup/java/com/android/server/backup/internal/BackupHandler.java
@@ -34,6 +34,7 @@ import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.EventLogTags;
 import com.android.server.backup.BackupAgentTimeoutParameters;
 import com.android.server.backup.BackupRestoreTask;
+import com.android.server.backup.BackupRestoreTask.CancellationReason;
 import com.android.server.backup.DataChangedJournal;
 import com.android.server.backup.OperationStorage;
 import com.android.server.backup.TransportManager;
@@ -410,8 +411,8 @@ public class BackupHandler extends Handler {
 
             case MSG_BACKUP_OPERATION_TIMEOUT:
             case MSG_RESTORE_OPERATION_TIMEOUT: {
-                Slog.d(TAG, "Timeout message received for token=" + Integer.toHexString(msg.arg1));
-                backupManagerService.handleCancel(msg.arg1, false);
+                Slog.d(TAG, "Timeout for token=" + Integer.toHexString(msg.arg1));
+                backupManagerService.handleCancel(msg.arg1, CancellationReason.TIMEOUT);
                 break;
             }
 
diff --git a/services/backup/java/com/android/server/backup/internal/LifecycleOperationStorage.java b/services/backup/java/com/android/server/backup/internal/LifecycleOperationStorage.java
index 0b974e2d0a8a..5aacb2f4f007 100644
--- a/services/backup/java/com/android/server/backup/internal/LifecycleOperationStorage.java
+++ b/services/backup/java/com/android/server/backup/internal/LifecycleOperationStorage.java
@@ -24,6 +24,7 @@ import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.backup.BackupRestoreTask;
+import com.android.server.backup.BackupRestoreTask.CancellationReason;
 import com.android.server.backup.OperationStorage;
 
 import com.google.android.collect.Sets;
@@ -296,20 +297,18 @@ public class LifecycleOperationStorage implements OperationStorage {
     }
 
     /**
-     * Cancel the operation associated with {@code token}.  Cancellation may be
-     * propagated to the operation's callback (a {@link BackupRestoreTask}) if
-     * the operation has one, and the cancellation is due to the operation
-     * timing out.
+     * Cancel the operation associated with {@code token}. Cancellation may be propagated to the
+     * operation's callback (a {@link BackupRestoreTask}) if the operation has one, and the
+     * cancellation is due to the operation timing out.
      *
      * @param token the operation token specified when registering the operation
-     * @param cancelAll this is passed on when propagating the cancellation
-     * @param operationTimedOutCallback a lambda that is invoked with the
-     *                                  operation type where the operation is
-     *                                  cancelled due to timeout, allowing the
-     *                                  caller to do type-specific clean-ups.
+     * @param operationTimedOutCallback a lambda that is invoked with the operation type where the
+     *     operation is cancelled due to timeout, allowing the caller to do type-specific clean-ups.
      */
     public void cancelOperation(
-            int token, boolean cancelAll, IntConsumer operationTimedOutCallback) {
+            int token,
+            IntConsumer operationTimedOutCallback,
+            @CancellationReason int cancellationReason) {
         // Notify any synchronous waiters
         Operation op = null;
         synchronized (mOperationsLock) {
@@ -343,7 +342,7 @@ public class LifecycleOperationStorage implements OperationStorage {
             if (DEBUG) {
                 Slog.v(TAG, "[UserID:" + mUserId + "   Invoking cancel on " + op.callback);
             }
-            op.callback.handleCancel(cancelAll);
+            op.callback.handleCancel(cancellationReason);
         }
     }
 }
diff --git a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
index 494b9d59a238..8e7a23ccbb25 100644
--- a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
+++ b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
@@ -171,7 +171,7 @@ import java.util.concurrent.atomic.AtomicInteger;
  * complete backup should be performed.
  *
  * <p>This task is designed to run on a dedicated thread, with the exception of the {@link
- * #handleCancel(boolean)} method, which can be called from any thread.
+ * BackupRestoreTask#handleCancel(int)} method, which can be called from any thread.
  */
 // TODO: Stop poking into BMS state and doing things for it (e.g. synchronizing on public locks)
 // TODO: Consider having the caller responsible for some clean-up (like resetting state)
@@ -1208,13 +1208,13 @@ public class KeyValueBackupTask implements BackupRestoreTask, Runnable {
      *
      * <p>Note: This method is inherently racy since there are no guarantees about how much of the
      * task will be executed after you made the call.
-     *
-     * @param cancelAll MUST be {@code true}. Will be removed.
      */
     @Override
-    public void handleCancel(boolean cancelAll) {
+    public void handleCancel(@CancellationReason int cancellationReason) {
         // This is called in a thread different from the one that executes method run().
-        Preconditions.checkArgument(cancelAll, "Can't partially cancel a key-value backup task");
+        Preconditions.checkArgument(
+                cancellationReason != CancellationReason.TIMEOUT,
+                "Key-value backup task cannot time out");
         markCancel();
         waitCancel();
     }
diff --git a/services/backup/java/com/android/server/backup/restore/AdbRestoreFinishedLatch.java b/services/backup/java/com/android/server/backup/restore/AdbRestoreFinishedLatch.java
index cb491c6f384e..f1829b6966a8 100644
--- a/services/backup/java/com/android/server/backup/restore/AdbRestoreFinishedLatch.java
+++ b/services/backup/java/com/android/server/backup/restore/AdbRestoreFinishedLatch.java
@@ -79,7 +79,7 @@ public class AdbRestoreFinishedLatch implements BackupRestoreTask {
     }
 
     @Override
-    public void handleCancel(boolean cancelAll) {
+    public void handleCancel(@CancellationReason int cancellationReason) {
         Slog.w(TAG, "adb onRestoreFinished() timed out");
         mLatch.countDown();
         mOperationStorage.removeOperation(mCurrentOpToken);
diff --git a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
index 707ae03b3964..1263146fe405 100644
--- a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
+++ b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
@@ -1307,7 +1307,7 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
 
         // The app has timed out handling a restoring file
         @Override
-        public void handleCancel(boolean cancelAll) {
+        public void handleCancel(@CancellationReason int cancellationReason) {
             mOperationStorage.removeOperation(mEphemeralOpToken);
             Slog.w(TAG, "Full-data restore target timed out; shutting down");
             Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
@@ -1555,7 +1555,7 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
 
     // A call to agent.doRestore() or agent.doRestoreFinished() has timed out
     @Override
-    public void handleCancel(boolean cancelAll) {
+    public void handleCancel(@CancellationReason int cancellationReason) {
         mOperationStorage.removeOperation(mEphemeralOpToken);
         Slog.e(TAG, "Timeout restoring application " + mCurrentPackage.packageName);
         Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
diff --git a/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java b/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
index 2dd16f68dc56..80a3a8788d80 100644
--- a/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
+++ b/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
@@ -109,6 +109,7 @@ import com.android.server.EventLogTags;
 import com.android.server.LocalServices;
 import com.android.server.backup.BackupAgentConnectionManager;
 import com.android.server.backup.BackupRestoreTask;
+import com.android.server.backup.BackupRestoreTask.CancellationReason;
 import com.android.server.backup.BackupWakeLock;
 import com.android.server.backup.DataChangedJournal;
 import com.android.server.backup.KeyValueBackupJob;
@@ -2412,7 +2413,7 @@ public class KeyValueBackupTaskTest  {
         KeyValueBackupTask task = spy(createKeyValueBackupTask(transportMock, PACKAGE_1));
         doNothing().when(task).waitCancel();
 
-        task.handleCancel(true);
+        task.handleCancel(CancellationReason.EXTERNAL);
 
         InOrder inOrder = inOrder(task);
         inOrder.verify(task).markCancel();
@@ -2420,12 +2421,14 @@ public class KeyValueBackupTaskTest  {
     }
 
     @Test
-    public void testHandleCancel_whenCancelAllFalse_throws() throws Exception {
+    public void testHandleCancel_timeout_throws() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
         KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
-        expectThrows(IllegalArgumentException.class, () -> task.handleCancel(false));
+        expectThrows(
+                IllegalArgumentException.class,
+                () -> task.handleCancel(CancellationReason.TIMEOUT));
     }
 
     /** Do not update backup token if no data was moved. */
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java
index 8aaa72339c5b..33bd95ec9f5b 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java
@@ -51,6 +51,7 @@ import android.platform.test.flag.junit.SetFlagsRule;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.server.LocalServices;
+import com.android.server.backup.BackupRestoreTask.CancellationReason;
 import com.android.server.backup.internal.LifecycleOperationStorage;
 
 import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
@@ -368,9 +369,12 @@ public class BackupAgentConnectionManagerTest {
         mConnectionManager.agentDisconnected(TEST_PACKAGE);
 
         mTestThread.join();
-        verify(mUserBackupManagerService).handleCancel(eq(123), eq(true));
-        verify(mUserBackupManagerService).handleCancel(eq(456), eq(true));
-        verify(mUserBackupManagerService).handleCancel(eq(789), eq(true));
+        verify(mUserBackupManagerService)
+                .handleCancel(eq(123), eq(CancellationReason.AGENT_DISCONNECTED));
+        verify(mUserBackupManagerService)
+                .handleCancel(eq(456), eq(CancellationReason.AGENT_DISCONNECTED));
+        verify(mUserBackupManagerService)
+                .handleCancel(eq(789), eq(CancellationReason.AGENT_DISCONNECTED));
     }
 
     @Test