Fix DreamService -> DreamOverlayService synchronization issues.

This change moves the DreamOverlayService connection management logic
out of DreamService so that it can be separately tested. It also updates
the connection logic to use a handler for synchronization of messages,
and correctly cancels pending messages when needed.

This avoids consumers from running after the service has been
unbound.

Fixes: 267335683
Test: atest DreamOverlayConnectionHandlerTest
Test: flashed device and verified no crash when enabling/disabling
lowlight, and pressing the power button

Change-Id: I93855885ad80410693b8dcb4c2e696e0cabc5737

core/java/android/service/dreams/DreamOverlayConnectionHandler.java

@@ -0,0 +1,242 @@
+/*
+ * Copyright (C) 2023 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 android.service.dreams;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.content.Intent;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.Message;
+import android.os.RemoteException;
+import android.util.Log;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.util.ObservableServiceConnection;
+import com.android.internal.util.PersistentServiceConnection;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+
+/**
+ * Handles the service connection to {@link IDreamOverlay}
+ *
+ * @hide
+ */
+@VisibleForTesting
+public final class DreamOverlayConnectionHandler {
+    private static final String TAG = "DreamOverlayConnection";
+
+    private static final int MSG_ADD_CONSUMER = 1;
+    private static final int MSG_REMOVE_CONSUMER = 2;
+    private static final int MSG_OVERLAY_CLIENT_READY = 3;
+
+    private final Handler mHandler;
+    private final PersistentServiceConnection<IDreamOverlay> mConnection;
+    // Retrieved Client
+    private IDreamOverlayClient mClient;
+    // A list of pending requests to execute on the overlay.
+    private final List<Consumer<IDreamOverlayClient>> mConsumers = new ArrayList<>();
+    private final OverlayConnectionCallback mCallback;
+
+    DreamOverlayConnectionHandler(
+            Context context,
+            Looper looper,
+            Intent serviceIntent,
+            int minConnectionDurationMs,
+            int maxReconnectAttempts,
+            int baseReconnectDelayMs) {
+        this(context, looper, serviceIntent, minConnectionDurationMs, maxReconnectAttempts,
+                baseReconnectDelayMs, new Injector());
+    }
+
+    @VisibleForTesting
+    public DreamOverlayConnectionHandler(
+            Context context,
+            Looper looper,
+            Intent serviceIntent,
+            int minConnectionDurationMs,
+            int maxReconnectAttempts,
+            int baseReconnectDelayMs,
+            Injector injector) {
+        mCallback = new OverlayConnectionCallback();
+        mHandler = new Handler(looper, new OverlayHandlerCallback());
+        mConnection = injector.buildConnection(
+                context,
+                mHandler,
+                serviceIntent,
+                minConnectionDurationMs,
+                maxReconnectAttempts,
+                baseReconnectDelayMs
+        );
+    }
+
+    /**
+     * Bind to the overlay service. If binding fails, we automatically call unbind to clean
+     * up resources.
+     *
+     * @return true if binding was successful, false otherwise.
+     */
+    public boolean bind() {
+        mConnection.addCallback(mCallback);
+        final boolean success = mConnection.bind();
+        if (!success) {
+            unbind();
+        }
+        return success;
+    }
+
+    /**
+     * Unbind from the overlay service, clearing any pending callbacks.
+     */
+    public void unbind() {
+        mConnection.removeCallback(mCallback);
+        // Remove any pending messages.
+        mHandler.removeCallbacksAndMessages(null);
+        mClient = null;
+        mConsumers.clear();
+        mConnection.unbind();
+    }
+
+    /**
+     * Adds a consumer to run once the overlay service has connected. If the overlay service
+     * disconnects (eg binding dies) and then reconnects, this consumer will be re-run unless
+     * removed.
+     *
+     * @param consumer The consumer to run. This consumer is always executed asynchronously.
+     */
+    public void addConsumer(Consumer<IDreamOverlayClient> consumer) {
+        final Message msg = mHandler.obtainMessage(MSG_ADD_CONSUMER, consumer);
+        mHandler.sendMessage(msg);
+    }
+
+    /**
+     * Removes the consumer, preventing this consumer from being called again.
+     *
+     * @param consumer The consumer to remove.
+     */
+    public void removeConsumer(Consumer<IDreamOverlayClient> consumer) {
+        final Message msg = mHandler.obtainMessage(MSG_REMOVE_CONSUMER, consumer);
+        mHandler.sendMessage(msg);
+        // Clear any pending messages to add this consumer
+        mHandler.removeMessages(MSG_ADD_CONSUMER, consumer);
+    }
+
+    private final class OverlayHandlerCallback implements Handler.Callback {
+        @Override
+        public boolean handleMessage(@NonNull Message msg) {
+            switch (msg.what) {
+                case MSG_OVERLAY_CLIENT_READY:
+                    onOverlayClientReady((IDreamOverlayClient) msg.obj);
+                    break;
+                case MSG_ADD_CONSUMER:
+                    onAddConsumer((Consumer<IDreamOverlayClient>) msg.obj);
+                    break;
+                case MSG_REMOVE_CONSUMER:
+                    onRemoveConsumer((Consumer<IDreamOverlayClient>) msg.obj);
+                    break;
+            }
+            return true;
+        }
+    }
+
+    private void onOverlayClientReady(IDreamOverlayClient client) {
+        mClient = client;
+        for (Consumer<IDreamOverlayClient> consumer : mConsumers) {
+            consumer.accept(mClient);
+        }
+    }
+
+    private void onAddConsumer(Consumer<IDreamOverlayClient> consumer) {
+        if (mClient != null) {
+            consumer.accept(mClient);
+        }
+        mConsumers.add(consumer);
+    }
+
+    private void onRemoveConsumer(Consumer<IDreamOverlayClient> consumer) {
+        mConsumers.remove(consumer);
+    }
+
+    private final class OverlayConnectionCallback implements
+            ObservableServiceConnection.Callback<IDreamOverlay> {
+
+        private final IDreamOverlayClientCallback mClientCallback =
+                new IDreamOverlayClientCallback.Stub() {
+                    @Override
+                    public void onDreamOverlayClient(IDreamOverlayClient client) {
+                        final Message msg =
+                                mHandler.obtainMessage(MSG_OVERLAY_CLIENT_READY, client);
+                        mHandler.sendMessage(msg);
+                    }
+                };
+
+        @Override
+        public void onConnected(
+                ObservableServiceConnection<IDreamOverlay> connection,
+                IDreamOverlay service) {
+            try {
+                service.getClient(mClientCallback);
+            } catch (RemoteException e) {
+                Log.e(TAG, "could not get DreamOverlayClient", e);
+            }
+        }
+
+        @Override
+        public void onDisconnected(ObservableServiceConnection<IDreamOverlay> connection,
+                int reason) {
+            mClient = null;
+            // Cancel any pending messages about the overlay being ready, since it is no
+            // longer ready.
+            mHandler.removeMessages(MSG_OVERLAY_CLIENT_READY);
+        }
+    }
+
+    /**
+     * Injector for testing
+     */
+    @VisibleForTesting
+    public static class Injector {
+        /**
+         * Returns milliseconds since boot, not counting time spent in deep sleep. Can be overridden
+         * in tests with a fake clock.
+         */
+        public PersistentServiceConnection<IDreamOverlay> buildConnection(
+                Context context,
+                Handler handler,
+                Intent serviceIntent,
+                int minConnectionDurationMs,
+                int maxReconnectAttempts,
+                int baseReconnectDelayMs) {
+            final Executor executor = handler::post;
+            final int flags = Context.BIND_AUTO_CREATE | Context.BIND_FOREGROUND_SERVICE;
+            return new PersistentServiceConnection<>(
+                    context,
+                    executor,
+                    handler,
+                    IDreamOverlay.Stub::asInterface,
+                    serviceIntent,
+                    flags,
+                    minConnectionDurationMs,
+                    maxReconnectAttempts,
+                    baseReconnectDelayMs
+            );
+        }
+    }
+}

