Merge "Cancel active Futures/Callbacks if transport dies."
diff --git a/services/backup/backuplib/java/com/android/server/backup/transport/BackupTransportClient.java b/services/backup/backuplib/java/com/android/server/backup/transport/BackupTransportClient.java
index 8963337..7e3ede1 100644
--- a/services/backup/backuplib/java/com/android/server/backup/transport/BackupTransportClient.java
+++ b/services/backup/backuplib/java/com/android/server/backup/transport/BackupTransportClient.java
@@ -22,20 +22,18 @@
 import android.app.backup.RestoreSet;
 import android.content.Intent;
 import android.content.pm.PackageInfo;
-import android.os.Binder;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
 import android.util.Slog;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.backup.IBackupTransport;
 import com.android.internal.infra.AndroidFuture;
 
 import java.util.ArrayDeque;
-import java.util.Deque;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Queue;
-import java.util.concurrent.ArrayBlockingQueue;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -49,17 +47,19 @@
 
     private final IBackupTransport mTransportBinder;
     private final TransportStatusCallbackPool mCallbackPool;
+    private final TransportFutures mTransportFutures;
 
     BackupTransportClient(IBackupTransport transportBinder) {
         mTransportBinder = transportBinder;
         mCallbackPool = new TransportStatusCallbackPool();
+        mTransportFutures = new TransportFutures();
     }
 
     /**
      * See {@link IBackupTransport#name()}.
      */
     public String name() throws RemoteException {
-        AndroidFuture<String> resultFuture = new AndroidFuture<>();
+        AndroidFuture<String> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.name(resultFuture);
         return getFutureResult(resultFuture);
     }
