Merge "Change BroadcastDispatcher to 1receiver/action" into rvc-dev am: d6b99a586a

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/11871757

Change-Id: Ieea3a890545df4ecc0db3d0f79ec166f63065917

packages/SystemUI/src/com/android/systemui/broadcast/ActionReceiver.kt

@@ -0,0 +1,133 @@
+/*
+ * Copyright (C) 2020 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.systemui.broadcast
+
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.util.ArraySet
+import com.android.systemui.Dumpable
+import com.android.systemui.broadcast.logging.BroadcastDispatcherLogger
+import com.android.systemui.util.indentIfPossible
+import java.io.FileDescriptor
+import java.io.PrintWriter
+import java.util.concurrent.Executor
+import java.util.concurrent.atomic.AtomicInteger
+
+/**
+ * Receiver for a given action-userId pair to be used by [UserBroadcastDispatcher].
+ *
+ * Each object of this class will take care of a single Action. It will register if it has at least
+ * one [BroadcastReceiver] added to it, and unregister when none are left.
+ *
+ * It will also re-register if filters with new categories are added. But this should not happen
+ * often.
+ *
+ * This class has no sync controls, so make sure to only make modifications from the background
+ * thread.
+ */
+class ActionReceiver(
+    private val action: String,
+    private val userId: Int,
+    private val registerAction: BroadcastReceiver.(IntentFilter) -> Unit,
+    private val unregisterAction: BroadcastReceiver.() -> Unit,
+    private val bgExecutor: Executor,
+    private val logger: BroadcastDispatcherLogger
+) : BroadcastReceiver(), Dumpable {
+
+    companion object {
+        val index = AtomicInteger(0)
+    }
+
+    var registered = false
+        private set
+    private val receiverDatas = ArraySet<ReceiverData>()
+    private val activeCategories = ArraySet<String>()
+
+    @Throws(IllegalArgumentException::class)
+    fun addReceiverData(receiverData: ReceiverData) {
+        if (!receiverData.filter.hasAction(action)) {
+            throw(IllegalArgumentException("Trying to attach to $action without correct action," +
+                "receiver: ${receiverData.receiver}"))
+        }
+        val addedCategories = activeCategories
+                .addAll(receiverData.filter.categoriesIterator()?.asSequence() ?: emptySequence())
+
+        if (receiverDatas.add(receiverData) && receiverDatas.size == 1) {
+            registerAction(createFilter())
+            registered = true
+        } else if (addedCategories) {
+            unregisterAction()
+            registerAction(createFilter())
+        }
+    }
+
+    fun hasReceiver(receiver: BroadcastReceiver): Boolean {
+        return receiverDatas.any { it.receiver == receiver }
+    }
+
+    private fun createFilter(): IntentFilter {
+        val filter = IntentFilter(action)
+        activeCategories.forEach(filter::addCategory)
+        return filter
+    }
+
+    fun removeReceiver(receiver: BroadcastReceiver) {
+        if (receiverDatas.removeAll { it.receiver == receiver } &&
+                receiverDatas.isEmpty() && registered) {
+            unregisterAction()
+            registered = false
+            activeCategories.clear()
+        }
+    }
+
+    @Throws(IllegalStateException::class)
+    override fun onReceive(context: Context, intent: Intent) {
+        if (intent.action != action) {
+            throw(IllegalStateException("Received intent for ${intent.action} " +
+                "in receiver for $action}"))
+        }
+        val id = index.getAndIncrement()
+        logger.logBroadcastReceived(id, userId, intent)
+        // Immediately return control to ActivityManager
+        bgExecutor.execute {
+            receiverDatas.forEach {
+                if (it.filter.matchCategories(intent.categories) == null) {
+                    it.executor.execute {
+                        it.receiver.pendingResult = pendingResult
+                        it.receiver.onReceive(context, intent)
+                        logger.logBroadcastDispatched(id, action, it.receiver)
+                    }
+                }
+            }
+        }
+    }
+
+    override fun dump(fd: FileDescriptor, pw: PrintWriter, args: Array<out String>) {
+        pw.indentIfPossible {
+            println("Registered: $registered")
+            println("Receivers:")
+            pw.indentIfPossible {
+                receiverDatas.forEach {
+                    println(it.receiver)
+                }
+            }
+            println("Categories: ${activeCategories.joinToString(", ")}")
+        }
+    }
+}

packages/SystemUI/src/com/android/systemui/broadcast/BroadcastDispatcher.kt

@@ -29,6 +29,7 @@ import android.os.UserHandle
 import android.text.TextUtils
 import android.util.SparseArray
 import com.android.internal.annotations.VisibleForTesting
+import com.android.internal.util.IndentingPrintWriter
 import com.android.systemui.Dumpable
 import com.android.systemui.broadcast.logging.BroadcastDispatcherLogger
 import com.android.systemui.dump.DumpManager
