Fix wake lock logic during Dream startup

Fixes a crash that would happen in all dreams
that did not have permission to acquire wake locks.
Instead moves the wake lock logic into the system
process. Also fixes a bug in DozeService where the
wake lock was not held until dozing was actually
properly initialized.

Fixes: 31612287
Bug: 31044352
Related-CL: I85955a2b7d6bad5171accbc336117a9660b1b198
Test: adb shell settings put secure screensaver_components com.android.dreams.basic/.Colors; adb shell service call dreams 1

Change-Id: Idb3f921ee71b6da6c2ab0c44c332ef91f93ddbc0
diff --git a/core/java/android/service/dreams/DreamService.java b/core/java/android/service/dreams/DreamService.java
index e958fbe..94505d3 100644
--- a/core/java/android/service/dreams/DreamService.java
+++ b/core/java/android/service/dreams/DreamService.java
@@ -27,6 +27,7 @@
 import android.graphics.drawable.ColorDrawable;
 import android.os.Handler;
 import android.os.IBinder;
+import android.os.IRemoteCallback;
 import android.os.PowerManager;
 import android.os.RemoteException;
 import android.os.ServiceManager;
@@ -192,9 +193,6 @@
 
     private boolean mDebug = false;
 
-    private PowerManager.WakeLock mWakeLock;
-    private boolean mWakeLockAcquired;
-
     public DreamService() {
         mSandman = IDreamManager.Stub.asInterface(ServiceManager.getService(DREAM_SERVICE));
     }
@@ -789,8 +787,6 @@
     public void onCreate() {
         if (mDebug) Slog.v(TAG, "onCreate()");
         super.onCreate();
-        mWakeLock = getSystemService(PowerManager.class)
-                .newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, "DreamService");
     }
 
     /**
@@ -830,21 +826,9 @@
     @Override
     public final IBinder onBind(Intent intent) {
         if (mDebug) Slog.v(TAG, "onBind() intent = " + intent);
-
-        // Need to stay awake until we dispatch onDreamingStarted. This is released either in
-        // attach() or onDestroy().
-        mWakeLock.acquire(5000);
-        mWakeLockAcquired = true;
         return new DreamServiceWrapper();
     }
 
-    private void releaseWakeLockIfNeeded() {
-        if (mWakeLockAcquired) {
-            mWakeLock.release();
-            mWakeLockAcquired = false;
-        }
-    }
-
     /**
      * Stops the dream and detaches from the window.
      * <p>
@@ -921,8 +905,6 @@
         detach();
 
         super.onDestroy();
-
-        releaseWakeLockIfNeeded(); // for acquire in onBind()
     }
 
     // end public api
@@ -961,90 +943,94 @@
      * Must run on mHandler.
      *
      * @param windowToken A window token that will allow a window to be created in the correct layer.
+     * @param started A callback that will be invoked once onDreamingStarted has completed.
      */