@@ -68,7 +68,7 @@
      * See {@link IBackupTransport#configurationIntent()}
      */
     public Intent configurationIntent() throws RemoteException {
-        AndroidFuture<Intent> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Intent> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.configurationIntent(resultFuture);
         return getFutureResult(resultFuture);
     }
@@ -77,7 +77,7 @@
      * See {@link IBackupTransport#currentDestinationString()}
      */
     public String currentDestinationString() throws RemoteException {
-        AndroidFuture<String> resultFuture = new AndroidFuture<>();
+        AndroidFuture<String> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.currentDestinationString(resultFuture);
         return getFutureResult(resultFuture);
     }
@@ -86,7 +86,7 @@
      * See {@link IBackupTransport#dataManagementIntent()}
      */
     public Intent dataManagementIntent() throws RemoteException {
-        AndroidFuture<Intent> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Intent> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.dataManagementIntent(resultFuture);
         return getFutureResult(resultFuture);
     }
@@ -96,7 +96,7 @@
      */
     @Nullable
     public CharSequence dataManagementIntentLabel() throws RemoteException {
-        AndroidFuture<CharSequence> resultFuture = new AndroidFuture<>();
+        AndroidFuture<CharSequence> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.dataManagementIntentLabel(resultFuture);
         return getFutureResult(resultFuture);
     }
@@ -105,7 +105,7 @@
      * See {@link IBackupTransport#transportDirName()}
      */
     public String transportDirName() throws RemoteException {
-        AndroidFuture<String> resultFuture = new AndroidFuture<>();
+        AndroidFuture<String> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.transportDirName(resultFuture);
         return getFutureResult(resultFuture);
     }
@@ -153,7 +153,7 @@
      * See {@link IBackupTransport#requestBackupTime()}
      */
     public long requestBackupTime() throws RemoteException {
-        AndroidFuture<Long> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Long> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.requestBackupTime(resultFuture);
         Long result = getFutureResult(resultFuture);
         return result == null ? BackupTransport.TRANSPORT_ERROR : result;
@@ -177,7 +177,7 @@
      * See {@link IBackupTransport#getAvailableRestoreSets()}
      */
     public RestoreSet[] getAvailableRestoreSets() throws RemoteException {
-        AndroidFuture<List<RestoreSet>> resultFuture = new AndroidFuture<>();
+        AndroidFuture<List<RestoreSet>> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.getAvailableRestoreSets(resultFuture);
         List<RestoreSet> result = getFutureResult(resultFuture);
         return result == null ? null : result.toArray(new RestoreSet[] {});
@@ -187,7 +187,7 @@
      * See {@link IBackupTransport#getCurrentRestoreSet()}
      */
     public long getCurrentRestoreSet() throws RemoteException {
-        AndroidFuture<Long> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Long> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.getCurrentRestoreSet(resultFuture);
         Long result = getFutureResult(resultFuture);
         return result == null ? BackupTransport.TRANSPORT_ERROR : result;
@@ -210,7 +210,7 @@
      * See {@link IBackupTransport#nextRestorePackage()}
      */
     public RestoreDescription nextRestorePackage() throws RemoteException {
-        AndroidFuture<RestoreDescription> resultFuture = new AndroidFuture<>();
+        AndroidFuture<RestoreDescription> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.nextRestorePackage(resultFuture);
         return getFutureResult(resultFuture);
     }
@@ -245,7 +245,7 @@
      * See {@link IBackupTransport#requestFullBackupTime()}
      */
     public long requestFullBackupTime() throws RemoteException {
-        AndroidFuture<Long> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Long> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.requestFullBackupTime(resultFuture);
         Long result = getFutureResult(resultFuture);
         return result == null ? BackupTransport.TRANSPORT_ERROR : result;
@@ -309,7 +309,7 @@
      */
     public boolean isAppEligibleForBackup(PackageInfo targetPackage, boolean isFullBackup)
             throws RemoteException {
-        AndroidFuture<Boolean> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Boolean> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.isAppEligibleForBackup(targetPackage, isFullBackup, resultFuture);
         Boolean result = getFutureResult(resultFuture);
         return result != null && result;
@@ -319,7 +319,7 @@
      * See {@link IBackupTransport#getBackupQuota(String, boolean)}
      */
     public long getBackupQuota(String packageName, boolean isFullBackup) throws RemoteException {
-        AndroidFuture<Long> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Long> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.getBackupQuota(packageName, isFullBackup, resultFuture);
         Long result = getFutureResult(resultFuture);
         return result == null ? BackupTransport.TRANSPORT_ERROR : result;
@@ -355,18 +355,58 @@
      * See {@link IBackupTransport#getTransportFlags()}
      */
     public int getTransportFlags() throws RemoteException {
-        AndroidFuture<Integer> resultFuture = new AndroidFuture<>();
+        AndroidFuture<Integer> resultFuture = mTransportFutures.newFuture();
         mTransportBinder.getTransportFlags(resultFuture);
         Integer result = getFutureResult(resultFuture);
         return result == null ? BackupTransport.TRANSPORT_ERROR : result;
     }
 
+    /**
+     * Allows the {@link TransportConnection} to notify this client
+     * if the underlying transport has become unusable.  If that happens
+     * we want to cancel all active futures or callbacks.
+     */
+    void onBecomingUnusable() {
+        mCallbackPool.cancelActiveCallbacks();
+        mTransportFutures.cancelActiveFutures();
+    }
+
     private <T> T getFutureResult(AndroidFuture<T> future) {
         try {
             return future.get(600, TimeUnit.SECONDS);
         } catch (InterruptedException | ExecutionException | TimeoutException e) {
             Slog.w(TAG, "Failed to get result from transport:", e);
             return null;
+        } finally {
+            mTransportFutures.remove(future);
+        }
+    }
+
+    private static class TransportFutures {
+        private final Object mActiveFuturesLock = new Object();
+        private final Set<AndroidFuture<?>> mActiveFutures = new HashSet<>();
+
+        <T> AndroidFuture<T> newFuture() {
+            AndroidFuture<T> future = new AndroidFuture<>();
+            synchronized (mActiveFuturesLock) {
+                mActiveFutures.add(future);
+            }
+            return future;
+        }
+
+        <T> void remove(AndroidFuture<T> future) {
+            synchronized (mActiveFuturesLock) {
+                mActiveFutures.remove(future);
+            }
+        }
+
+        void cancelActiveFutures() {
+            synchronized (mActiveFuturesLock) {
+                for (AndroidFuture<?> future : mActiveFutures) {
+                    future.cancel(true);
+                }
+                mActiveFutures.clear();
+            }
         }
     }
 
@@ -375,27 +415,47 @@
 
         private final Object mPoolLock = new Object();
         private final Queue<TransportStatusCallback> mCallbackPool = new ArrayDeque<>();
+        private final Set<TransportStatusCallback> mActiveCallbacks = new HashSet<>();
 
         TransportStatusCallback acquire() {
             synchronized (mPoolLock) {
-                if (mCallbackPool.isEmpty()) {
-                    return new TransportStatusCallback();
-                } else {
-                    return mCallbackPool.poll();
+                TransportStatusCallback callback = mCallbackPool.poll();
+                if (callback == null) {
+                    callback = new TransportStatusCallback();
                 }
+                callback.reset();
+                mActiveCallbacks.add(callback);
+                return callback;
             }
         }
 
         void recycle(TransportStatusCallback callback) {
             synchronized (mPoolLock) {
+                mActiveCallbacks.remove(callback);
                 if (mCallbackPool.size() > MAX_POOL_SIZE) {
                     Slog.d(TAG, "TransportStatusCallback pool size exceeded");
                     return;
                 }
-
-                callback.reset();
                 mCallbackPool.add(callback);
             }
         }
+
+        void cancelActiveCallbacks() {
+            synchronized (mPoolLock) {
+                for (TransportStatusCallback callback : mActiveCallbacks) {
+                    try {
+                        callback.onOperationCompleteWithStatus(BackupTransport.TRANSPORT_ERROR);
+                        // This waits for status to propagate before the callback is reset.
+                        callback.getOperationStatus();
+                    } catch (RemoteException ex) {
+                        // Nothing we can do.
+                    }
+                    if (mCallbackPool.size() < MAX_POOL_SIZE) {
+                        mCallbackPool.add(callback);
+                    }
+                }
+                mActiveCallbacks.clear();
+            }
+        }
     }
 }
diff --git a/services/backup/backuplib/java/com/android/server/backup/transport/TransportConnection.java b/services/backup/backuplib/java/com/android/server/backup/transport/TransportConnection.java
index f9a3c36..1009787 100644
--- a/services/backup/backuplib/java/com/android/server/backup/transport/TransportConnection.java
+++ b/services/backup/backuplib/java/com/android/server/backup/transport/TransportConnection.java
@@ -449,6 +449,9 @@
     private void onServiceDisconnected() {
         synchronized (mStateLock) {
             log(Priority.ERROR, "Service disconnected: client UNUSABLE");
+            if (mTransport != null) {
+                mTransport.onBecomingUnusable();
+            }
             setStateLocked(State.UNUSABLE, null);
             try {
                 // After unbindService() no calls back to mConnection
@@ -473,6 +476,9 @@
             checkStateIntegrityLocked();
 
             log(Priority.ERROR, "Binding died: client UNUSABLE");
+            if (mTransport != null) {
+                mTransport.onBecomingUnusable();
+            }
             // After unbindService() no calls back to mConnection
             switch (mState) {
                 case State.UNUSABLE:
diff --git a/services/backup/backuplib/java/com/android/server/backup/transport/TransportStatusCallback.java b/services/backup/backuplib/java/com/android/server/backup/transport/TransportStatusCallback.java
index a55178c..bc5cb02 100644
--- a/services/backup/backuplib/java/com/android/server/backup/transport/TransportStatusCallback.java
+++ b/services/backup/backuplib/java/com/android/server/backup/transport/TransportStatusCallback.java
@@ -75,13 +75,11 @@
             }
 
             Slog.w(TAG, "Couldn't get operation status from transport");
-            return BackupTransport.TRANSPORT_ERROR;
         } catch (InterruptedException e) {
             Slog.w(TAG, "Couldn't get operation status from transport: ", e);
-            return BackupTransport.TRANSPORT_ERROR;
-        } finally {
-            reset();
         }
+
+        return BackupTransport.TRANSPORT_ERROR;
     }
 
     synchronized void reset() {
diff --git a/services/tests/servicestests/src/com/android/server/backup/transport/BackupTransportClientTest.java b/services/tests/servicestests/src/com/android/server/backup/transport/BackupTransportClientTest.java
new file mode 100644
index 0000000..b154d6f
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/backup/transport/BackupTransportClientTest.java
@@ -0,0 +1,265 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.backup.transport;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.fail;
+
+import android.app.backup.BackupTransport;
+import android.app.backup.RestoreDescription;
+import android.app.backup.RestoreSet;
+import android.content.Intent;
+import android.content.pm.PackageInfo;
+import android.os.IBinder;
+import android.os.ParcelFileDescriptor;
+import android.os.RemoteException;
+import android.platform.test.annotations.Presubmit;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.backup.IBackupTransport;
+import com.android.internal.backup.ITransportStatusCallback;
+import com.android.internal.infra.AndroidFuture;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.List;
+import java.util.concurrent.CancellationException;
+
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class BackupTransportClientTest {
+
+    private static class TestFuturesFakeTransportBinder extends FakeTransportBinderBase {
+        public final Object mLock = new Object();
+
+        public String mNameCompletedImmediately;
+
+        @GuardedBy("mLock")
+        public AndroidFuture<String> mNameCompletedInFuture;
+
+        @Override public void name(AndroidFuture<String> name) throws RemoteException {
+            name.complete(mNameCompletedImmediately);
+        }
+        @Override public void transportDirName(AndroidFuture<String> name) throws RemoteException {
+            synchronized (mLock) {
+                mNameCompletedInFuture = name;
+                mLock.notifyAll();
+            }
+        }
+    }
+
+    @Test
+    public void testName_completesImmediately_returnsName() throws Exception {
+        TestFuturesFakeTransportBinder binder = new TestFuturesFakeTransportBinder();
+        binder.mNameCompletedImmediately = "fake name";
+
+        BackupTransportClient client = new BackupTransportClient(binder);
+        String name = client.name();
+
+        assertThat(name).isEqualTo("fake name");
+    }
+
+    @Test
+    public void testTransportDirName_completesLater_returnsName() throws Exception {
+        TestFuturesFakeTransportBinder binder = new TestFuturesFakeTransportBinder();
+        BackupTransportClient client = new BackupTransportClient(binder);
+
+        Thread thread = new Thread(() -> {
+            try {
+                String name = client.transportDirName();
+                assertThat(name).isEqualTo("fake name");
+            } catch (Exception ex) {
+                fail("unexpected Exception: " + ex.getClass().getCanonicalName());
+            }
+        });
+
+        thread.start();
+
+        synchronized (binder.mLock) {
+            while (binder.mNameCompletedInFuture == null) {
+                binder.mLock.wait();
+            }
+            assertThat(binder.mNameCompletedInFuture.complete("fake name")).isTrue();
+        }
+
+        thread.join();
+    }
+
+    @Test
+    public void testTransportDirName_canceledBeforeCompletion_throwsException() throws Exception {
+        TestFuturesFakeTransportBinder binder = new TestFuturesFakeTransportBinder();
+        BackupTransportClient client = new BackupTransportClient(binder);
+
+        Thread thread = new Thread(() -> {
+            try {
+                /*String name =*/ client.transportDirName();
+                fail("transportDirName should be cancelled");
+            } catch (CancellationException ex) {
+                // This is expected.
+            } catch (Exception ex) {
+                fail("unexpected Exception: " + ex.getClass().getCanonicalName());
+            }
+        });
+
+        thread.start();
+
+        synchronized (binder.mLock) {
+            while (binder.mNameCompletedInFuture == null) {
+                binder.mLock.wait();
+            }
+            client.onBecomingUnusable();
+        }
+
+        thread.join();
+    }
+
+    private static class TestCallbacksFakeTransportBinder extends FakeTransportBinderBase {
+        public final Object mLock = new Object();
+
+        public int mStatusCompletedImmediately;
+
+        @GuardedBy("mLock")
+        public ITransportStatusCallback mStatusCompletedInFuture;
+
+        @Override public void initializeDevice(ITransportStatusCallback c) throws RemoteException {
+            c.onOperationCompleteWithStatus(mStatusCompletedImmediately);
+        }
+        @Override public void finishBackup(ITransportStatusCallback c) throws RemoteException {
+            synchronized (mLock) {
+                mStatusCompletedInFuture = c;
+                mLock.notifyAll();
+            }
+        }
+    }
+
+    @Test
+    public void testInitializeDevice_completesImmediately_returnsStatus() throws Exception {
+        TestCallbacksFakeTransportBinder binder = new TestCallbacksFakeTransportBinder();
+        binder.mStatusCompletedImmediately = 123;
+
+        BackupTransportClient client = new BackupTransportClient(binder);
+        int status = client.initializeDevice();
+
+        assertThat(status).isEqualTo(123);
+    }
+
+
+    @Test
+    public void testFinishBackup_completesLater_returnsStatus() throws Exception {
+        TestCallbacksFakeTransportBinder binder = new TestCallbacksFakeTransportBinder();
+        BackupTransportClient client = new BackupTransportClient(binder);
+
+        Thread thread = new Thread(() -> {
+            try {
+                int status = client.finishBackup();
+                assertThat(status).isEqualTo(456);
+            } catch (Exception ex) {
+                fail("unexpected Exception: " + ex.getClass().getCanonicalName());
+            }
+        });
+
+        thread.start();
+
+        synchronized (binder.mLock) {
+            while (binder.mStatusCompletedInFuture == null) {
+                binder.mLock.wait();
+            }
+            binder.mStatusCompletedInFuture.onOperationCompleteWithStatus(456);
+        }
+
+        thread.join();
+    }
+
+    @Test
+    public void testFinishBackup_canceledBeforeCompletion_throwsException() throws Exception {
+        TestCallbacksFakeTransportBinder binder = new TestCallbacksFakeTransportBinder();
+        BackupTransportClient client = new BackupTransportClient(binder);
+
+        Thread thread = new Thread(() -> {
+            try {
+                int status = client.finishBackup();
+                assertThat(status).isEqualTo(BackupTransport.TRANSPORT_ERROR);
+            } catch (Exception ex) {
+                fail("unexpected Exception: " + ex.getClass().getCanonicalName());
+            }
+        });
+
+        thread.start();
+
+        synchronized (binder.mLock) {
+            while (binder.mStatusCompletedInFuture == null) {
+                binder.mLock.wait();
+            }
+            client.onBecomingUnusable();
+        }
+
+        thread.join();
+    }
+
+    // Convenience layer so we only need to fake specific methods useful for each test case.
+    private static class FakeTransportBinderBase implements IBackupTransport {
+        @Override public void name(AndroidFuture<String> f) throws RemoteException {}
+        @Override public void transportDirName(AndroidFuture<String> f) throws RemoteException {}
+        @Override public void configurationIntent(AndroidFuture<Intent> f) throws RemoteException {}
+        @Override public void currentDestinationString(AndroidFuture<String> f)
+                throws RemoteException {}
+        @Override public void dataManagementIntent(AndroidFuture<Intent> f)
+                throws RemoteException {}
+        @Override public void dataManagementIntentLabel(AndroidFuture<CharSequence> f)
+                throws RemoteException {}
+        @Override public void initializeDevice(ITransportStatusCallback c) throws RemoteException {}
+        @Override public void clearBackupData(PackageInfo i, ITransportStatusCallback c)
+                throws RemoteException {}
+        @Override public void finishBackup(ITransportStatusCallback c) throws RemoteException {}
+        @Override public void requestBackupTime(AndroidFuture<Long> f) throws RemoteException {}
+        @Override public void performBackup(PackageInfo i, ParcelFileDescriptor fd, int f,
+            ITransportStatusCallback c) throws RemoteException {}
+        @Override public void getAvailableRestoreSets(AndroidFuture<List<RestoreSet>> f)
+                throws RemoteException {}
+        @Override public void getCurrentRestoreSet(AndroidFuture<Long> f) throws RemoteException {}
+        @Override public void startRestore(long t, PackageInfo[] p, ITransportStatusCallback c)
+                throws RemoteException {}
+        @Override public void nextRestorePackage(AndroidFuture<RestoreDescription> f)
+                throws RemoteException {}
+        @Override public void getRestoreData(ParcelFileDescriptor fd, ITransportStatusCallback c)
+                throws RemoteException {}
+        @Override public void finishRestore(ITransportStatusCallback c) throws RemoteException {}
+        @Override public void requestFullBackupTime(AndroidFuture<Long> f) throws RemoteException {}
+        @Override public void performFullBackup(PackageInfo i, ParcelFileDescriptor fd, int f,
+            ITransportStatusCallback c) throws RemoteException {}
+        @Override public void checkFullBackupSize(long s, ITransportStatusCallback c)
+                throws RemoteException {}
+        @Override public void sendBackupData(int n, ITransportStatusCallback c)
+                throws RemoteException {}
+        @Override public void cancelFullBackup(ITransportStatusCallback c) throws RemoteException {}
+        @Override public void isAppEligibleForBackup(PackageInfo p, boolean b,
+            AndroidFuture<Boolean> f) throws RemoteException {}
+        @Override public void getBackupQuota(String s, boolean b, AndroidFuture<Long> f)
+                throws RemoteException {}
+        @Override public void getNextFullRestoreDataChunk(ParcelFileDescriptor fd,
+            ITransportStatusCallback c) throws RemoteException {}
+        @Override public void abortFullRestore(ITransportStatusCallback c) throws RemoteException {}
+        @Override public void getTransportFlags(AndroidFuture<Integer> f) throws RemoteException {}
+        @Override public IBinder asBinder() {
+            return null;
+        }
+    }
+}