core/java/android/service/dreams/DreamService.java

@@ -68,8 +68,6 @@ import android.view.accessibility.AccessibilityEvent;
 
 import com.android.internal.R;
 import com.android.internal.util.DumpUtils;
-import com.android.internal.util.ObservableServiceConnection;
-import com.android.internal.util.PersistentServiceConnection;
 
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
@@ -77,8 +75,6 @@ import org.xmlpull.v1.XmlPullParserException;
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.io.PrintWriter;
-import java.util.ArrayList;
-import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
 /**
@@ -234,7 +230,6 @@ public class DreamService extends Service implements Window.Callback {
     private boolean mCanDoze;
     private boolean mDozing;
     private boolean mWindowless;
-    private boolean mOverlayFinishing;
     private int mDozeScreenState = Display.STATE_UNKNOWN;
     private int mDozeScreenBrightness = PowerManager.BRIGHTNESS_DEFAULT;
 
@@ -246,88 +241,7 @@ public class DreamService extends Service implements Window.Callback {
     private DreamServiceWrapper mDreamServiceWrapper;
     private Runnable mDispatchAfterOnAttachedToWindow;
 
-    private OverlayConnection mOverlayConnection;
-
-    private static class OverlayConnection extends PersistentServiceConnection<IDreamOverlay> {
-        // Retrieved Client
-        private IDreamOverlayClient mClient;
-
-        // A list of pending requests to execute on the overlay.
-        private final ArrayList<Consumer<IDreamOverlayClient>> mConsumers = new ArrayList<>();
-
-        private final IDreamOverlayClientCallback mClientCallback =
-                new IDreamOverlayClientCallback.Stub() {
-            @Override
-            public void onDreamOverlayClient(IDreamOverlayClient client) {
-                mClient = client;
-
-                for (Consumer<IDreamOverlayClient> consumer : mConsumers) {
-                    consumer.accept(mClient);
-                }
-            }
-        };
-
-        private final Callback<IDreamOverlay> mCallback = new Callback<IDreamOverlay>() {
-            @Override
-            public void onConnected(ObservableServiceConnection<IDreamOverlay> connection,
-                    IDreamOverlay service) {
-                try {
-                    service.getClient(mClientCallback);
-                } catch (RemoteException e) {
-                    Log.e(TAG, "could not get DreamOverlayClient", e);
-                }
-            }
-
-            @Override
-            public void onDisconnected(ObservableServiceConnection<IDreamOverlay> connection,
-                    int reason) {
-                mClient = null;
-            }
-        };
-
-        OverlayConnection(Context context,
-                Executor executor,
-                Handler handler,
-                ServiceTransformer<IDreamOverlay> transformer,
-                Intent serviceIntent,
-                int flags,
-                int minConnectionDurationMs,
-                int maxReconnectAttempts,
-                int baseReconnectDelayMs) {
-            super(context, executor, handler, transformer, serviceIntent, flags,
-                    minConnectionDurationMs,
-                    maxReconnectAttempts, baseReconnectDelayMs);
-        }
-
-        @Override
-        public boolean bind() {
-            addCallback(mCallback);
-            return super.bind();
-        }
-
-        @Override
-        public void unbind() {
-            removeCallback(mCallback);
-            super.unbind();
-        }
-
-        public void addConsumer(Consumer<IDreamOverlayClient> consumer) {
-            execute(() -> {
-                mConsumers.add(consumer);
-                if (mClient != null) {
-                    consumer.accept(mClient);
-                }
-            });
-        }
-
-        public void removeConsumer(Consumer<IDreamOverlayClient> consumer) {
-            execute(() -> mConsumers.remove(consumer));
-        }
-
-        public void clearConsumers() {
-            execute(() -> mConsumers.clear());
-        }
-    }
+    private DreamOverlayConnectionHandler mOverlayConnection;
 
     private final IDreamOverlayCallback mOverlayCallback = new IDreamOverlayCallback.Stub() {
         @Override
@@ -1030,18 +944,18 @@ public class DreamService extends Service implements Window.Callback {
             final Resources resources = getResources();
             final Intent overlayIntent = new Intent().setComponent(overlayComponent);
 
-            mOverlayConnection = new OverlayConnection(
+            mOverlayConnection = new DreamOverlayConnectionHandler(
                     /* context= */ this,