@@ -65,6 +66,7 @@ private const val DEBUG = true
 open class BroadcastDispatcher constructor (
     private val context: Context,
     private val bgLooper: Looper,
+    private val bgExecutor: Executor,
     private val dumpManager: DumpManager,
     private val logger: BroadcastDispatcherLogger
 ) : Dumpable, BroadcastReceiver() {
@@ -173,15 +175,18 @@ open class BroadcastDispatcher constructor (
 
     @VisibleForTesting
     protected open fun createUBRForUser(userId: Int) =
-            UserBroadcastDispatcher(context, userId, bgLooper, logger)
+            UserBroadcastDispatcher(context, userId, bgLooper, bgExecutor, logger)
 
     override fun dump(fd: FileDescriptor, pw: PrintWriter, args: Array<out String>) {
         pw.println("Broadcast dispatcher:")
-        pw.println("  Current user: ${handler.currentUser}")
+        val ipw = IndentingPrintWriter(pw, "  ")
+        ipw.increaseIndent()
+        ipw.println("Current user: ${handler.currentUser}")
         for (index in 0 until receiversByUser.size()) {
-            pw.println("  User ${receiversByUser.keyAt(index)}")
+            ipw.println("User ${receiversByUser.keyAt(index)}")
-            receiversByUser.valueAt(index).dump(fd, pw, args)
+            receiversByUser.valueAt(index).dump(fd, ipw, args)
         }
+        ipw.decreaseIndent()
     }
 
     private val handler = object : Handler(bgLooper) {

packages/SystemUI/src/com/android/systemui/broadcast/UserBroadcastDispatcher.kt

@@ -18,8 +18,6 @@ package com.android.systemui.broadcast
 
 import android.content.BroadcastReceiver
 import android.content.Context
-import android.content.Intent
-import android.content.IntentFilter
 import android.os.Handler
 import android.os.Looper
 import android.os.Message
@@ -31,11 +29,10 @@ import androidx.annotation.VisibleForTesting
 import com.android.internal.util.Preconditions
 import com.android.systemui.Dumpable
 import com.android.systemui.broadcast.logging.BroadcastDispatcherLogger
+import com.android.systemui.util.indentIfPossible
 import java.io.FileDescriptor
 import java.io.PrintWriter
-import java.lang.IllegalArgumentException
+import java.util.concurrent.Executor
-import java.lang.IllegalStateException
-import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicInteger
 
 private const val MSG_REGISTER_RECEIVER = 0
@@ -48,16 +45,14 @@ private const val DEBUG = false
  *
  * Created by [BroadcastDispatcher] as needed by users. The value of [userId] can be
  * [UserHandle.USER_ALL].
- *
- * Each instance of this class will register itself exactly once with [Context]. Updates to the
- * [IntentFilter] will be done in the background thread.
  */
-class UserBroadcastDispatcher(
+open class UserBroadcastDispatcher(
     private val context: Context,
     private val userId: Int,
     private val bgLooper: Looper,
+    private val bgExecutor: Executor,
     private val logger: BroadcastDispatcherLogger
-) : BroadcastReceiver(), Dumpable {
+) : Dumpable {
 
     companion object {
         // Used only for debugging. If not debugging, this variable will not be accessed and all
@@ -76,47 +71,16 @@ class UserBroadcastDispatcher(
         }
     }
 
-    private val registered = AtomicBoolean(false)
-
-    internal fun isRegistered() = registered.get()
-
     // Only modify in BG thread
-    private val actionsToReceivers = ArrayMap<String, MutableSet<ReceiverData>>()
+    @VisibleForTesting
-    private val receiverToReceiverData = ArrayMap<BroadcastReceiver, MutableSet<ReceiverData>>()
+    internal val actionsToActionsReceivers = ArrayMap<String, ActionReceiver>()
+    private val receiverToActions = ArrayMap<BroadcastReceiver, MutableSet<String>>()
 
     @VisibleForTesting
     internal fun isReceiverReferenceHeld(receiver: BroadcastReceiver): Boolean {
-        return receiverToReceiverData.contains(receiver) ||
+        return actionsToActionsReceivers.values.any {
-                actionsToReceivers.any {
+            it.hasReceiver(receiver)
-            it.value.any { it.receiver == receiver }
+        } || (receiver in receiverToActions)
-        }
-    }
-
-    // Only call on BG thread as it reads from the maps
-    private fun createFilter(): IntentFilter {
-        Preconditions.checkState(bgHandler.looper.isCurrentThread,
-                "This method should only be called from BG thread")
-        val categories = mutableSetOf<String>()
-        receiverToReceiverData.values.flatten().forEach {
-            it.filter.categoriesIterator()?.asSequence()?.let {
-                categories.addAll(it)
-            }
-        }
-        val intentFilter = IntentFilter().apply {
-            // The keys of the arrayMap are of type String! so null check is needed
-            actionsToReceivers.keys.forEach { if (it != null) addAction(it) else Unit }
-            categories.forEach { addCategory(it) }
-        }
-        return intentFilter
-    }
-
-    override fun onReceive(context: Context, intent: Intent) {
-        val id = index.getAndIncrement()
-        if (DEBUG) Log.w(TAG, "[$id] Received $intent")
-        logger.logBroadcastReceived(id, userId, intent)
-        bgHandler.post(
-                HandleBroadcastRunnable(
-                        actionsToReceivers, context, intent, pendingResult, id, logger))
     }
 
     /**
@@ -137,109 +101,57 @@ class UserBroadcastDispatcher(
         Preconditions.checkState(bgHandler.looper.isCurrentThread,
                 "This method should only be called from BG thread")
         if (DEBUG) Log.w(TAG, "Register receiver: ${receiverData.receiver}")
-        receiverToReceiverData.getOrPut(receiverData.receiver, { ArraySet() }).add(receiverData)
+        receiverToActions
-        var changed = false
+                .getOrPut(receiverData.receiver, { ArraySet() })
-        // Index the BroadcastReceiver by all its actions, that way it's easier to dispatch given
+                .addAll(receiverData.filter.actionsIterator()?.asSequence() ?: emptySequence())
-        // a received intent.
         receiverData.filter.actionsIterator().forEach {
-            actionsToReceivers.getOrPut(it) {
+            actionsToActionsReceivers
-                changed = true
+                    .getOrPut(it, { createActionReceiver(it) })
-                ArraySet()
+                    .addReceiverData(receiverData)
-            }.add(receiverData)
         }
         logger.logReceiverRegistered(userId, receiverData.receiver)
-        if (changed) {
+    }
-            createFilterAndRegisterReceiverBG()
+
-        }
+    @VisibleForTesting
+    internal open fun createActionReceiver(action: String): ActionReceiver {
+        return ActionReceiver(
+                action,
+                userId,
+                {
+                    context.registerReceiverAsUser(this, UserHandle.of(userId), it, null, bgHandler)
+                    logger.logContextReceiverRegistered(userId, it)
+                },
+                {
+                    try {
+                        context.unregisterReceiver(this)
+                        logger.logContextReceiverUnregistered(userId, action)
+                    } catch (e: IllegalArgumentException) {
+                        Log.e(TAG, "Trying to unregister unregistered receiver for user $userId, " +
+                                "action $action",
+                                IllegalStateException(e))
+                    }
+                },
+                bgExecutor,
+                logger
+        )
     }
 
     private fun handleUnregisterReceiver(receiver: BroadcastReceiver) {
         Preconditions.checkState(bgHandler.looper.isCurrentThread,
                 "This method should only be called from BG thread")
         if (DEBUG) Log.w(TAG, "Unregister receiver: $receiver")
-        val actions = receiverToReceiverData.getOrElse(receiver) { return }
+        receiverToActions.getOrDefault(receiver, mutableSetOf()).forEach {
-                .flatMap { it.filter.actionsIterator().asSequence().asIterable() }.toSet()
+            actionsToActionsReceivers.get(it)?.removeReceiver(receiver)
-        receiverToReceiverData.remove(receiver)?.clear()
-        var changed = false
-        actions.forEach { action ->
-            actionsToReceivers.get(action)?.removeIf { it.receiver == receiver }
-            if (actionsToReceivers.get(action)?.isEmpty() ?: false) {
-                changed = true
-                actionsToReceivers.remove(action)
-            }
         }
+        receiverToActions.remove(receiver)
         logger.logReceiverUnregistered(userId, receiver)
-        if (changed) {
-            createFilterAndRegisterReceiverBG()
-        }
-    }
-
-    // Only call this from a BG thread
-    private fun createFilterAndRegisterReceiverBG() {
-        val intentFilter = createFilter()
-        bgHandler.post(RegisterReceiverRunnable(intentFilter))
     }
 
     override fun dump(fd: FileDescriptor, pw: PrintWriter, args: Array<out String>) {
-        pw.println("  Registered=${registered.get()}")
+        pw.indentIfPossible {
-        actionsToReceivers.forEach { (action, list) ->
+            actionsToActionsReceivers.forEach { (action, actionReceiver) ->
-            pw.println("    $action:")
+                println("$action:")
-            list.forEach { pw.println("      ${it.receiver}") }
+                actionReceiver.dump(fd, pw, args)
-        }
-    }
-
-    private class HandleBroadcastRunnable(
-        val actionsToReceivers: Map<String, Set<ReceiverData>>,
-        val context: Context,
-        val intent: Intent,
-        val pendingResult: PendingResult,
-        val index: Int,
-        val logger: BroadcastDispatcherLogger
-    ) : Runnable {
-        override fun run() {
-            if (DEBUG) Log.w(TAG, "[$index] Dispatching $intent")
-            actionsToReceivers.get(intent.action)
-                    ?.filter {
-                        it.filter.hasAction(intent.action) &&
-                            it.filter.matchCategories(intent.categories) == null }
-                    ?.forEach {
-                        it.executor.execute {
-                            if (DEBUG) Log.w(TAG,
-                                    "[$index] Dispatching ${intent.action} to ${it.receiver}")
-                            logger.logBroadcastDispatched(index, intent.action, it.receiver)
-                            it.receiver.pendingResult = pendingResult
-                            it.receiver.onReceive(context, intent)
-                        }
-                    }
-        }
-    }
-
-    private inner class RegisterReceiverRunnable(val intentFilter: IntentFilter) : Runnable {
-
-        /*
-         * Registers and unregisters the BroadcastReceiver
-         */
-        override fun run() {
-            if (registered.get()) {
-                try {
-                    context.unregisterReceiver(this@UserBroadcastDispatcher)
-                    logger.logContextReceiverUnregistered(userId)
-                } catch (e: IllegalArgumentException) {
-                    Log.e(TAG, "Trying to unregister unregistered receiver for user $userId",
-                            IllegalStateException(e))
-                }
-                registered.set(false)
-            }
-            // Short interval without receiver, this can be problematic
-            if (intentFilter.countActions() > 0 && !registered.get()) {
-                context.registerReceiverAsUser(
-                        this@UserBroadcastDispatcher,
-                        UserHandle.of(userId),
-                        intentFilter,
-                        null,
-                        bgHandler)
-                registered.set(true)
-                logger.logContextReceiverRegistered(userId, intentFilter)
             }
         }
     }

packages/SystemUI/src/com/android/systemui/broadcast/logging/BroadcastDispatcherLogger.kt

@@ -99,11 +99,12 @@ class BroadcastDispatcherLogger @Inject constructor(
         })
     }
 
-    fun logContextReceiverUnregistered(user: Int) {
+    fun logContextReceiverUnregistered(user: Int, action: String) {
         log(INFO, {
             int1 = user
+            str1 = action
         }, {
-            "Receiver unregistered with Context for user $int1."
+            "Receiver unregistered with Context for user $int1, action $str1"
         })
     }
 

packages/SystemUI/src/com/android/systemui/dagger/DependencyProvider.java

@@ -61,6 +61,8 @@ import com.android.systemui.statusbar.policy.DataSaverController;
 import com.android.systemui.statusbar.policy.NetworkController;
 import com.android.systemui.util.leak.LeakDetector;
 
+import java.util.concurrent.Executor;
+
 import javax.inject.Named;
 import javax.inject.Singleton;
 
@@ -188,11 +190,12 @@ public class DependencyProvider {
     public BroadcastDispatcher providesBroadcastDispatcher(
             Context context,
             @Background Looper backgroundLooper,
+            @Background Executor backgroundExecutor,
             DumpManager dumpManager,
             BroadcastDispatcherLogger logger
     ) {
-        BroadcastDispatcher bD =
+        BroadcastDispatcher bD = new BroadcastDispatcher(context, backgroundLooper,
-                new BroadcastDispatcher(context, backgroundLooper, dumpManager, logger);
+                backgroundExecutor, dumpManager, logger);
         bD.initialize();
         return bD;
     }

packages/SystemUI/src/com/android/systemui/util/ConvenienceExtensions.kt

@@ -17,6 +17,8 @@
 package com.android.systemui.util
 
 import android.view.ViewGroup
+import com.android.internal.util.IndentingPrintWriter
+import java.io.PrintWriter
 
 /** [Sequence] that yields all of the direct children of this [ViewGroup] */
 val ViewGroup.children
@@ -33,3 +35,14 @@ fun <T> Sequence<T>.takeUntil(pred: (T) -> Boolean): Sequence<T> = sequence {
         }
     }
 }
+
+/**
+ * If `this` is an [IndentingPrintWriter], it will process block inside an indentation level.
+ *
+ * If not, this will just process block.
+ */
+inline fun PrintWriter.indentIfPossible(block: PrintWriter.() -> Unit) {
+    if (this is IndentingPrintWriter) increaseIndent()
+    block()
+    if (this is IndentingPrintWriter) decreaseIndent()
+}

packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java

@@ -49,6 +49,7 @@ import org.junit.Rule;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.Future;
 
 /**
@@ -74,7 +75,8 @@ public abstract class SysuiTestCase {
         SystemUIFactory.createFromConfig(mContext);
         mDependency = new TestableDependency(mContext);
         mFakeBroadcastDispatcher = new FakeBroadcastDispatcher(mContext, mock(Looper.class),
-                mock(DumpManager.class), mock(BroadcastDispatcherLogger.class));
+                mock(Executor.class), mock(DumpManager.class),
+                mock(BroadcastDispatcherLogger.class));
 
         mRealInstrumentation = InstrumentationRegistry.getInstrumentation();
         Instrumentation inst = spy(mRealInstrumentation);

packages/SystemUI/tests/src/com/android/systemui/broadcast/ActionReceiverTest.kt

@@ -0,0 +1,256 @@
+/*
+ * Copyright (C) 2020 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.systemui.broadcast
+
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.os.UserHandle
+import android.test.suitebuilder.annotation.SmallTest
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.broadcast.logging.BroadcastDispatcherLogger
+import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.capture
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.time.FakeSystemClock
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.ArgumentMatchers.anyString
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.Mockito
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+import java.lang.IllegalArgumentException
+import java.lang.IllegalStateException
+import java.util.concurrent.Executor
+
+@RunWith(AndroidTestingRunner::class)
+@TestableLooper.RunWithLooper
+@SmallTest
+class ActionReceiverTest : SysuiTestCase() {
+
+    companion object {
+        private const val ACTION1 = "TEST_ACTION1"
+        private const val ACTION2 = "TEST_ACTION2"
+        private const val CATEGORY = "TEST_CATEGORY"
+        private val USER = UserHandle.of(0)
+        private fun <T : Any> sameNotNull(arg: T): T = Mockito.same(arg) ?: arg
+
+        fun IntentFilter.matchesOther(it: IntentFilter): Boolean {
+            val actions = actionsIterator()?.asSequence()?.toSet() ?: emptySet()
+            val categories = categoriesIterator()?.asSequence()?.toSet() ?: emptySet()
+            return (it.actionsIterator()?.asSequence()?.toSet() ?: emptySet()) == actions &&
+                    (it.categoriesIterator()?.asSequence()?.toSet() ?: emptySet()) == categories &&
+                    it.countDataAuthorities() == 0 &&
+                    it.countDataPaths() == 0 &&
+                    it.countDataSchemes() == 0 &&
+                    it.countDataTypes() == 0 &&
+                    it.countMimeGroups() == 0 &&
+                    it.priority == 0
+        }
+    }
+
+    @Mock
+    private lateinit var registerFunction: BroadcastReceiver.(IntentFilter) -> Unit
+    @Mock
+    private lateinit var unregisterFunction: BroadcastReceiver.() -> Unit
+    @Mock
+    private lateinit var receiver1: BroadcastReceiver
+    @Mock
+    private lateinit var receiver2: BroadcastReceiver
+    @Mock
+    private lateinit var logger: BroadcastDispatcherLogger
+    @Captor
+    private lateinit var intentFilterCaptor: ArgumentCaptor<IntentFilter>
+
+    private lateinit var executor: FakeExecutor
+    private lateinit var actionReceiver: ActionReceiver
+    private val directExecutor = Executor { it.run() }
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        executor = FakeExecutor(FakeSystemClock())
+
+        actionReceiver = ActionReceiver(
+                ACTION1,
+                USER.identifier,
+                registerFunction,
+                unregisterFunction,
+                executor,
+                logger
+        )
+    }
+
+    @Test
+    fun testStartsUnregistered() {
+        assertFalse(actionReceiver.registered)
+        verify(registerFunction, never()).invoke(sameNotNull(actionReceiver),
+                any(IntentFilter::class.java))
+    }
+
+    @Test
+    fun testRegistersOnFirstAdd() {
+        val receiverData = ReceiverData(receiver1, IntentFilter(ACTION1), directExecutor, USER)
+
+        actionReceiver.addReceiverData(receiverData)
+
+        assertTrue(actionReceiver.registered)
+        verify(registerFunction).invoke(sameNotNull(actionReceiver), capture(intentFilterCaptor))
+
+        assertTrue(IntentFilter(ACTION1).matchesOther(intentFilterCaptor.value))
+    }
+
+    @Test
+    fun testRegistersOnlyOnce() {
+        val receiverData1 = ReceiverData(receiver1, IntentFilter(ACTION1), directExecutor, USER)
+        val receiverData2 = ReceiverData(receiver2, IntentFilter(ACTION1), directExecutor, USER)
+
+        actionReceiver.addReceiverData(receiverData1)
+        actionReceiver.addReceiverData(receiverData2)
+
+        verify(registerFunction).invoke(sameNotNull(actionReceiver), any(IntentFilter::class.java))
+    }
+
+    @Test
+    fun testRemovingLastReceiverUnregisters() {
+        val receiverData = ReceiverData(receiver1, IntentFilter(ACTION1), directExecutor, USER)
+
+        actionReceiver.addReceiverData(receiverData)
+
+        actionReceiver.removeReceiver(receiver1)
+
+        assertFalse(actionReceiver.registered)
+        verify(unregisterFunction).invoke(sameNotNull(actionReceiver))
+    }
+
+    @Test
+    fun testRemovingWhileOtherReceiversDoesntUnregister() {
+        val receiverData1 = ReceiverData(receiver1, IntentFilter(ACTION1), directExecutor, USER)
+        val receiverData2 = ReceiverData(receiver2, IntentFilter(ACTION1), directExecutor, USER)
+
+        actionReceiver.addReceiverData(receiverData1)
+        actionReceiver.addReceiverData(receiverData2)
+
+        actionReceiver.removeReceiver(receiver1)
+
+        assertTrue(actionReceiver.registered)
+        verify(unregisterFunction, never()).invoke(any(BroadcastReceiver::class.java))
+    }
+
+    @Test
+    fun testReceiverHasCategories() {
+        val filter = IntentFilter(ACTION1)
+        filter.addCategory(CATEGORY)
+
+        val receiverData = ReceiverData(receiver1, filter, directExecutor, USER)
+
+        actionReceiver.addReceiverData(receiverData)
+
+        verify(registerFunction).invoke(sameNotNull(actionReceiver), capture(intentFilterCaptor))
+        assertTrue(intentFilterCaptor.value.hasCategory(CATEGORY))
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun testNotRegisteredWithWrongAction_throwsException() {
+        val receiverData = ReceiverData(receiver1, IntentFilter(ACTION2), directExecutor, USER)
+
+        actionReceiver.addReceiverData(receiverData)
+    }
+
+    @Test
+    fun testReceiverGetsBroadcast() {
+        val receiverData = ReceiverData(receiver1, IntentFilter(ACTION1), directExecutor, USER)
+        actionReceiver.addReceiverData(receiverData)
+
+        val intent = Intent(ACTION1)
+
+        actionReceiver.onReceive(mContext, intent)
+
+        executor.runAllReady()
+
+        verify(receiver1).onReceive(any(Context::class.java), sameNotNull(intent))
+    }
+
+    @Test
+    fun testReceiverGetsPendingResult() {
+        val receiverData = ReceiverData(receiver1, IntentFilter(ACTION1), directExecutor, USER)
+        actionReceiver.addReceiverData(receiverData)
+
+        val intent = Intent(ACTION1)
+        val pendingResult = mock(BroadcastReceiver.PendingResult::class.java)
+
+        actionReceiver.pendingResult = pendingResult
+        actionReceiver.onReceive(mContext, intent)
+
+        executor.runAllReady()
+        verify(receiver1).pendingResult = pendingResult
+    }
+
+    @Test
+    fun testBroadcastIsDispatchedInExecutor() {
+        val executor = FakeExecutor(FakeSystemClock())
+        val receiverData = ReceiverData(receiver1, IntentFilter(ACTION1), executor, USER)
+        actionReceiver.addReceiverData(receiverData)
+
+        val intent = Intent(ACTION1)
+        actionReceiver.onReceive(mContext, intent)
+
+        this.executor.runAllReady()
+
+        verify(receiver1, never()).onReceive(mContext, intent)
+
+        executor.runAllReady()
+        // Dispatched after executor is processed
+        verify(receiver1).onReceive(mContext, intent)
+    }
+
+    @Test
+    fun testBroadcastReceivedDispatched_logger() {
+        val receiverData = ReceiverData(receiver1, IntentFilter(ACTION1), directExecutor, USER)
+
+        actionReceiver.addReceiverData(receiverData)
+
+        val intent = Intent(ACTION1)
+        actionReceiver.onReceive(mContext, intent)
+        verify(logger).logBroadcastReceived(anyInt(), eq(USER.identifier), eq(intent))
+
+        verify(logger, never()).logBroadcastDispatched(anyInt(), anyString(),
+                any(BroadcastReceiver::class.java))
+
+        executor.runAllReady()
+
+        verify(logger).logBroadcastDispatched(anyInt(), eq(ACTION1), sameNotNull(receiver1))
+    }
+
+    @Test(expected = IllegalStateException::class)
+    fun testBroadcastWithWrongAction_throwsException() {
+        actionReceiver.onReceive(mContext, Intent(ACTION2))
+    }
+}

packages/SystemUI/tests/src/com/android/systemui/broadcast/BroadcastDispatcherTest.kt

@@ -98,6 +98,7 @@ class BroadcastDispatcherTest : SysuiTestCase() {
         broadcastDispatcher = TestBroadcastDispatcher(
                 mockContext,
                 testableLooper.looper,
+                mock(Executor::class.java),
                 mock(DumpManager::class.java),
                 logger,
                 mapOf(0 to mockUBRUser0, 1 to mockUBRUser1))
@@ -246,10 +247,11 @@ class BroadcastDispatcherTest : SysuiTestCase() {
     private class TestBroadcastDispatcher(
         context: Context,
         bgLooper: Looper,
+        executor: Executor,
         dumpManager: DumpManager,
         logger: BroadcastDispatcherLogger,
         var mockUBRMap: Map<Int, UserBroadcastDispatcher>
-    ) : BroadcastDispatcher(context, bgLooper, dumpManager, logger) {
+    ) : BroadcastDispatcher(context, bgLooper, executor, dumpManager, logger) {
         override fun createUBRForUser(userId: Int): UserBroadcastDispatcher {
             return mockUBRMap.getOrDefault(userId, mock(UserBroadcastDispatcher::class.java))
         }

packages/SystemUI/tests/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt

@@ -31,9 +31,10 @@ import java.util.concurrent.Executor
 class FakeBroadcastDispatcher(
     context: SysuiTestableContext,
     looper: Looper,
+    executor: Executor,
     dumpManager: DumpManager,
     logger: BroadcastDispatcherLogger
-) : BroadcastDispatcher(context, looper, dumpManager, logger) {
+) : BroadcastDispatcher(context, looper, executor, dumpManager, logger) {
 
     private val registeredReceivers = ArraySet<BroadcastReceiver>()
 

packages/SystemUI/tests/src/com/android/systemui/broadcast/UserBroadcastDispatcherTest.kt

@@ -18,7 +18,6 @@ package com.android.systemui.broadcast
 
 import android.content.BroadcastReceiver
 import android.content.Context
-import android.content.Intent
 import android.content.IntentFilter
 import android.os.Handler
 import android.os.UserHandle
@@ -29,26 +28,18 @@ import com.android.systemui.SysuiTestCase
 import com.android.systemui.broadcast.logging.BroadcastDispatcherLogger
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.time.FakeSystemClock
-import junit.framework.Assert.assertEquals
 import junit.framework.Assert.assertFalse
-import junit.framework.Assert.assertTrue
+import junit.framework.Assert.assertNotNull
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
-import org.mockito.ArgumentMatchers.any
-import org.mockito.ArgumentMatchers.anyInt
-import org.mockito.ArgumentMatchers.eq
-import org.mockito.Captor
 import org.mockito.Mock
 import org.mockito.Mockito
-import org.mockito.Mockito.anyString
+import org.mockito.Mockito.mock
-import org.mockito.Mockito.atLeastOnce
-import org.mockito.Mockito.never
-import org.mockito.Mockito.reset
-import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import java.util.concurrent.Executor
 
 @RunWith(AndroidTestingRunner::class)
 @TestableLooper.RunWithLooper
@@ -58,8 +49,6 @@ class UserBroadcastDispatcherTest : SysuiTestCase() {
     companion object {
         private const val ACTION_1 = "com.android.systemui.tests.ACTION_1"
         private const val ACTION_2 = "com.android.systemui.tests.ACTION_2"
-        private const val CATEGORY_1 = "com.android.systemui.tests.CATEGORY_1"
-        private const val CATEGORY_2 = "com.android.systemui.tests.CATEGORY_2"
         private const val USER_ID = 0
         private val USER_HANDLE = UserHandle.of(USER_ID)
 
@@ -75,13 +64,8 @@ class UserBroadcastDispatcherTest : SysuiTestCase() {
     @Mock
     private lateinit var mockContext: Context
     @Mock
-    private lateinit var mPendingResult: BroadcastReceiver.PendingResult
-    @Mock
     private lateinit var logger: BroadcastDispatcherLogger
 
-    @Captor
-    private lateinit var argumentCaptor: ArgumentCaptor<IntentFilter>
-
     private lateinit var testableLooper: TestableLooper
     private lateinit var userBroadcastDispatcher: UserBroadcastDispatcher
     private lateinit var intentFilter: IntentFilter
@@ -96,46 +80,25 @@ class UserBroadcastDispatcherTest : SysuiTestCase() {
         handler = Handler(testableLooper.looper)
         fakeExecutor = FakeExecutor(FakeSystemClock())
 
-        userBroadcastDispatcher = UserBroadcastDispatcher(
+        userBroadcastDispatcher = object : UserBroadcastDispatcher(
-                mockContext, USER_ID, testableLooper.looper, logger)
+                mockContext, USER_ID, testableLooper.looper, mock(Executor::class.java), logger) {
-        userBroadcastDispatcher.pendingResult = mPendingResult
+            override fun createActionReceiver(action: String): ActionReceiver {
-    }
+                return mock(ActionReceiver::class.java)
-
+            }
-    @Test
+        }
-    fun testNotRegisteredOnStart() {
-        testableLooper.processAllMessages()
-        verify(mockContext, never()).registerReceiver(any(), any())
-        verify(mockContext, never()).registerReceiver(any(), any(), anyInt())
-        verify(mockContext, never()).registerReceiver(any(), any(), anyString(), any())
-        verify(mockContext, never()).registerReceiver(any(), any(), anyString(), any(), anyInt())
-        verify(mockContext, never()).registerReceiverAsUser(any(), any(), any(), anyString(), any())
-    }
-
-    @Test
-    fun testNotRegisteredOnStart_logging() {
-        testableLooper.processAllMessages()
-
-        verify(logger, never()).logContextReceiverRegistered(anyInt(), any())
     }
 
     @Test
     fun testSingleReceiverRegistered() {
         intentFilter = IntentFilter(ACTION_1)
+        val receiverData = ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE)
 
-        userBroadcastDispatcher.registerReceiver(
+        userBroadcastDispatcher.registerReceiver(receiverData)
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
         testableLooper.processAllMessages()
 
-        assertTrue(userBroadcastDispatcher.isRegistered())
+        val actionReceiver = userBroadcastDispatcher.getActionReceiver(ACTION_1)
-        verify(mockContext).registerReceiverAsUser(
+        assertNotNull(actionReceiver)
-                any(),
+        verify(actionReceiver)?.addReceiverData(receiverData)
-                eq(USER_HANDLE),
-                capture(argumentCaptor),
-                any(),
-                any())
-        assertEquals(1, argumentCaptor.value.countActions())
-        assertTrue(argumentCaptor.value.hasAction(ACTION_1))
-        assertEquals(0, argumentCaptor.value.countCategories())
     }
 
     @Test
@@ -147,7 +110,6 @@ class UserBroadcastDispatcherTest : SysuiTestCase() {
         testableLooper.processAllMessages()
 
         verify(logger).logReceiverRegistered(USER_HANDLE.identifier, broadcastReceiver)
-        verify(logger).logContextReceiverRegistered(eq(USER_HANDLE.identifier), any())
     }
 
     @Test
@@ -157,16 +119,13 @@ class UserBroadcastDispatcherTest : SysuiTestCase() {
         userBroadcastDispatcher.registerReceiver(
                 ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
         testableLooper.processAllMessages()
-        reset(mockContext)
-
-        assertTrue(userBroadcastDispatcher.isRegistered())
 
         userBroadcastDispatcher.unregisterReceiver(broadcastReceiver)
         testableLooper.processAllMessages()
 
-        verify(mockContext, atLeastOnce()).unregisterReceiver(any())
+        val actionReceiver = userBroadcastDispatcher.getActionReceiver(ACTION_1)
-        verify(mockContext, never()).registerReceiverAsUser(any(), any(), any(), any(), any())
+        assertNotNull(actionReceiver)
-        assertFalse(userBroadcastDispatcher.isRegistered())
+        verify(actionReceiver)?.removeReceiver(broadcastReceiver)
     }
 
     @Test
@@ -181,139 +140,6 @@ class UserBroadcastDispatcherTest : SysuiTestCase() {
         testableLooper.processAllMessages()
 
         verify(logger).logReceiverUnregistered(USER_HANDLE.identifier, broadcastReceiver)
-        verify(logger).logContextReceiverUnregistered(USER_HANDLE.identifier)
-    }
-
-    @Test
-    fun testFilterHasAllActionsAndCategories_twoReceivers() {
-        intentFilter = IntentFilter(ACTION_1)
-        intentFilterOther = IntentFilter(ACTION_2).apply {
-            addCategory(CATEGORY_1)
-            addCategory(CATEGORY_2)
-        }
-
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiverOther, intentFilterOther, fakeExecutor, USER_HANDLE))
-
-        testableLooper.processAllMessages()
-        assertTrue(userBroadcastDispatcher.isRegistered())
-
-        verify(mockContext, times(2)).registerReceiverAsUser(
-                any(),
-                eq(USER_HANDLE),
-                capture(argumentCaptor),
-                any(),
-                any())
-
-        val lastFilter = argumentCaptor.value
-
-        assertTrue(lastFilter.hasAction(ACTION_1))
-        assertTrue(lastFilter.hasAction(ACTION_2))
-        assertTrue(lastFilter.hasCategory(CATEGORY_1))
-        assertTrue(lastFilter.hasCategory(CATEGORY_1))
-    }
-
-    @Test
-    fun testDispatchToCorrectReceiver() {
-        intentFilter = IntentFilter(ACTION_1)
-        intentFilterOther = IntentFilter(ACTION_2)
-
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiverOther, intentFilterOther, fakeExecutor, USER_HANDLE))
-
-        val intent = Intent(ACTION_2)
-
-        userBroadcastDispatcher.onReceive(mockContext, intent)
-        testableLooper.processAllMessages()
-        fakeExecutor.runAllReady()
-
-        verify(broadcastReceiver, never()).onReceive(any(), any())
-        verify(broadcastReceiverOther).onReceive(mockContext, intent)
-    }
-
-    @Test
-    fun testDispatch_logger() {
-        intentFilter = IntentFilter(ACTION_1)
-        intentFilterOther = IntentFilter(ACTION_2)
-
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiverOther, intentFilterOther, fakeExecutor, USER_HANDLE))
-
-        val intent = Intent(ACTION_2)
-
-        userBroadcastDispatcher.onReceive(mockContext, intent)
-        testableLooper.processAllMessages()
-        fakeExecutor.runAllReady()
-
-        val captor = ArgumentCaptor.forClass(Int::class.java)
-        verify(logger)
-                .logBroadcastReceived(captor.capture(), eq(USER_HANDLE.identifier), eq(intent))
-        verify(logger).logBroadcastDispatched(captor.value, ACTION_2, broadcastReceiverOther)
-        verify(logger, never())
-                .logBroadcastDispatched(eq(captor.value), any(), eq(broadcastReceiver))
-    }
-
-    @Test
-    fun testDispatchToCorrectReceiver_differentFiltersSameReceiver() {
-        intentFilter = IntentFilter(ACTION_1)
-        intentFilterOther = IntentFilter(ACTION_2)
-
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilterOther, fakeExecutor, USER_HANDLE))
-
-        val intent = Intent(ACTION_2)
-
-        userBroadcastDispatcher.onReceive(mockContext, intent)
-        testableLooper.processAllMessages()
-        fakeExecutor.runAllReady()
-
-        verify(broadcastReceiver).onReceive(mockContext, intent)
-    }
-
-    @Test
-    fun testDispatchIntentWithoutCategories() {
-        intentFilter = IntentFilter(ACTION_1)
-        intentFilter.addCategory(CATEGORY_1)
-        intentFilterOther = IntentFilter(ACTION_1)
-        intentFilterOther.addCategory(CATEGORY_2)
-
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiverOther, intentFilterOther, fakeExecutor, USER_HANDLE))
-
-        val intent = Intent(ACTION_1)
-
-        userBroadcastDispatcher.onReceive(mockContext, intent)
-        testableLooper.processAllMessages()
-        fakeExecutor.runAllReady()
-
-        verify(broadcastReceiver).onReceive(mockContext, intent)
-        verify(broadcastReceiverOther).onReceive(mockContext, intent)
-    }
-
-    @Test
-    fun testPendingResult() {
-        intentFilter = IntentFilter(ACTION_1)
-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
-
-        val intent = Intent(ACTION_1)
-        userBroadcastDispatcher.onReceive(mockContext, intent)
-
-        testableLooper.processAllMessages()
-        fakeExecutor.runAllReady()
-
-        verify(broadcastReceiver).onReceive(mockContext, intent)
-        verify(broadcastReceiver).pendingResult = mPendingResult
     }
 
     @Test
@@ -333,4 +159,8 @@ class UserBroadcastDispatcherTest : SysuiTestCase() {
 
         assertFalse(userBroadcastDispatcher.isReceiverReferenceHeld(broadcastReceiver))
     }
+
+    private fun UserBroadcastDispatcher.getActionReceiver(action: String): ActionReceiver? {
+        return actionsToActionsReceivers.get(action)
+    }
 }