-    private final void attach(IBinder windowToken, boolean canDoze) {
-        try {
-            if (mWindowToken != null) {
-                Slog.e(TAG, "attach() called when already attached with token=" + mWindowToken);
-                return;
-            }
-            if (mFinished || mWaking) {
-                Slog.w(TAG, "attach() called after dream already finished");
-                try {
-                    mSandman.finishSelf(windowToken, true /*immediate*/);
-                } catch (RemoteException ex) {
-                    // system server died
-                }
-                return;
-            }
-
-            mWindowToken = windowToken;
-            mCanDoze = canDoze;
-            if (mWindowless && !mCanDoze) {
-                throw new IllegalStateException("Only doze dreams can be windowless");
-            }
-            if (!mWindowless) {
-                mWindow = new PhoneWindow(this);
-                mWindow.setCallback(this);
-                mWindow.requestFeature(Window.FEATURE_NO_TITLE);
-                mWindow.setBackgroundDrawable(new ColorDrawable(0xFF000000));
-                mWindow.setFormat(PixelFormat.OPAQUE);
-
-                if (mDebug) {
-                    Slog.v(TAG, String.format("Attaching window token: %s to window of type %s",
-                            windowToken, WindowManager.LayoutParams.TYPE_DREAM));
-                }
-
-                WindowManager.LayoutParams lp = mWindow.getAttributes();
-                lp.type = WindowManager.LayoutParams.TYPE_DREAM;
-                lp.token = windowToken;
-                lp.windowAnimations = com.android.internal.R.style.Animation_Dream;
-                lp.flags |= (WindowManager.LayoutParams.FLAG_LAYOUT_IN_SCREEN
-                            | WindowManager.LayoutParams.FLAG_LAYOUT_INSET_DECOR
-                            | WindowManager.LayoutParams.FLAG_SHOW_WHEN_LOCKED
-                            | WindowManager.LayoutParams.FLAG_DISMISS_KEYGUARD
-                            | WindowManager.LayoutParams.FLAG_ALLOW_LOCK_WHILE_SCREEN_ON
-                            | (mFullscreen ? WindowManager.LayoutParams.FLAG_FULLSCREEN : 0)
-                            | (mScreenBright ? WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON : 0)
-                            );
-                mWindow.setAttributes(lp);
-                // Workaround: Currently low-profile and in-window system bar backgrounds don't go
-                // along well. Dreams usually don't need such bars anyways, so disable them by default.
-                mWindow.clearFlags(WindowManager.LayoutParams.FLAG_DRAWS_SYSTEM_BAR_BACKGROUNDS);
-                mWindow.setWindowManager(null, windowToken, "dream", true);
-
-                applySystemUiVisibilityFlags(
-                        (mLowProfile ? View.SYSTEM_UI_FLAG_LOW_PROFILE : 0),
-                        View.SYSTEM_UI_FLAG_LOW_PROFILE);
-
-                try {
-                    getWindowManager().addView(mWindow.getDecorView(), mWindow.getAttributes());
-                } catch (WindowManager.BadTokenException ex) {
-                    // This can happen because the dream manager service will remove the token
-                    // immediately without necessarily waiting for the dream to start.
-                    // We should receive a finish message soon.
-                    Slog.i(TAG, "attach() called after window token already removed, dream will "
-                            + "finish soon");
-                    mWindow = null;
-                    return;
-                }
-            }
-            // We need to defer calling onDreamingStarted until after onWindowAttached,
-            // which is posted to the handler by addView, so we post onDreamingStarted
-            // to the handler also.  Need to watch out here in case detach occurs before
-            // this callback is invoked.
-            mHandler.post(mWakeLock.wrap(() -> {
-                if (mWindow != null || mWindowless) {
-                    if (mDebug) {
-                        Slog.v(TAG, "Calling onDreamingStarted()");
-                    }
-                    mStarted = true;
-                    onDreamingStarted();
-                }
-            }));
-        } finally {
-            releaseWakeLockIfNeeded(); // for acquire in onBind
+    private final void attach(IBinder windowToken, boolean canDoze, IRemoteCallback started) {
+        if (mWindowToken != null) {
+            Slog.e(TAG, "attach() called when already attached with token=" + mWindowToken);
+            return;
         }
+        if (mFinished || mWaking) {
+            Slog.w(TAG, "attach() called after dream already finished");
+            try {
+                mSandman.finishSelf(windowToken, true /*immediate*/);
+            } catch (RemoteException ex) {
+                // system server died
+            }
+            return;
+        }
+
+        mWindowToken = windowToken;
+        mCanDoze = canDoze;
+        if (mWindowless && !mCanDoze) {
+            throw new IllegalStateException("Only doze dreams can be windowless");
+        }
+        if (!mWindowless) {
+            mWindow = new PhoneWindow(this);
+            mWindow.setCallback(this);
+            mWindow.requestFeature(Window.FEATURE_NO_TITLE);
+            mWindow.setBackgroundDrawable(new ColorDrawable(0xFF000000));
+            mWindow.setFormat(PixelFormat.OPAQUE);
+
+            if (mDebug) Slog.v(TAG, String.format("Attaching window token: %s to window of type %s",
+                    windowToken, WindowManager.LayoutParams.TYPE_DREAM));
+
+            WindowManager.LayoutParams lp = mWindow.getAttributes();
+            lp.type = WindowManager.LayoutParams.TYPE_DREAM;
+            lp.token = windowToken;
+            lp.windowAnimations = com.android.internal.R.style.Animation_Dream;
+            lp.flags |= ( WindowManager.LayoutParams.FLAG_LAYOUT_IN_SCREEN
+                        | WindowManager.LayoutParams.FLAG_LAYOUT_INSET_DECOR
+                        | WindowManager.LayoutParams.FLAG_SHOW_WHEN_LOCKED
+                        | WindowManager.LayoutParams.FLAG_DISMISS_KEYGUARD
+                        | WindowManager.LayoutParams.FLAG_ALLOW_LOCK_WHILE_SCREEN_ON
+                        | (mFullscreen ? WindowManager.LayoutParams.FLAG_FULLSCREEN : 0)
+                        | (mScreenBright ? WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON : 0)
+                        );
+            mWindow.setAttributes(lp);
+            // Workaround: Currently low-profile and in-window system bar backgrounds don't go
+            // along well. Dreams usually don't need such bars anyways, so disable them by default.
+            mWindow.clearFlags(WindowManager.LayoutParams.FLAG_DRAWS_SYSTEM_BAR_BACKGROUNDS);
+            mWindow.setWindowManager(null, windowToken, "dream", true);
+
+            applySystemUiVisibilityFlags(
+                    (mLowProfile ? View.SYSTEM_UI_FLAG_LOW_PROFILE : 0),
+                    View.SYSTEM_UI_FLAG_LOW_PROFILE);
+
+            try {
+                getWindowManager().addView(mWindow.getDecorView(), mWindow.getAttributes());
+            } catch (WindowManager.BadTokenException ex) {
+                // This can happen because the dream manager service will remove the token
+                // immediately without necessarily waiting for the dream to start.
+                // We should receive a finish message soon.
+                Slog.i(TAG, "attach() called after window token already removed, dream will "
+                        + "finish soon");
+                mWindow = null;
+                return;
+            }
+        }
+        // We need to defer calling onDreamingStarted until after onWindowAttached,
+        // which is posted to the handler by addView, so we post onDreamingStarted
+        // to the handler also.  Need to watch out here in case detach occurs before
+        // this callback is invoked.
+        mHandler.post(new Runnable() {
+            @Override
+            public void run() {
+                if (mWindow != null || mWindowless) {
+                    if (mDebug) Slog.v(TAG, "Calling onDreamingStarted()");
+                    mStarted = true;
+                    try {
+                        onDreamingStarted();
+                    } finally {
+                        try {
+                            started.sendResult(null);
+                        } catch (RemoteException e) {
+                            throw e.rethrowFromSystemServer();
+                        }
+                    }
+                }
+            }
+        });
     }
 
     private boolean getWindowFlagValue(int flag, boolean defaultValue) {
@@ -1116,11 +1102,12 @@
 
     private final class DreamServiceWrapper extends IDreamService.Stub {
         @Override
-        public void attach(final IBinder windowToken, final boolean canDoze) {
+        public void attach(final IBinder windowToken, final boolean canDoze,
+                IRemoteCallback started) {
             mHandler.post(new Runnable() {
                 @Override
                 public void run() {
-                    DreamService.this.attach(windowToken, canDoze);
+                    DreamService.this.attach(windowToken, canDoze, started);
                 }
             });
         }
diff --git a/core/java/android/service/dreams/IDreamService.aidl b/core/java/android/service/dreams/IDreamService.aidl
index 9bb1804..ce04354 100644
--- a/core/java/android/service/dreams/IDreamService.aidl
+++ b/core/java/android/service/dreams/IDreamService.aidl
@@ -16,11 +16,13 @@
 
 package android.service.dreams;
 
+import android.os.IRemoteCallback;
+
 /**
  * @hide
  */
 oneway interface IDreamService {
-    void attach(IBinder windowToken, boolean canDoze);
+    void attach(IBinder windowToken, boolean canDoze, IRemoteCallback started);
     void detach();
     void wakeUp();
 }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeService.java b/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
index ec4f447..4edcb4b 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
@@ -159,18 +159,15 @@
         // Ask the host to get things ready to start dozing.
         // Once ready, we call startDozing() at which point the CPU may suspend
         // and we will need to acquire a wakelock to do work.
-        mHost.startDozing(new Runnable() {
-            @Override
-            public void run() {
-                if (mDreaming) {
-                    startDozing();
+        mHost.startDozing(mWakeLock.wrap(() -> {
+            if (mDreaming) {
+                startDozing();
 
-                    // From this point until onDreamingStopped we will need to hold a
-                    // wakelock whenever we are doing work.  Note that we never call
-                    // stopDozing because can we just keep dozing until the bitter end.
-                }
+                // From this point until onDreamingStopped we will need to hold a
+                // wakelock whenever we are doing work.  Note that we never call
+                // stopDozing because can we just keep dozing until the bitter end.
             }
-        });
+        }));
     }
 
     @Override
diff --git a/services/core/java/com/android/server/dreams/DreamController.java b/services/core/java/com/android/server/dreams/DreamController.java
index 9fa93f4..3072f43 100644
--- a/services/core/java/com/android/server/dreams/DreamController.java
+++ b/services/core/java/com/android/server/dreams/DreamController.java
@@ -24,8 +24,10 @@
 import android.content.Intent;
 import android.content.ServiceConnection;
 import android.os.Binder;
+import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
+import android.os.IRemoteCallback;
 import android.os.PowerManager;
 import android.os.RemoteException;
 import android.os.IBinder.DeathRecipient;
@@ -253,7 +255,8 @@
     private void attach(IDreamService service) {
         try {
             service.asBinder().linkToDeath(mCurrentDream, 0);
-            service.attach(mCurrentDream.mToken, mCurrentDream.mCanDoze);
+            service.attach(mCurrentDream.mToken, mCurrentDream.mCanDoze,
+                    mCurrentDream.mDreamingStartedCallback);
         } catch (RemoteException ex) {
             Slog.e(TAG, "The dream service died unexpectedly.", ex);
             stopDream(true /*immediate*/);
@@ -298,10 +301,10 @@
             mCanDoze = canDoze;
             mUserId  = userId;
             mWakeLock = wakeLock;
-            // Hold the lock while we're waiting for the service to connect. Released either when
-            // DreamService connects (and is then responsible for keeping the device awake) or
-            // dreaming stops.
+            // Hold the lock while we're waiting for the service to connect and start dreaming.
+            // Released after the service has started dreaming, we stop dreaming, or it timed out.
             mWakeLock.acquire();
+            mHandler.postDelayed(mReleaseWakeLockIfNeeded, 10000);
         }
 
         // May be called on any thread.
@@ -324,25 +327,17 @@
             mHandler.post(new Runnable() {
                 @Override
                 public void run() {
-                    try {
-                        mConnected = true;
-                        if (mCurrentDream == DreamRecord.this && mService == null) {
-                            attach(IDreamService.Stub.asInterface(service));
-                        }
-                    } finally {
+                    mConnected = true;
+                    if (mCurrentDream == DreamRecord.this && mService == null) {
+                        attach(IDreamService.Stub.asInterface(service));
+                        // Wake lock will be released once dreaming starts.
+                    } else {
                         releaseWakeLockIfNeeded();
                     }
                 }
             });
         }
 
-        private void releaseWakeLockIfNeeded() {
-            if (mWakeLock != null) {
-                mWakeLock.release();
-                mWakeLock = null;
-            }
-        }
-
         // May be called on any thread.
         @Override
         public void onServiceDisconnected(ComponentName name) {
@@ -356,5 +351,23 @@
                 }
             });
         }
+
+        void releaseWakeLockIfNeeded() {
+            if (mWakeLock != null) {
+                mWakeLock.release();
+                mWakeLock = null;
+                mHandler.removeCallbacks(mReleaseWakeLockIfNeeded);
+            }
+        }
+
+        final Runnable mReleaseWakeLockIfNeeded = this::releaseWakeLockIfNeeded;
+
+        final IRemoteCallback mDreamingStartedCallback = new IRemoteCallback.Stub() {
+            // May be called on any thread.
+            @Override
+            public void sendResult(Bundle data) throws RemoteException {
+                mHandler.post(mReleaseWakeLockIfNeeded);
+            }
+        };
     }
 }
\ No newline at end of file