-                    getMainExecutor(),
-                    mHandler,
-                    IDreamOverlay.Stub::asInterface,
+                    Looper.getMainLooper(),
                     overlayIntent,
-                    /* flags= */ Context.BIND_AUTO_CREATE | Context.BIND_FOREGROUND_SERVICE,
                     resources.getInteger(R.integer.config_minDreamOverlayDurationMs),
                     resources.getInteger(R.integer.config_dreamOverlayMaxReconnectAttempts),
                     resources.getInteger(R.integer.config_dreamOverlayReconnectTimeoutMs));
 
-            mOverlayConnection.bind();
+            if (!mOverlayConnection.bind()) {
+                // Binding failed.
+                mOverlayConnection = null;
+            }
         }
 
         return mDreamServiceWrapper;
@@ -1069,9 +983,7 @@ public class DreamService extends Service implements Window.Callback {
         // If there is an active overlay connection, signal that the dream is ending before
         // continuing. Note that the overlay cannot rely on the unbound state, since another dream
         // might have bound to it in the meantime.
-        if (mOverlayConnection != null && !mOverlayFinishing) {
-            // Set mOverlayFinish to true to only allow this consumer to be added once.
-            mOverlayFinishing = true;
+        if (mOverlayConnection != null) {
             mOverlayConnection.addConsumer(overlay -> {
                 try {
                     overlay.endDream();
@@ -1082,7 +994,6 @@ public class DreamService extends Service implements Window.Callback {
                     Log.e(mTag, "could not inform overlay of dream end:" + e);
                 }
             });
-            mOverlayConnection.clearConsumers();
             return;
         }
 

services/tests/mockingservicestests/src/android/service/dreams/DreamOverlayConnectionHandlerTest.java

@@ -0,0 +1,245 @@
+/*
+ * Copyright (C) 2023 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 android.service.dreams;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+import android.content.Intent;
+import android.os.Handler;
+import android.os.RemoteException;
+import android.os.test.TestLooper;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.internal.util.ObservableServiceConnection;
+import com.android.internal.util.PersistentServiceConnection;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class DreamOverlayConnectionHandlerTest {
+    private static final int MIN_CONNECTION_DURATION_MS = 100;
+    private static final int MAX_RECONNECT_ATTEMPTS = 3;
+    private static final int BASE_RECONNECT_DELAY_MS = 50;
+
+    @Mock
+    private Context mContext;
+    @Mock
+    private PersistentServiceConnection<IDreamOverlay> mConnection;
+    @Mock
+    private Intent mServiceIntent;
+    @Mock
+    private IDreamOverlay mOverlayService;
+    @Mock
+    private IDreamOverlayClient mOverlayClient;
+
+    private TestLooper mTestLooper;
+    private DreamOverlayConnectionHandler mDreamOverlayConnectionHandler;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mTestLooper = new TestLooper();
+        mDreamOverlayConnectionHandler = new DreamOverlayConnectionHandler(
+                mContext,
+                mTestLooper.getLooper(),
+                mServiceIntent,
+                MIN_CONNECTION_DURATION_MS,
+                MAX_RECONNECT_ATTEMPTS,
+                BASE_RECONNECT_DELAY_MS,
+                new TestInjector(mConnection));
+    }
+
+    @Test
+    public void consumerShouldRunImmediatelyWhenClientAvailable() throws RemoteException {
+        mDreamOverlayConnectionHandler.bind();
+        connectService();
+        provideClient();
+
+        final Consumer<IDreamOverlayClient> consumer = Mockito.mock(Consumer.class);
+        mDreamOverlayConnectionHandler.addConsumer(consumer);
+        mTestLooper.dispatchAll();
+        verify(consumer).accept(mOverlayClient);
+    }
+
+    @Test
+    public void consumerShouldRunAfterClientAvailable() throws RemoteException {
+        mDreamOverlayConnectionHandler.bind();
+        connectService();
+
+        final Consumer<IDreamOverlayClient> consumer = Mockito.mock(Consumer.class);
+        mDreamOverlayConnectionHandler.addConsumer(consumer);
+        mTestLooper.dispatchAll();
+        // No client yet, so we shouldn't have executed
+        verify(consumer, never()).accept(mOverlayClient);
+
+        provideClient();
+        mTestLooper.dispatchAll();
+        verify(consumer).accept(mOverlayClient);
+    }
+
+    @Test
+    public void consumerShouldNeverRunIfClientConnectsAndDisconnects() throws RemoteException {
+        mDreamOverlayConnectionHandler.bind();
+        connectService();
+
+        final Consumer<IDreamOverlayClient> consumer = Mockito.mock(Consumer.class);
+        mDreamOverlayConnectionHandler.addConsumer(consumer);
+        mTestLooper.dispatchAll();
+        // No client yet, so we shouldn't have executed
+        verify(consumer, never()).accept(mOverlayClient);
+
+        provideClient();
+        // Service disconnected before looper could handle the message.
+        disconnectService();
+        mTestLooper.dispatchAll();
+        verify(consumer, never()).accept(mOverlayClient);
+    }
+
+    @Test
+    public void consumerShouldNeverRunIfUnbindCalled() throws RemoteException {
+        mDreamOverlayConnectionHandler.bind();
+        connectService();
+        provideClient();
+
+        final Consumer<IDreamOverlayClient> consumer = Mockito.mock(Consumer.class);
+        mDreamOverlayConnectionHandler.addConsumer(consumer);
+        mDreamOverlayConnectionHandler.unbind();
+        mTestLooper.dispatchAll();
+        // We unbinded immediately after adding consumer, so should never have run.
+        verify(consumer, never()).accept(mOverlayClient);
+    }
+
+    @Test
+    public void consumersOnlyRunOnceIfUnbound() throws RemoteException {
+        mDreamOverlayConnectionHandler.bind();
+        connectService();
+        provideClient();
+
+        AtomicInteger counter = new AtomicInteger();
+        // Add 10 consumers in a row which call unbind within the consumer.
+        for (int i = 0; i < 10; i++) {
+            mDreamOverlayConnectionHandler.addConsumer(client -> {
+                counter.getAndIncrement();
+                mDreamOverlayConnectionHandler.unbind();
+            });
+        }
+        mTestLooper.dispatchAll();
+        // Only the first consumer should have run, since we unbinded.
+        assertThat(counter.get()).isEqualTo(1);
+    }
+
+    @Test
+    public void consumerShouldRunAgainAfterReconnect() throws RemoteException {
+        mDreamOverlayConnectionHandler.bind();
+        connectService();
+        provideClient();
+
+        final Consumer<IDreamOverlayClient> consumer = Mockito.mock(Consumer.class);
+        mDreamOverlayConnectionHandler.addConsumer(consumer);
+        mTestLooper.dispatchAll();
+        verify(consumer, times(1)).accept(mOverlayClient);
+
+        disconnectService();
+        mTestLooper.dispatchAll();
+        // No new calls should happen when service disconnected.
+        verify(consumer, times(1)).accept(mOverlayClient);
+
+        connectService();
+        provideClient();
+        mTestLooper.dispatchAll();
+        // We should trigger the consumer again once the server reconnects.
+        verify(consumer, times(2)).accept(mOverlayClient);
+    }
+
+    @Test
+    public void consumerShouldNeverRunIfRemovedImmediately() throws RemoteException {
+        mDreamOverlayConnectionHandler.bind();
+        connectService();
+        provideClient();
+
+        final Consumer<IDreamOverlayClient> consumer = Mockito.mock(Consumer.class);
+        mDreamOverlayConnectionHandler.addConsumer(consumer);
+        mDreamOverlayConnectionHandler.removeConsumer(consumer);
+        mTestLooper.dispatchAll();
+        verify(consumer, never()).accept(mOverlayClient);
+    }
+
+    private void connectService() {
+        final ObservableServiceConnection.Callback<IDreamOverlay> callback =
+                captureConnectionCallback();
+        callback.onConnected(mConnection, mOverlayService);
+    }
+
+    private void disconnectService() {
+        final ObservableServiceConnection.Callback<IDreamOverlay> callback =
+                captureConnectionCallback();
+        callback.onDisconnected(mConnection, /* reason= */ 0);
+    }
+
+    private void provideClient() throws RemoteException {
+        final IDreamOverlayClientCallback callback = captureClientCallback();
+        callback.onDreamOverlayClient(mOverlayClient);
+    }
+
+    private ObservableServiceConnection.Callback<IDreamOverlay> captureConnectionCallback() {
+        ArgumentCaptor<ObservableServiceConnection.Callback<IDreamOverlay>>
+                callbackCaptor =
+                ArgumentCaptor.forClass(ObservableServiceConnection.Callback.class);
+        verify(mConnection).addCallback(callbackCaptor.capture());
+        return callbackCaptor.getValue();
+    }
+
+    private IDreamOverlayClientCallback captureClientCallback() throws RemoteException {
+        ArgumentCaptor<IDreamOverlayClientCallback> callbackCaptor =
+                ArgumentCaptor.forClass(IDreamOverlayClientCallback.class);
+        verify(mOverlayService, atLeastOnce()).getClient(callbackCaptor.capture());
+        return callbackCaptor.getValue();
+    }
+
+    static class TestInjector extends DreamOverlayConnectionHandler.Injector {
+        private final PersistentServiceConnection<IDreamOverlay> mConnection;
+
+        TestInjector(PersistentServiceConnection<IDreamOverlay> connection) {
+            mConnection = connection;
+        }
+
+        @Override
+        public PersistentServiceConnection<IDreamOverlay> buildConnection(Context context,
+                Handler handler, Intent serviceIntent, int minConnectionDurationMs,
+                int maxReconnectAttempts, int baseReconnectDelayMs) {
+            return mConnection;
+        }
+    }
+}