Ensure dream logic is not suspended on handler.

It is possible for work to be queued on a handler and then
suspended in doze mode. This can cause dream operations to
be noticeably deferred. In order to address this, a wake
lock is now acquired in this scenario and wraps these
calls on the handler. Note that in the normal, non-doze
dream case, the wakelock is not necessary.

Bug: 404420534
Flag: EXEMPT bugfix
Test: atest DreamServiceTest
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:d958722427762581b9a39b3518d21334e063d99b)
Merged-In: Ib7802e951ea793ca4484ef19972a9f2aed4bfd81
Change-Id: Ib7802e951ea793ca4484ef19972a9f2aed4bfd81

2 files changed, 106 insertions, 25 deletions

diff --git a/core/java/android/service/dreams/DreamService.java b/core/java/android/service/dreams/DreamService.java
index 50a646284a8e..31143a559060 100644
--- a/core/java/android/service/dreams/DreamService.java
+++ b/core/java/android/service/dreams/DreamService.java
@@ -21,6 +21,7 @@ import static android.service.dreams.Flags.dreamHandlesBeingObscured;
 import static android.service.dreams.Flags.dreamHandlesConfirmKeys;
 import static android.service.dreams.Flags.startAndStopDozingInBackground;
 
+import android.Manifest;
 import android.annotation.FlaggedApi;
 import android.annotation.IdRes;
 import android.annotation.IntDef;
