Binding on-demand #5: PerformUnifiedRestoreTask usage

Migrate restore flow and related.

Change-Id: Ib61863e401067d7d4a9669982be8b3d87af0caa2
Ref: http://go/br-binding-on-demand
Bug: 17140907
Test: adb shell bmgr restore <set> <package> and observed logs
Test: adb shell bmgr restore <set> and observed logs
Test: Backed-up and re-installed app, observing logs
Test: gts-tradefed run commandAndExit gts-dev -m GtsBackupTestCases
Test: gts-tradefed run commandAndExit gts-dev -m GtsBackupHostTestCases
Test: cts-tradefed run commandAndExit cts-dev -m CtsBackupTestCases
Test: Looking into adding GTS/CTS for restore scenarios

services/backup/java/com/android/server/backup/RefactoredBackupManagerService.java

@@ -3233,10 +3233,10 @@ public class RefactoredBackupManagerService implements BackupManagerServiceInter
             skip = true;
         }
 
-        // Do we have a transport to fetch data for us?
-        IBackupTransport transport = mTransportManager.getCurrentTransportBinder();
-        if (transport == null) {
-            if (DEBUG) Slog.w(TAG, "No transport");
+        TransportClient transportClient =
+                mTransportManager.getCurrentTransportClient("BMS.restoreAtInstall()");
+        if (transportClient == null) {
+            if (DEBUG) Slog.w(TAG, "No transport client");
             skip = true;
         }
 
@@ -3253,16 +3253,26 @@ public class RefactoredBackupManagerService implements BackupManagerServiceInter
                 // The eventual message back into the Package Manager to run the post-install
                 // steps for 'token' will be issued from the restore handling code.
 
-                // This can throw and so *must* happen before the wakelock is acquired
-                String dirName = transport.transportDirName();
-
                 mWakelock.acquire();
+
+                OnTaskFinishedListener listener = caller -> {
+                        mTransportManager.disposeOfTransportClient(transportClient, caller);
+                        mWakelock.release();
+                };
+
                 if (MORE_DEBUG) {
                     Slog.d(TAG, "Restore at install of " + packageName);
                 }
                 Message msg = mBackupHandler.obtainMessage(MSG_RUN_RESTORE);
-                msg.obj = new RestoreParams(transport, dirName, null, null,
-                        restoreSet, packageName, token);
+                msg.obj =
+                        RestoreParams.createForRestoreAtInstall(
+                                transportClient,
+                                /* observer */ null,
+                                /* monitor */ null,
+                                restoreSet,
+                                packageName,
+                                token,
+                                listener);
                 mBackupHandler.sendMessage(msg);
             } catch (Exception e) {
                 // Calling into the transport broke; back off and proceed with the installation.
@@ -3272,8 +3282,14 @@ public class RefactoredBackupManagerService implements BackupManagerServiceInter
         }
 
         if (skip) {
-            // Auto-restore disabled or no way to attempt a restore; just tell the Package
-            // Manager to proceed with the post-install handling for this package.
+            // Auto-restore disabled or no way to attempt a restore
+
+            if (transportClient != null) {
+                mTransportManager.disposeOfTransportClient(
+                        transportClient, "BMS.restoreAtInstall()");
+            }
+
+            // Tell the PackageManager to proceed with the post-install handling for this package.
             if (DEBUG) Slog.v(TAG, "Finishing install immediately");
             try {
                 mPackageManagerBinder.finishPackageInstall(token, false);

services/backup/java/com/android/server/backup/TransportManager.java

@@ -308,6 +308,12 @@ public class TransportManager {
         }
     }
 
+    public boolean isTransportRegistered(String transportName) {
+        synchronized (mTransportLock) {
+            return getRegisteredTransportEntryLocked(transportName) != null;
+        }
+    }
+
     /**
      * Returns a {@link TransportClient} for the current transport or null if not found.
      *

services/backup/java/com/android/server/backup/internal/BackupHandler.java

@@ -189,7 +189,7 @@ public class BackupHandler extends Handler {
                     }
                     task.execute();
                 } catch (ClassCastException e) {
-                    Slog.e(TAG, "Invalid backup task in flight, obj=" + msg.obj);
+                    Slog.e(TAG, "Invalid backup/restore task in flight, obj=" + msg.obj);
                 }
                 break;
             }
@@ -229,10 +229,18 @@ public class BackupHandler extends Handler {
                 RestoreParams params = (RestoreParams) msg.obj;
                 Slog.d(TAG, "MSG_RUN_RESTORE observer=" + params.observer);
 
-                PerformUnifiedRestoreTask task = new PerformUnifiedRestoreTask(backupManagerService,
-                        params.transport,
-                        params.observer, params.monitor, params.token, params.pkgInfo,
-                        params.pmToken, params.isSystemRestore, params.filterSet);
+                PerformUnifiedRestoreTask task =
+                        new PerformUnifiedRestoreTask(
+                                backupManagerService,
+                                params.transportClient,
+                                params.observer,
+                                params.monitor,
+                                params.token,
+                                params.packageInfo,
+                                params.pmToken,
+                                params.isSystemRestore,
+                                params.filterSet,
+                                params.listener);
 
                 synchronized (backupManagerService.getPendingRestores()) {
                     if (backupManagerService.isRestoreInProgress()) {
@@ -294,8 +302,11 @@ public class BackupHandler extends Handler {
                 // Like other async operations, this is entered with the wakelock held
                 RestoreSet[] sets = null;
                 RestoreGetSetsParams params = (RestoreGetSetsParams) msg.obj;
+                String callerLogString = "BH/MSG_RUN_GET_RESTORE_SETS";
                 try {
-                    sets = params.transport.getAvailableRestoreSets();
+                    IBackupTransport transport =
+                            params.transportClient.connectOrThrow(callerLogString);
+                    sets = transport.getAvailableRestoreSets();
                     // cache the result in the active session
                     synchronized (params.session) {
                         params.session.mRestoreSets = sets;
@@ -320,7 +331,7 @@ public class BackupHandler extends Handler {
                     removeMessages(MSG_RESTORE_SESSION_TIMEOUT);
                     sendEmptyMessageDelayed(MSG_RESTORE_SESSION_TIMEOUT, TIMEOUT_RESTORE_INTERVAL);
 
-                    backupManagerService.getWakelock().release();
+                    params.listener.onFinished(callerLogString);
                 }
                 break;
             }

services/backup/java/com/android/server/backup/internal/PerformBackupTask.java

@@ -907,7 +907,7 @@ public class PerformBackupTask implements BackupRestoreTask {
             mStatus = BackupTransport.TRANSPORT_OK;
             long size = 0;
             try {
-                TransportUtils.checkTransport(transport);
+                TransportUtils.checkTransportNotNull(transport);
                 size = mBackupDataName.length();
                 if (size > 0) {
                     if (mStatus == BackupTransport.TRANSPORT_OK) {
@@ -997,7 +997,7 @@ public class PerformBackupTask implements BackupRestoreTask {
                 }
                 if (mAgentBinder != null) {
                     try {
-                        TransportUtils.checkTransport(transport);
+                        TransportUtils.checkTransportNotNull(transport);
                         long quota = transport.getBackupQuota(mCurrentPackage.packageName, false);
                         mAgentBinder.doQuotaExceeded(size, quota);
                     } catch (Exception e) {

services/backup/java/com/android/server/backup/params/RestoreGetSetsParams.java

@@ -20,20 +20,24 @@ import android.app.backup.IBackupManagerMonitor;
 import android.app.backup.IRestoreObserver;
 
 import com.android.internal.backup.IBackupTransport;
+import com.android.server.backup.internal.OnTaskFinishedListener;
 import com.android.server.backup.restore.ActiveRestoreSession;
+import com.android.server.backup.transport.TransportClient;
 
 public class RestoreGetSetsParams {
+    public final TransportClient transportClient;
+    public final ActiveRestoreSession session;
+    public final IRestoreObserver observer;
+    public final IBackupManagerMonitor monitor;
+    public final OnTaskFinishedListener listener;
 
-    public IBackupTransport transport;
-    public ActiveRestoreSession session;
-    public IRestoreObserver observer;
-    public IBackupManagerMonitor monitor;
-
-    public RestoreGetSetsParams(IBackupTransport _transport, ActiveRestoreSession _session,
-            IRestoreObserver _observer, IBackupManagerMonitor _monitor) {
-        transport = _transport;
+    public RestoreGetSetsParams(TransportClient _transportClient, ActiveRestoreSession _session,
+            IRestoreObserver _observer, IBackupManagerMonitor _monitor,
+            OnTaskFinishedListener _listener) {
+        transportClient = _transportClient;
         session = _session;
         observer = _observer;
         monitor = _monitor;
+        listener = _listener;
     }
 }

services/backup/java/com/android/server/backup/params/RestoreParams.java

@@ -20,84 +20,128 @@ import android.app.backup.IBackupManagerMonitor;
 import android.app.backup.IRestoreObserver;
 import android.content.pm.PackageInfo;
 
-import com.android.internal.backup.IBackupTransport;
+import com.android.server.backup.internal.OnTaskFinishedListener;
+import com.android.server.backup.transport.TransportClient;
 
 public class RestoreParams {
-
-    public IBackupTransport transport;
-    public String dirName;
-    public IRestoreObserver observer;
-    public IBackupManagerMonitor monitor;
-    public long token;
-    public PackageInfo pkgInfo;
-    public int pmToken; // in post-install restore, the PM's token for this transaction
-    public boolean isSystemRestore;
-    public String[] filterSet;
+    public final TransportClient transportClient;
+    public final IRestoreObserver observer;
+    public final IBackupManagerMonitor monitor;
+    public final long token;
+    public final PackageInfo packageInfo;
+    public final int pmToken; // in post-install restore, the PM's token for this transaction
+    public final boolean isSystemRestore;
+    public final String[] filterSet;
+    public final OnTaskFinishedListener listener;
 
     /**
-     * Restore a single package; no kill after restore
+     * No kill after restore.
      */
-    public RestoreParams(IBackupTransport _transport, String _dirName, IRestoreObserver _obs,
-            IBackupManagerMonitor _monitor, long _token, PackageInfo _pkg) {
-        transport = _transport;
-        dirName = _dirName;
-        observer = _obs;
-        monitor = _monitor;
-        token = _token;
-        pkgInfo = _pkg;
-        pmToken = 0;
-        isSystemRestore = false;
-        filterSet = null;
+    public static RestoreParams createForSinglePackage(
+            TransportClient transportClient,
+            IRestoreObserver observer,
+            IBackupManagerMonitor monitor,
+            long token,
+            PackageInfo packageInfo,
+            OnTaskFinishedListener listener) {
+        return new RestoreParams(
+                transportClient,
+                observer,
+                monitor,
+                token,
+                packageInfo,
+                /* pmToken */ 0,
+                /* isSystemRestore */ false,
+                /* filterSet */ null,
+                listener);
     }
 
     /**
-     * Restore at install: PM token needed, kill after restore
+     * Kill after restore.
      */
-    public RestoreParams(IBackupTransport _transport, String _dirName, IRestoreObserver _obs,
-            IBackupManagerMonitor _monitor, long _token, String _pkgName, int _pmToken) {
-        transport = _transport;
-        dirName = _dirName;
-        observer = _obs;
-        monitor = _monitor;
-        token = _token;
-        pkgInfo = null;
-        pmToken = _pmToken;
-        isSystemRestore = false;
-        filterSet = new String[]{_pkgName};
+    public static RestoreParams createForRestoreAtInstall(
+            TransportClient transportClient,
+            IRestoreObserver observer,
+            IBackupManagerMonitor monitor,
+            long token,
+            String packageName,
+            int pmToken,
+            OnTaskFinishedListener listener) {
+        String[] filterSet = {packageName};
+        return new RestoreParams(
+                transportClient,
+                observer,
+                monitor,
+                token,
+                /* packageInfo */ null,
+                pmToken,
+                /* isSystemRestore */ false,
+                filterSet,
+                listener);
     }
 
     /**
-     * Restore everything possible.  This is the form that Setup Wizard or similar
-     * restore UXes use.
+     * This is the form that Setup Wizard or similar restore UXes use.
      */
-    public RestoreParams(IBackupTransport _transport, String _dirName, IRestoreObserver _obs,
-            IBackupManagerMonitor _monitor, long _token) {
-        transport = _transport;
-        dirName = _dirName;
-        observer = _obs;
-        monitor = _monitor;
-        token = _token;
-        pkgInfo = null;
-        pmToken = 0;
-        isSystemRestore = true;
-        filterSet = null;
+    public static RestoreParams createForRestoreAll(
+            TransportClient transportClient,
+            IRestoreObserver observer,
+            IBackupManagerMonitor monitor,
+            long token,
+            OnTaskFinishedListener listener) {
+        return new RestoreParams(
+                transportClient,
+                observer,
+                monitor,
+                token,
+                /* packageInfo */ null,
+                /* pmToken */ 0,
+                /* isSystemRestore */ true,
+                /* filterSet */ null,
+                listener);
     }
 
     /**
-     * Restore some set of packages.  Leave this one up to the caller to specify
-     * whether it's to be considered a system-level restore.
+     * Caller specifies whether is considered a system-level restore.
      */
-    public RestoreParams(IBackupTransport _transport, String _dirName, IRestoreObserver _obs,
-            IBackupManagerMonitor _monitor, long _token,
-            String[] _filterSet, boolean _isSystemRestore) {
-        transport = _transport;
-        dirName = _dirName;
-        observer = _obs;
-        monitor = _monitor;
-        token = _token;
-        pkgInfo = null;
-        pmToken = 0;
-        isSystemRestore = _isSystemRestore;
-        filterSet = _filterSet;
+    public static RestoreParams createForRestoreSome(
+            TransportClient transportClient,
+            IRestoreObserver observer,
+            IBackupManagerMonitor monitor,
+            long token,
+            String[] filterSet,
+            boolean isSystemRestore,
+            OnTaskFinishedListener listener) {
+        return new RestoreParams(
+                transportClient,
+                observer,
+                monitor,
+                token,
+                /* packageInfo */ null,
+                /* pmToken */ 0,
+                isSystemRestore,
+                filterSet,
+                listener);
+    }
+
+    private RestoreParams(
+            TransportClient transportClient,
+            IRestoreObserver observer,
+            IBackupManagerMonitor monitor,
+            long token,
+            PackageInfo packageInfo,
+            int pmToken,
+            boolean isSystemRestore,
+            String[] filterSet,
+            OnTaskFinishedListener listener) {
+        this.transportClient = transportClient;
+        this.observer = observer;
+        this.monitor = monitor;
+        this.token = token;
+        this.packageInfo = packageInfo;
+        this.pmToken = pmToken;
+        this.isSystemRestore = isSystemRestore;
+        this.filterSet = filterSet;
+        this.listener = listener;
     }
 }

services/backup/java/com/android/server/backup/restore/ActiveRestoreSession.java

@@ -31,33 +31,39 @@ import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.os.Binder;
 import android.os.Message;
+import android.os.PowerManager;
 import android.util.Slog;
 
-import com.android.internal.backup.IBackupTransport;
 import com.android.server.backup.RefactoredBackupManagerService;
+import com.android.server.backup.TransportManager;
+import com.android.server.backup.internal.BackupHandler;
+import com.android.server.backup.internal.OnTaskFinishedListener;
 import com.android.server.backup.params.RestoreGetSetsParams;
 import com.android.server.backup.params.RestoreParams;
+import com.android.server.backup.transport.TransportClient;
+
+import java.util.function.BiFunction;
 
 /**
  * Restore session.
  */
 public class ActiveRestoreSession extends IRestoreSession.Stub {
-
     private static final String TAG = "RestoreSession";
 
-    private RefactoredBackupManagerService backupManagerService;
-    private String mPackageName;
-    private IBackupTransport mRestoreTransport = null;
+    private final TransportManager mTransportManager;
+    private final String mTransportName;
+    private final RefactoredBackupManagerService mBackupManagerService;
+    private final String mPackageName;
     public RestoreSet[] mRestoreSets = null;
     boolean mEnded = false;
     boolean mTimedOut = false;
 
     public ActiveRestoreSession(RefactoredBackupManagerService backupManagerService,
-            String packageName, String transport) {
-        this.backupManagerService = backupManagerService;
+            String packageName, String transportName) {
+        mBackupManagerService = backupManagerService;
         mPackageName = packageName;
-        mRestoreTransport = backupManagerService.getTransportManager().getTransportBinder(
-                transport);
+        mTransportManager = backupManagerService.getTransportManager();
+        mTransportName = transportName;
     }
 
     public void markTimedOut() {
@@ -67,7 +73,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
     // --- Binder interface ---
     public synchronized int getAvailableRestoreSets(IRestoreObserver observer,
             IBackupManagerMonitor monitor) {
-        backupManagerService.getContext().enforceCallingOrSelfPermission(
+        mBackupManagerService.getContext().enforceCallingOrSelfPermission(
                 android.Manifest.permission.BACKUP,
                 "getAvailableRestoreSets");
         if (observer == null) {
@@ -85,23 +91,32 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
 
         long oldId = Binder.clearCallingIdentity();
         try {
-            if (mRestoreTransport == null) {
-                Slog.w(TAG, "Null transport getting restore sets");
+            TransportClient transportClient =
+                    mTransportManager.getTransportClient(
+                                    mTransportName, "RestoreSession.getAvailableRestoreSets()");
+            if (transportClient == null) {
+                Slog.w(TAG, "Null transport client getting restore sets");
                 return -1;
             }
 
             // We know we're doing legit work now, so halt the timeout
             // until we're done.  It gets started again when the result
             // comes in.
-            backupManagerService.getBackupHandler().removeMessages(MSG_RESTORE_SESSION_TIMEOUT);
+            mBackupManagerService.getBackupHandler().removeMessages(MSG_RESTORE_SESSION_TIMEOUT);
 
-            // spin off the transport request to our service thread
-            backupManagerService.getWakelock().acquire();
-            Message msg = backupManagerService.getBackupHandler().obtainMessage(
+            PowerManager.WakeLock wakelock = mBackupManagerService.getWakelock();
+            wakelock.acquire();
+
+            // Prevent lambda from leaking 'this'
+            TransportManager transportManager = mTransportManager;
+            OnTaskFinishedListener listener = caller -> {
+                    transportManager.disposeOfTransportClient(transportClient, caller);
+                    wakelock.release();
+            };
+            Message msg = mBackupManagerService.getBackupHandler().obtainMessage(
                     MSG_RUN_GET_RESTORE_SETS,
-                    new RestoreGetSetsParams(mRestoreTransport, this, observer,
-                            monitor));
-            backupManagerService.getBackupHandler().sendMessage(msg);
+                    new RestoreGetSetsParams(transportClient, this, observer, monitor, listener));
+            mBackupManagerService.getBackupHandler().sendMessage(msg);
             return 0;
         } catch (Exception e) {
             Slog.e(TAG, "Error in getAvailableRestoreSets", e);
@@ -113,7 +128,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
 
     public synchronized int restoreAll(long token, IRestoreObserver observer,
             IBackupManagerMonitor monitor) {
-        backupManagerService.getContext().enforceCallingOrSelfPermission(
+        mBackupManagerService.getContext().enforceCallingOrSelfPermission(
                 android.Manifest.permission.BACKUP,
                 "performRestore");
 
@@ -131,7 +146,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
             return -1;
         }
 
-        if (mRestoreTransport == null || mRestoreSets == null) {
+        if (mRestoreSets == null) {
             Slog.e(TAG, "Ignoring restoreAll() with no restore set");
             return -1;
         }
@@ -141,37 +156,28 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
             return -1;
         }
 
-        String dirName;
-        try {
-            dirName = mRestoreTransport.transportDirName();
-        } catch (Exception e) {
-            // Transport went AWOL; fail.
-            Slog.e(TAG, "Unable to get transport dir for restore: " + e.getMessage());
+        if (!mTransportManager.isTransportRegistered(mTransportName)) {
+            Slog.e(TAG, "Transport " + mTransportName + " not registered");
             return -1;
         }
 
-        synchronized (backupManagerService.getQueueLock()) {
+        synchronized (mBackupManagerService.getQueueLock()) {
             for (int i = 0; i < mRestoreSets.length; i++) {
                 if (token == mRestoreSets[i].token) {
-                    // Real work, so stop the session timeout until we finalize the restore
-                    backupManagerService.getBackupHandler().removeMessages(
-                            MSG_RESTORE_SESSION_TIMEOUT);
-
                     long oldId = Binder.clearCallingIdentity();
                     try {
-                        backupManagerService.getWakelock().acquire();
-                        if (MORE_DEBUG) {
-                            Slog.d(TAG, "restoreAll() kicking off");
-                        }
-                        Message msg = backupManagerService.getBackupHandler().obtainMessage(
-                                MSG_RUN_RESTORE);
-                        msg.obj = new RestoreParams(mRestoreTransport, dirName,
-                                observer, monitor, token);
-                        backupManagerService.getBackupHandler().sendMessage(msg);
+                        return sendRestoreToHandlerLocked(
+                                (transportClient, listener) ->
+                                        RestoreParams.createForRestoreAll(
+                                                transportClient,
+                                                observer,
+                                                monitor,
+                                                token,
+                                                listener),
+                                "RestoreSession.restoreAll()");
                     } finally {
                         Binder.restoreCallingIdentity(oldId);
                     }
-                    return 0;
                 }
             }
         }
@@ -183,7 +189,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
     // Restores of more than a single package are treated as 'system' restores
     public synchronized int restoreSome(long token, IRestoreObserver observer,
             IBackupManagerMonitor monitor, String[] packages) {
-        backupManagerService.getContext().enforceCallingOrSelfPermission(
+        mBackupManagerService.getContext().enforceCallingOrSelfPermission(
                 android.Manifest.permission.BACKUP,
                 "performRestore");
 
@@ -227,7 +233,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
             return -1;
         }
 
-        if (mRestoreTransport == null || mRestoreSets == null) {
+        if (mRestoreSets == null) {
             Slog.e(TAG, "Ignoring restoreAll() with no restore set");
             return -1;
         }
@@ -237,34 +243,30 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
             return -1;
         }
 
-        String dirName;
-        try {
-            dirName = mRestoreTransport.transportDirName();
-        } catch (Exception e) {
-            // Transport went AWOL; fail.
-            Slog.e(TAG, "Unable to get transport name for restoreSome: " + e.getMessage());
+        if (!mTransportManager.isTransportRegistered(mTransportName)) {
+            Slog.e(TAG, "Transport " + mTransportName + " not registered");
             return -1;
         }
 
-        synchronized (backupManagerService.getQueueLock()) {
+        synchronized (mBackupManagerService.getQueueLock()) {
             for (int i = 0; i < mRestoreSets.length; i++) {
                 if (token == mRestoreSets[i].token) {
-                    // Stop the session timeout until we finalize the restore
-                    backupManagerService.getBackupHandler().removeMessages(
-                            MSG_RESTORE_SESSION_TIMEOUT);
-
                     long oldId = Binder.clearCallingIdentity();
-                    backupManagerService.getWakelock().acquire();
-                    if (MORE_DEBUG) {
-                        Slog.d(TAG, "restoreSome() of " + packages.length + " packages");
+                    try {
+                        return sendRestoreToHandlerLocked(
+                                (transportClient, listener) ->
+                                        RestoreParams.createForRestoreSome(
+                                                transportClient,
+                                                observer,
+                                                monitor,
+                                                token,
+                                                packages,
+                                                /* isSystemRestore */ packages.length > 1,
+                                                listener),
+                                "RestoreSession.restoreSome(" + packages.length + " packages)");
+                    } finally {
+                        Binder.restoreCallingIdentity(oldId);
                     }
-                    Message msg = backupManagerService.getBackupHandler().obtainMessage(
-                            MSG_RUN_RESTORE);
-                    msg.obj = new RestoreParams(mRestoreTransport, dirName, observer, monitor,
-                            token, packages, packages.length > 1);
-                    backupManagerService.getBackupHandler().sendMessage(msg);
-                    Binder.restoreCallingIdentity(oldId);
-                    return 0;
                 }
             }
         }
@@ -297,9 +299,9 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
             }
         }
 
-        PackageInfo app = null;
+        final PackageInfo app;
         try {
-            app = backupManagerService.getPackageManager().getPackageInfo(packageName, 0);
+            app = mBackupManagerService.getPackageManager().getPackageInfo(packageName, 0);
         } catch (NameNotFoundException nnf) {
             Slog.w(TAG, "Asked to restore nonexistent pkg " + packageName);
             return -1;
@@ -307,7 +309,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
 
         // If the caller is not privileged and is not coming from the target
         // app's uid, throw a permission exception back to the caller.
-        int perm = backupManagerService.getContext().checkPermission(
+        int perm = mBackupManagerService.getContext().checkPermission(
                 android.Manifest.permission.BACKUP,
                 Binder.getCallingPid(), Binder.getCallingUid());
         if ((perm == PackageManager.PERMISSION_DENIED) &&
@@ -317,12 +319,17 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
             throw new SecurityException("No permission to restore other packages");
         }
 
+        if (!mTransportManager.isTransportRegistered(mTransportName)) {
+            Slog.e(TAG, "Transport " + mTransportName + " not registered");
+            return -1;
+        }
+
         // So far so good; we're allowed to try to restore this package.
         long oldId = Binder.clearCallingIdentity();
         try {
             // Check whether there is data for it in the current dataset, falling back
             // to the ancestral dataset if not.
-            long token = backupManagerService.getAvailableRestoreToken(packageName);
+            long token = mBackupManagerService.getAvailableRestoreToken(packageName);
             if (DEBUG) {
                 Slog.v(TAG, "restorePackage pkg=" + packageName
                         + " token=" + Long.toHexString(token));
@@ -338,30 +345,53 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
                 return -1;
             }
 
-            String dirName;
-            try {
-                dirName = mRestoreTransport.transportDirName();
-            } catch (Exception e) {
-                // Transport went AWOL; fail.
-                Slog.e(TAG, "Unable to get transport dir for restorePackage: " + e.getMessage());
-                return -1;
-            }
-
-            // Stop the session timeout until we finalize the restore
-            backupManagerService.getBackupHandler().removeMessages(MSG_RESTORE_SESSION_TIMEOUT);
-
-            // Ready to go:  enqueue the restore request and claim success
-            backupManagerService.getWakelock().acquire();
-            if (MORE_DEBUG) {
-                Slog.d(TAG, "restorePackage() : " + packageName);
-            }
-            Message msg = backupManagerService.getBackupHandler().obtainMessage(MSG_RUN_RESTORE);
-            msg.obj = new RestoreParams(mRestoreTransport, dirName, observer, monitor,
-                    token, app);
-            backupManagerService.getBackupHandler().sendMessage(msg);
+            return sendRestoreToHandlerLocked(
+                    (transportClient, listener) ->
+                            RestoreParams.createForSinglePackage(
+                                    transportClient,
+                                    observer,
+                                    monitor,
+                                    token,
+                                    app,
+                                    listener),
+                    "RestoreSession.restorePackage(" + packageName + ")");
         } finally {
             Binder.restoreCallingIdentity(oldId);
         }
+    }
+
+    /**
+     * Returns 0 if operation sent or -1 otherwise.
+     */
+    private int sendRestoreToHandlerLocked(
+            BiFunction<TransportClient, OnTaskFinishedListener, RestoreParams> restoreParamsBuilder,
+            String callerLogString) {
+        TransportClient transportClient =
+                mTransportManager.getTransportClient(mTransportName, callerLogString);
+        if (transportClient == null) {
+            Slog.e(TAG, "Transport " + mTransportName + " got unregistered");
+            return -1;
+        }
+
+        // Stop the session timeout until we finalize the restore
+        BackupHandler backupHandler = mBackupManagerService.getBackupHandler();
+        backupHandler.removeMessages(MSG_RESTORE_SESSION_TIMEOUT);
+
+        PowerManager.WakeLock wakelock = mBackupManagerService.getWakelock();
+        wakelock.acquire();
+        if (MORE_DEBUG) {
+            Slog.d(TAG, callerLogString);
+        }
+
+        // Prevent lambda from leaking 'this'
+        TransportManager transportManager = mTransportManager;
+        OnTaskFinishedListener listener = caller -> {
+                transportManager.disposeOfTransportClient(transportClient, caller);
+                wakelock.release();
+        };
+        Message msg = backupHandler.obtainMessage(MSG_RUN_RESTORE);
+        msg.obj = restoreParamsBuilder.apply(transportClient, listener);
+        backupHandler.sendMessage(msg);
         return 0;
     }
 
@@ -380,7 +410,6 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
         public void run() {
             // clean up the session's bookkeeping
             synchronized (mSession) {
-                mSession.mRestoreTransport = null;
                 mSession.mEnded = true;
             }
 
@@ -404,7 +433,7 @@ public class ActiveRestoreSession extends IRestoreSession.Stub {
             throw new IllegalStateException("Restore session already ended");
         }
 
-        backupManagerService.getBackupHandler().post(
-                new EndRestoreRunnable(backupManagerService, this));
+        mBackupManagerService.getBackupHandler().post(
+                new EndRestoreRunnable(mBackupManagerService, this));
     }
 }

services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java

@@ -61,6 +61,8 @@ import com.android.server.backup.BackupUtils;
 import com.android.server.backup.PackageManagerBackupAgent;
 import com.android.server.backup.PackageManagerBackupAgent.Metadata;
 import com.android.server.backup.RefactoredBackupManagerService;
+import com.android.server.backup.internal.OnTaskFinishedListener;
+import com.android.server.backup.transport.TransportClient;
 import com.android.server.backup.utils.AppBackupUtils;
 import com.android.server.backup.utils.BackupManagerMonitorUtils;
 
@@ -76,8 +78,8 @@ import java.util.List;
 public class PerformUnifiedRestoreTask implements BackupRestoreTask {
 
     private RefactoredBackupManagerService backupManagerService;
-    // Transport we're working with to do the restore
-    private IBackupTransport mTransport;
+    // Transport client we're working with to do the restore
+    private final TransportClient mTransportClient;
 
     // Where per-transport saved state goes
     File mStateDir;
@@ -141,6 +143,9 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
     // Done?
     private boolean mFinished;
 
+    // When finished call listener
+    private final OnTaskFinishedListener mListener;
+
     // Key/value: bookkeeping about staged data and files for agent access
     private File mBackupDataName;
     private File mStageName;
@@ -154,15 +159,16 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
     // Invariant: mWakelock is already held, and this task is responsible for
     // releasing it at the end of the restore operation.
     public PerformUnifiedRestoreTask(RefactoredBackupManagerService backupManagerService,
-            IBackupTransport transport, IRestoreObserver observer,
+            TransportClient transportClient, IRestoreObserver observer,
             IBackupManagerMonitor monitor, long restoreSetToken, PackageInfo targetPackage,
-            int pmToken, boolean isFullSystemRestore, String[] filterSet) {
+            int pmToken, boolean isFullSystemRestore, String[] filterSet,
+            OnTaskFinishedListener listener) {
         this.backupManagerService = backupManagerService;
         mEphemeralOpToken = backupManagerService.generateRandomIntegerToken();
         mState = UnifiedRestoreState.INITIAL;
         mStartRealtime = SystemClock.elapsedRealtime();
 
-        mTransport = transport;
+        mTransportClient = transportClient;
         mObserver = observer;
         mMonitor = monitor;
         mToken = restoreSetToken;
@@ -171,6 +177,7 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
         mIsSystemRestore = isFullSystemRestore;
         mFinished = false;
         mDidLaunch = false;
+        mListener = listener;
 
         if (targetPackage != null) {
             // Single package restore
@@ -342,7 +349,7 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
         }
 
         try {
-            String transportDir = mTransport.transportDirName();
+            String transportDir = mTransportClient.getTransportDirName();
             mStateDir = new File(backupManagerService.getBaseStateDir(), transportDir);
 
             // Fetch the current metadata from the dataset first
@@ -351,7 +358,11 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
             mAcceptSet.add(0, pmPackage);
 
             PackageInfo[] packages = mAcceptSet.toArray(new PackageInfo[0]);
-            mStatus = mTransport.startRestore(mToken, packages);
+
+            IBackupTransport transport =
+                    mTransportClient.connectOrThrow("PerformUnifiedRestoreTask.startRestore()");
+
+            mStatus = transport.startRestore(mToken, packages);
             if (mStatus != BackupTransport.TRANSPORT_OK) {
                 Slog.e(TAG, "Transport error " + mStatus + "; no restore possible");
                 mStatus = BackupTransport.TRANSPORT_ERROR;
@@ -359,7 +370,7 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
                 return;
             }
 
-            RestoreDescription desc = mTransport.nextRestorePackage();
+            RestoreDescription desc = transport.nextRestorePackage();
             if (desc == null) {
                 Slog.e(TAG, "No restore metadata available; halting");
                 mMonitor = BackupManagerMonitorUtils.monitorEvent(mMonitor,
@@ -444,7 +455,10 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
     private void dispatchNextRestore() {
         UnifiedRestoreState nextState = UnifiedRestoreState.FINAL;
         try {
-            mRestoreDescription = mTransport.nextRestorePackage();
+            IBackupTransport transport =
+                    mTransportClient.connectOrThrow(
+                            "PerformUnifiedRestoreTask.dispatchNextRestore()");
+            mRestoreDescription = transport.nextRestorePackage();
             final String pkgName = (mRestoreDescription != null)
                     ? mRestoreDescription.getPackageName() : null;
             if (pkgName == null) {
@@ -657,13 +671,17 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
         File downloadFile = (staging) ? mStageName : mBackupDataName;
 
         try {
+            IBackupTransport transport =
+                    mTransportClient.connectOrThrow(
+                            "PerformUnifiedRestoreTask.initiateOneRestore()");
+
             // Run the transport's restore pass
             stage = ParcelFileDescriptor.open(downloadFile,
                     ParcelFileDescriptor.MODE_READ_WRITE |
                             ParcelFileDescriptor.MODE_CREATE |
                             ParcelFileDescriptor.MODE_TRUNCATE);
 
-            if (mTransport.getRestoreData(stage) != BackupTransport.TRANSPORT_OK) {
+            if (transport.getRestoreData(stage) != BackupTransport.TRANSPORT_OK) {
                 // Transport-level failure, so we wind everything up and
                 // terminate the restore operation.
                 Slog.e(TAG, "Error getting restore data for " + packageName);
@@ -750,7 +768,7 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
         // None of this can run on the work looper here, so we spin asynchronous
         // work like this:
         //
-        //   StreamFeederThread: read data from mTransport.getNextFullRestoreDataChunk()
+        //   StreamFeederThread: read data from transport.getNextFullRestoreDataChunk()
         //                       write it into the pipe to the engine
         //   EngineThread: FullRestoreEngine thread communicating with the target app
         //
@@ -844,10 +862,12 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
             // spin up the engine and start moving data to it
             new Thread(mEngineThread, "unified-restore-engine").start();
 
+            String callerLogString = "PerformUnifiedRestoreTask$StreamFeederThread.run()";
             try {
+                IBackupTransport transport = mTransportClient.connectOrThrow(callerLogString);
                 while (status == BackupTransport.TRANSPORT_OK) {
                     // have the transport write some of the restoring data to us
-                    int result = mTransport.getNextFullRestoreDataChunk(tWriteEnd);
+                    int result = transport.getNextFullRestoreDataChunk(tWriteEnd);
                     if (result > 0) {
                         // The transport wrote this many bytes of restore data to the
                         // pipe, so pass it along to the engine.
@@ -936,7 +956,9 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
                     // Something went wrong somewhere.  Whether it was at the transport
                     // level is immaterial; we need to tell the transport to bail
                     try {
-                        mTransport.abortFullRestore();
+                        IBackupTransport transport =
+                                mTransportClient.connectOrThrow(callerLogString);
+                        transport.abortFullRestore();
                     } catch (Exception e) {
                         // transport itself is dead; make sure we handle this as a
                         // fatal error
@@ -1039,8 +1061,11 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
             Slog.d(TAG, "finishing restore mObserver=" + mObserver);
         }
 
+        String callerLogString = "PerformUnifiedRestoreTask.finalizeRestore()";
         try {
-            mTransport.finishRestore();
+            IBackupTransport transport =
+                    mTransportClient.connectOrThrow(callerLogString);
+            transport.finishRestore();
         } catch (Exception e) {
             Slog.e(TAG, "Error finishing restore", e);
         }
@@ -1087,9 +1112,6 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
             backupManagerService.writeRestoreTokens();
         }
 
-        // done; we can finally release the wakelock and be legitimately done.
-        Slog.i(TAG, "Restore complete.");
-
         synchronized (backupManagerService.getPendingRestores()) {
             if (backupManagerService.getPendingRestores().size() > 0) {
                 if (DEBUG) {
@@ -1108,7 +1130,8 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
             }
         }
 
-        backupManagerService.getWakelock().release();
+        Slog.i(TAG, "Restore complete.");
+        mListener.onFinished(callerLogString);
     }
 
     void keyValueAgentErrorCleanup() {
@@ -1301,13 +1324,11 @@ public class PerformUnifiedRestoreTask implements BackupRestoreTask {
 
     void sendOnRestorePackage(String name) {
         if (mObserver != null) {
-            if (mObserver != null) {
-                try {
-                    mObserver.onUpdate(mCount, name);
-                } catch (RemoteException e) {
-                    Slog.d(TAG, "Restore observer died in onUpdate");
-                    mObserver = null;
-                }
+            try {
+                mObserver.onUpdate(mCount, name);
+            } catch (RemoteException e) {
+                Slog.d(TAG, "Restore observer died in onUpdate");
+                mObserver = null;
             }
         }
     }

services/backup/java/com/android/server/backup/restore/RestoreEngine.java

@@ -30,8 +30,8 @@ public abstract class RestoreEngine {
     public static final int TARGET_FAILURE = -2;
     public static final int TRANSPORT_FAILURE = -3;
 
-    private AtomicBoolean mRunning = new AtomicBoolean(false);
-    private AtomicInteger mResult = new AtomicInteger(SUCCESS);
+    private final AtomicBoolean mRunning = new AtomicBoolean(false);
+    private final AtomicInteger mResult = new AtomicInteger(SUCCESS);
 
     public boolean isRunning() {
         return mRunning.get();

services/backup/java/com/android/server/backup/transport/TransportUtils.java

@@ -31,7 +31,7 @@ public class TransportUtils {
      * Throws {@link TransportNotAvailableException} if {@param transport} is null. The semantics is
      * similar to a {@link DeadObjectException} coming from a dead transport binder.
      */
-    public static IBackupTransport checkTransport(@Nullable IBackupTransport transport)
+    public static IBackupTransport checkTransportNotNull(@Nullable IBackupTransport transport)
             throws TransportNotAvailableException {
         if (transport == null) {
             log(Log.ERROR, TAG, "Transport not available");

services/robotests/src/com/android/server/backup/TransportManagerTest.java

@@ -532,6 +532,18 @@ public class TransportManagerTest {
         assertThat(transportName).isEqualTo("newName");
     }
 
+    @Test
+    public void isTransportRegistered_returnsCorrectly() throws Exception {
+        TransportManager transportManager =
+                createTransportManagerAndSetUpTransports(
+                        Collections.singletonList(mTransport1),
+                        Collections.singletonList(mTransport2),
+                        mTransport1.name);
+
+        assertThat(transportManager.isTransportRegistered(mTransport1.name)).isTrue();
+        assertThat(transportManager.isTransportRegistered(mTransport2.name)).isFalse();
+    }
+
     private void setUpPackageWithTransports(String packageName, List<TransportInfo> transports,
             int flags) throws Exception {
         PackageInfo packageInfo = new PackageInfo();