Merge changes from topic "BackportUiAutomatorRetry"

* changes:
  Let #getUiAutomation return null if UiAutomation fails to connect (3/n)
  Allow #disconnect to be called safely on connection timeout (2/n)
  Add #connectWithTimeout (1/n)

2 files changed, 126 insertions, 36 deletions

diff --git a/core/java/android/app/Instrumentation.java b/core/java/android/app/Instrumentation.java
index 721525d9af9d..2ef147b3e903 100644
--- a/core/java/android/app/Instrumentation.java
+++ b/core/java/android/app/Instrumentation.java
@@ -62,6 +62,7 @@ import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Base class for implementing application instrumentation code.  When running
@@ -90,6 +91,8 @@ public class Instrumentation {
 
     private static final String TAG = "Instrumentation";
 
+    private static final long CONNECT_TIMEOUT_MILLIS = 5000;
+
     /**
      * @hide
      */
@@ -2125,6 +2128,13 @@ public class Instrumentation {
      * Equivalent to {@code getUiAutomation(0)}. If a {@link UiAutomation} exists with different
      * flags, the flags on that instance will be changed, and then it will be returned.
      * </p>
+     * <p>
+     * Compatibility mode: This method is infallible for apps targeted for
+     * {@link Build.VERSION_CODES#R} and earlier versions; for apps targeted for later versions, it
+     * will return null if {@link UiAutomation} fails to connect. The caller can check the return
+     * value and retry on error.
+     * </p>
+     *
      * @return The UI automation instance.
      *
      * @see UiAutomation
@@ -2152,6 +2162,12 @@ public class Instrumentation {
      * If a {@link UiAutomation} exists with different flags, the flags on that instance will be
      * changed, and then it will be returned.
      * </p>
+     * <p>
+     * Compatibility mode: This method is infallible for apps targeted for
+     * {@link Build.VERSION_CODES#R} and earlier versions; for apps targeted for later versions, it
+     * will return null if {@link UiAutomation} fails to connect. The caller can check the return
+     * value and retry on error.
+     * </p>
      *
      * @param flags The flags to be passed to the UiAutomation, for example
      *        {@link UiAutomation#FLAG_DONT_SUPPRESS_ACCESSIBILITY_SERVICES}.
@@ -2173,8 +2189,17 @@ public class Instrumentation {
             } else {
                 mUiAutomation.disconnect();
             }
-            mUiAutomation.connect(flags);
-            return mUiAutomation;
+            if (getTargetContext().getApplicationInfo().targetSdkVersion <= Build.VERSION_CODES.R) {
+                mUiAutomation.connect(flags);
+                return mUiAutomation;
+            }
+            try {
+                mUiAutomation.connectWithTimeout(flags, CONNECT_TIMEOUT_MILLIS);
+                return mUiAutomation;
+            } catch (TimeoutException e) {
+                mUiAutomation.destroy();
+                mUiAutomation = null;
+            }
         }
         return null;
     }
diff --git a/core/java/android/app/UiAutomation.java b/core/java/android/app/UiAutomation.java
index a9a06dabc049..e0951bf3f4d2 100644
--- a/core/java/android/app/UiAutomation.java
+++ b/core/java/android/app/UiAutomation.java
@@ -22,6 +22,7 @@ import android.accessibilityservice.AccessibilityService.IAccessibilityServiceCl
 import android.accessibilityservice.AccessibilityServiceInfo;
 import android.accessibilityservice.IAccessibilityServiceClient;
 import android.accessibilityservice.IAccessibilityServiceConnection;
+import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.TestApi;
@@ -60,6 +61,8 @@ import com.android.internal.util.function.pooled.PooledLambda;
 import libcore.io.IoUtils;
 
 import java.io.IOException;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeoutException;
@@ -116,6 +119,28 @@ public final class UiAutomation {
     /** Rotation constant: Freeze rotation to 270 degrees . */
     public static final int ROTATION_FREEZE_270 = Surface.ROTATION_270;
 
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef(value = {
+            ConnectionState.DISCONNECTED,
+            ConnectionState.CONNECTING,
+            ConnectionState.CONNECTED,
+            ConnectionState.FAILED
+    })
+    private @interface ConnectionState {
+        /** The initial state before {@link #connect} or after {@link #disconnect} is called. */
+        int DISCONNECTED = 0;
+        /**
+         * The temporary state after {@link #connect} is called. Will transition to
+         * {@link #CONNECTED} or {@link #FAILED} depending on whether {@link #connect} succeeds or
+         * not.
+         */
+        int CONNECTING = 1;
+        /** The state when {@link #connect} has succeeded. */
+        int CONNECTED = 2;
+        /** The state when {@link #connect} has failed. */
+        int FAILED = 3;
+    }
+
     /**
      * UiAutomation supresses accessibility services by default. This flag specifies that
      * existing accessibility services should continue to run, and that new ones may start.
@@ -144,12 +169,14 @@ public final class UiAutomation {
 
     private long mLastEventTimeMillis;
 
-    private boolean mIsConnecting;
+    private @ConnectionState int mConnectionState = ConnectionState.DISCONNECTED;
 
     private boolean mIsDestroyed;
 
     private int mFlags;
 
+    private int mGenerationId = 0;
+
     /**
      * Listener for observing the {@link AccessibilityEvent} stream.
      */
@@ -210,32 +237,55 @@ public final class UiAutomation {
     }
 
     /**
-     * Connects this UiAutomation to the accessibility introspection APIs with default flags.
+     * Connects this UiAutomation to the accessibility introspection APIs with default flags
+     * and default timeout.
      *
      * @hide
      */
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.P, trackingBug = 115609023)
     public void connect() {
-        connect(0);
+        try {
+            connectWithTimeout(0, CONNECT_TIMEOUT_MILLIS);
+        } catch (TimeoutException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    /**
+     * Connects this UiAutomation to the accessibility introspection APIs with default timeout.
+     *
+     * @hide
+     */
+    public void connect(int flags) {
+        try {
+            connectWithTimeout(flags, CONNECT_TIMEOUT_MILLIS);
+        } catch (TimeoutException e) {
+            throw new RuntimeException(e);
+        }
     }
 
     /**
      * Connects this UiAutomation to the accessibility introspection APIs.
      *
      * @param flags Any flags to apply to the automation as it gets connected
+     * @param timeoutMillis The wait timeout in milliseconds
+     *
+     * @throws TimeoutException If not connected within the timeout
      *
      * @hide
      */
-    public void connect(int flags) {
+    public void connectWithTimeout(int flags, long timeoutMillis) throws TimeoutException {
         synchronized (mLock) {
             throwIfConnectedLocked();
-            if (mIsConnecting) {
+            if (mConnectionState == ConnectionState.CONNECTING) {
                 return;
             }
-            mIsConnecting = true;
+            mConnectionState = ConnectionState.CONNECTING;
             mRemoteCallbackThread = new HandlerThread("UiAutomation");
             mRemoteCallbackThread.start();
-            mClient = new IAccessibilityServiceClientImpl(mRemoteCallbackThread.getLooper());
+            // Increment the generation since we are about to interact with a new client
+            mClient = new IAccessibilityServiceClientImpl(
+                    mRemoteCallbackThread.getLooper(), ++mGenerationId);
         }
 
         try {
@@ -248,24 +298,21 @@ public final class UiAutomation {
 
         synchronized (mLock) {
             final long startTimeMillis = SystemClock.uptimeMillis();
-            try {
-                while (true) {
-                    if (isConnectedLocked()) {
-                        break;
-                    }
-                    final long elapsedTimeMillis = SystemClock.uptimeMillis() - startTimeMillis;
-                    final long remainingTimeMillis = CONNECT_TIMEOUT_MILLIS - elapsedTimeMillis;
-                    if (remainingTimeMillis <= 0) {
-                        throw new RuntimeException("Error while connecting " + this);
-                    }
-                    try {
-                        mLock.wait(remainingTimeMillis);
-                    } catch (InterruptedException ie) {
-                        /* ignore */
-                    }
+            while (true) {
+                if (mConnectionState == ConnectionState.CONNECTED) {
+                    break;
+                }
+                final long elapsedTimeMillis = SystemClock.uptimeMillis() - startTimeMillis;
+                final long remainingTimeMillis = timeoutMillis - elapsedTimeMillis;
+                if (remainingTimeMillis <= 0) {
+                    mConnectionState = ConnectionState.FAILED;
+                    throw new TimeoutException("Timeout while connecting " + this);
+                }
+                try {
+                    mLock.wait(remainingTimeMillis);
+                } catch (InterruptedException ie) {
+                    /* ignore */
                 }
-            } finally {
-                mIsConnecting = false;
             }
         }
     }
@@ -289,12 +336,17 @@ public final class UiAutomation {
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.P, trackingBug = 115609023)
     public void disconnect() {
         synchronized (mLock) {
-            if (mIsConnecting) {
+            if (mConnectionState == ConnectionState.CONNECTING) {
                 throw new IllegalStateException(
                         "Cannot call disconnect() while connecting " + this);
             }
-            throwIfNotConnectedLocked();
+            if (mConnectionState == ConnectionState.DISCONNECTED) {
+                return;
+            }
+            mConnectionState = ConnectionState.DISCONNECTED;
             mConnectionId = CONNECTION_ID_UNDEFINED;
+            // Increment the generation so we no longer interact with the existing client
+            ++mGenerationId;
         }
         try {
             // Calling out without a lock held.
@@ -1224,18 +1276,14 @@ public final class UiAutomation {
         return stringBuilder.toString();
     }
 
-    private boolean isConnectedLocked() {
-        return mConnectionId != CONNECTION_ID_UNDEFINED;
-    }
-
     private void throwIfConnectedLocked() {
-        if (mConnectionId != CONNECTION_ID_UNDEFINED) {
-            throw new IllegalStateException("UiAutomation not connected, " + this);
+        if (mConnectionState == ConnectionState.CONNECTED) {
+            throw new IllegalStateException("UiAutomation connected, " + this);
         }
     }
 
     private void throwIfNotConnectedLocked() {
-        if (!isConnectedLocked()) {
+        if (mConnectionState != ConnectionState.CONNECTED) {
             throw new IllegalStateException("UiAutomation not connected, " + this);
         }
     }
@@ -1252,11 +1300,25 @@ public final class UiAutomation {
 
     private class IAccessibilityServiceClientImpl extends IAccessibilityServiceClientWrapper {
 
-        public IAccessibilityServiceClientImpl(Looper looper) {
+        public IAccessibilityServiceClientImpl(Looper looper, int generationId) {
             super(null, looper, new Callbacks() {
+                private final int mGenerationId = generationId;
+                /**
+                 * True if UiAutomation doesn't interact with this client anymore.
+                 * Used by methods below to stop sending notifications or changing members
+                 * of {@link UiAutomation}.
+                 */
+                private boolean isGenerationChangedLocked() {
+                    return mGenerationId != UiAutomation.this.mGenerationId;
+                }
+
                 @Override
                 public void init(int connectionId, IBinder windowToken) {
                     synchronized (mLock) {
+                        if (isGenerationChangedLocked()) {
+                            return;
+                        }
+                        mConnectionState = ConnectionState.CONNECTED;
                         mConnectionId = connectionId;
                         mLock.notifyAll();
                     }
@@ -1290,6 +1352,9 @@ public final class UiAutomation {
                 public void onAccessibilityEvent(AccessibilityEvent event) {
                     final OnAccessibilityEventListener listener;
                     synchronized (mLock) {
+                        if (isGenerationChangedLocked()) {
+                            return;
+                        }
                         mLastEventTimeMillis = event.getEventTime();
                         if (mWaitingForEventDelivery) {
                             mEventQueue.add(AccessibilityEvent.obtain(event));