@@ -254,7 +255,6 @@ public class DreamService extends Service implements Window.Callback {
             "android.service.dream.DreamService.dream_overlay_component";
 
     private final IDreamManager mDreamManager;
-    private final Handler mHandler;
     private IBinder mDreamToken;
     private Window mWindow;
     private Activity mActivity;
@@ -323,19 +323,87 @@ public class DreamService extends Service implements Window.Callback {
         /** Returns the associated service info */
         ServiceInfo getServiceInfo();
 
-        /** Returns the handler to be used for any posted operation */
-        Handler getHandler();
-
         /** Returns the package manager */
         PackageManager getPackageManager();
 
         /** Returns the resources */
         Resources getResources();
+
+        /** Returns a specialized handler to ensure Runnables are not suspended */
+        WakefulHandler getWakefulHandler();
+    }
+
+    /**
+     * {@link WakefulHandler} is an interface for defining an object that helps post work without
+     * being interrupted by doze state.
+     *
+     * @hide
+     */
+    public interface WakefulHandler {
+        /** Posts a {@link Runnable} to be ran on the underlying {@link Handler}. */
+        void post(Runnable r);
+
+        /**
+         * Returns the underlying {@link Handler}. Should only be used for passing the handler into
+         * a function and not for directly calling methods on it.
+         */
+        Handler getHandler();
+    }
+
+    /**
+     * {@link WakefulHandlerImpl} ensures work on a handler is not suspended by wrapping the call
+     * with a partial wakelock. Note that this is only needed for Doze DreamService implementations.
+     * In this case, the component should have wake lock permissions. When such permission is not
+     * available, this class behaves like an ordinary handler.
+     */
+    private static final class WakefulHandlerImpl implements WakefulHandler {
+        private static final String SERVICE_HANDLER_WAKE_LOCK_TAG = "dream:service:handler";
+        private Context mContext;
+        private Handler mHandler;
+
+        private PowerManager.WakeLock mWakeLock;
+
+        private PowerManager.WakeLock getWakeLock() {
+            if (mContext.checkCallingOrSelfPermission(Manifest.permission.WAKE_LOCK)
+                    != PERMISSION_GRANTED) {
+                return null;
+            }
+
+            final PowerManager powerManager = mContext.getSystemService(PowerManager.class);
+
+            if (powerManager == null) {
+                return null;
+            }
+
+            return powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK,
+                    SERVICE_HANDLER_WAKE_LOCK_TAG);
+        }
+
+        WakefulHandlerImpl(Context context) {
+            mContext = context;
+            mHandler = new Handler(Looper.getMainLooper());
+            mWakeLock = getWakeLock();
+        }
+
+        @Override
+        public void post(Runnable r) {
+            if (mWakeLock != null) {
+                mHandler.post(mWakeLock.wrap(r));
+            } else {
+                mHandler.post(r);
+            }
+        }
+
+        @Override
+        public Handler getHandler() {
+            return mHandler;
+        }
     }
 
     private static final class DefaultInjector implements Injector {
         private Context mContext;
         private Class<?> mClassName;
+        private WakefulHandler mWakefulHandler;
 
         public void init(Context context) {
             mContext = context;
@@ -346,8 +414,6 @@ public class DreamService extends Service implements Window.Callback {
         public DreamOverlayConnectionHandler createOverlayConnection(
                 ComponentName overlayComponent,
                 Runnable onDisconnected) {
-            final Resources resources = mContext.getResources();
-
             return new DreamOverlayConnectionHandler(
                     /* context= */ mContext,
                     Looper.getMainLooper(),
@@ -381,8 +447,14 @@ public class DreamService extends Service implements Window.Callback {
         }
 
         @Override
-        public Handler getHandler() {
-            return new Handler(Looper.getMainLooper());
+        public WakefulHandler getWakefulHandler() {
+            synchronized (this) {
+                if (mWakefulHandler == null) {
+                    mWakefulHandler = new WakefulHandlerImpl(mContext);
+                }
+            }
+
+            return mWakefulHandler;
         }
 
         @Override
@@ -394,7 +466,6 @@ public class DreamService extends Service implements Window.Callback {
         public Resources getResources() {
             return mContext.getResources();
         }
-
     }
 
     public DreamService() {
@@ -412,7 +483,6 @@ public class DreamService extends Service implements Window.Callback {
         mInjector = injector;
         mInjector.init(this);
         mDreamManager = mInjector.getDreamManager();
-        mHandler = mInjector.getHandler();
     }
 
     /**
@@ -925,11 +995,17 @@ public class DreamService extends Service implements Window.Callback {
         }
     }
 
+    private void post(Runnable runnable) {
+        // The handler is based on the populated context is not ready at construction time.
+        // therefore we fetch on demand.
+        mInjector.getWakefulHandler().post(runnable);
+    }
+
     /**
      * Updates doze state. Note that this must be called on the mHandler.
      */
     private void updateDoze() {
-        mHandler.post(() -> {
+        post(() -> {
             if (mDreamToken == null) {
                 Slog.w(mTag, "Updating doze without a dream token.");
                 return;
@@ -971,7 +1047,7 @@ public class DreamService extends Service implements Window.Callback {
      */
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.P, trackingBug = 115609023)
     public void stopDozing() {
-        mHandler.post(() -> {
+        post(() -> {
             if (mDreamToken == null) {
                 return;
             }
@@ -1214,7 +1290,7 @@ public class DreamService extends Service implements Window.Callback {
                 final long token = Binder.clearCallingIdentity();
                 try {
                     // Simply finish dream when exit is requested.
-                    mHandler.post(() -> finishInternal());
+                    post(() -> finishInternal());
                 } finally {
                     Binder.restoreCallingIdentity(token);
                 }
@@ -1320,7 +1396,7 @@ public class DreamService extends Service implements Window.Callback {
      * </p>
      */
     public final void finish() {
-        mHandler.post(this::finishInternal);
+        post(this::finishInternal);
     }
 
     private void finishInternal() {
@@ -1382,7 +1458,7 @@ public class DreamService extends Service implements Window.Callback {
      * </p>
      */
     public final void wakeUp() {
-        mHandler.post(()-> wakeUp(false));
+        post(()-> wakeUp(false));
     }
 
     /**
@@ -1831,7 +1907,8 @@ public class DreamService extends Service implements Window.Callback {
 
     @Override
     protected void dump(final FileDescriptor fd, PrintWriter pw, final String[] args) {
-        DumpUtils.dumpAsync(mHandler, (pw1, prefix) -> dumpOnHandler(fd, pw1, args), pw, "", 1000);
+        DumpUtils.dumpAsync(mInjector.getWakefulHandler().getHandler(),
+                (pw1, prefix) -> dumpOnHandler(fd, pw1, args), pw, "", 1000);
     }
 
     /** @hide */
@@ -1887,7 +1964,7 @@ public class DreamService extends Service implements Window.Callback {
                 return;
             }
 
-            service.mHandler.post(() -> consumer.accept(service));
+            service.post(() -> consumer.accept(service));
         }
 
         @Override
diff --git a/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java b/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java
index 586ff52aa78c..a16d463cae71 100644
--- a/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java
+++ b/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java
@@ -33,7 +33,6 @@ import android.content.pm.ServiceInfo;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
 import android.graphics.drawable.Drawable;
-import android.os.Handler;
 import android.os.IBinder;
 import android.os.IRemoteCallback;
 import android.os.Looper;
@@ -176,7 +175,9 @@ public class TestDreamEnvironment {
         @Mock
         private ServiceInfo mServiceInfo;
 
-        private final Handler mHandler;
+        @Mock
+        private DreamService.WakefulHandler mWakefulHandler;
+
         private final IDreamManager mDreamManager;
         private final DreamOverlayConnectionHandler mDreamOverlayConnectionHandler;
 
@@ -185,7 +186,6 @@ public class TestDreamEnvironment {
                 DreamOverlayConnectionHandler dreamOverlayConnectionHandler,
                 boolean shouldShowComplications) {
             MockitoAnnotations.initMocks(this);
-            mHandler = new Handler(looper);
             mDreamManager = dreamManager;
             mDreamOverlayConnectionHandler = dreamOverlayConnectionHandler;
             mServiceInfo.packageName = FAKE_DREAM_PACKAGE_NAME;
@@ -198,6 +198,10 @@ public class TestDreamEnvironment {
                     .thenReturn(FAKE_DREAM_SETTINGS_ACTIVITY);
             when(mPackageManager.extractPackageItemInfoAttributes(any(), any(), any(), any()))
                     .thenReturn(mAttributes);
+            doAnswer(invocation -> {
+                ((Runnable) invocation.getArgument(0)).run();
+                return null;
+            }).when(mWakefulHandler).post(any());
         }
         @Override
         public void init(Context context) {
@@ -235,11 +239,6 @@ public class TestDreamEnvironment {
         }
 
         @Override
-        public Handler getHandler() {
-            return mHandler;
-        }
-
-        @Override
         public PackageManager getPackageManager() {
             return mPackageManager;
         }
@@ -248,6 +247,11 @@ public class TestDreamEnvironment {
         public Resources getResources() {
             return mResources;
         }
+
+        @Override
+        public DreamService.WakefulHandler getWakefulHandler() {
+            return mWakefulHandler;
+        }
     }
 
     @Mock