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
2020-06-19 16:57:33 +00:00
parent dfaa7932a2 d6b99a586a
commit 73d6f232f8
11 changed files with 496 additions and 338 deletions
--- a/packages/SystemUI/src/com/android/systemui/broadcast/ActionReceiver.kt
+++ b/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(", ")}")
+        }
+    }
+}
--- a/packages/SystemUI/src/com/android/systemui/broadcast/BroadcastDispatcher.kt
+++ b/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)}")
-            receiversByUser.valueAt(index).dump(fd, pw, args)
+            ipw.println("User ${receiversByUser.keyAt(index)}")
+            receiversByUser.valueAt(index).dump(fd, ipw, args)
        }
+        ipw.decreaseIndent()
    }

    private val handler = object : Handler(bgLooper) {
--- a/packages/SystemUI/src/com/android/systemui/broadcast/UserBroadcastDispatcher.kt
+++ b/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.lang.IllegalStateException
-import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.Executor
 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>>()
-    private val receiverToReceiverData = ArrayMap<BroadcastReceiver, MutableSet<ReceiverData>>()
+    @VisibleForTesting
+    internal val actionsToActionsReceivers = ArrayMap<String, ActionReceiver>()
+    private val receiverToActions = ArrayMap<BroadcastReceiver, MutableSet<String>>()

    @VisibleForTesting
    internal fun isReceiverReferenceHeld(receiver: BroadcastReceiver): Boolean {
-        return receiverToReceiverData.contains(receiver) ||
-                actionsToReceivers.any {
-            it.value.any { it.receiver == receiver }
-        }
-    }
-
-    // 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))
+        return actionsToActionsReceivers.values.any {
+            it.hasReceiver(receiver)
+        } || (receiver in receiverToActions)
    }

    /**
@@ -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)
-        var changed = false
-        // Index the BroadcastReceiver by all its actions, that way it's easier to dispatch given
-        // a received intent.
+        receiverToActions
+                .getOrPut(receiverData.receiver, { ArraySet() })
+                .addAll(receiverData.filter.actionsIterator()?.asSequence() ?: emptySequence())
        receiverData.filter.actionsIterator().forEach {
-            actionsToReceivers.getOrPut(it) {
-                changed = true
-                ArraySet()
-            }.add(receiverData)
+            actionsToActionsReceivers
+                    .getOrPut(it, { createActionReceiver(it) })
+                    .addReceiverData(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 }
-                .flatMap { it.filter.actionsIterator().asSequence().asIterable() }.toSet()
-        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.getOrDefault(receiver, mutableSetOf()).forEach {
+            actionsToActionsReceivers.get(it)?.removeReceiver(receiver)
        }
+        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()}")
-        actionsToReceivers.forEach { (action, list) ->
-            pw.println("    $action:")
-            list.forEach { pw.println("      ${it.receiver}") }
-        }
-    }
-
-    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)
+        pw.indentIfPossible {
+            actionsToActionsReceivers.forEach { (action, actionReceiver) ->
+                println("$action:")
+                actionReceiver.dump(fd, pw, args)
            }
        }
    }
--- a/packages/SystemUI/src/com/android/systemui/broadcast/logging/BroadcastDispatcherLogger.kt
+++ b/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"
        })
    }

--- a/packages/SystemUI/src/com/android/systemui/dagger/DependencyProvider.java
+++ b/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 =
-                new BroadcastDispatcher(context, backgroundLooper, dumpManager, logger);
+        BroadcastDispatcher bD = new BroadcastDispatcher(context, backgroundLooper,
+                backgroundExecutor, dumpManager, logger);
        bD.initialize();
        return bD;
    }
--- a/packages/SystemUI/src/com/android/systemui/util/ConvenienceExtensions.kt
+++ b/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
@@ -32,4 +34,15 @@ fun <T> Sequence<T>.takeUntil(pred: (T) -> Boolean): Sequence<T> = sequence {
            break
        }
    }
+}
+
+/**
+ * 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()
 }
--- a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
+++ b/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);
--- a/packages/SystemUI/tests/src/com/android/systemui/broadcast/ActionReceiverTest.kt
+++ b/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))
+    }
+}
--- a/packages/SystemUI/tests/src/com/android/systemui/broadcast/BroadcastDispatcherTest.kt
+++ b/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))
        }
--- a/packages/SystemUI/tests/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt
+++ b/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>()

--- a/packages/SystemUI/tests/src/com/android/systemui/broadcast/UserBroadcastDispatcherTest.kt
+++ b/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.atLeastOnce
-import org.mockito.Mockito.never
-import org.mockito.Mockito.reset
-import org.mockito.Mockito.times
+import org.mockito.Mockito.mock
 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(
-                mockContext, USER_ID, testableLooper.looper, logger)
-        userBroadcastDispatcher.pendingResult = mPendingResult
-    }
-
-    @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())
+        userBroadcastDispatcher = object : UserBroadcastDispatcher(
+                mockContext, USER_ID, testableLooper.looper, mock(Executor::class.java), logger) {
+            override fun createActionReceiver(action: String): ActionReceiver {
+                return mock(ActionReceiver::class.java)
+            }
+        }
    }

    @Test
    fun testSingleReceiverRegistered() {
        intentFilter = IntentFilter(ACTION_1)
+        val receiverData = ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE)

-        userBroadcastDispatcher.registerReceiver(
-                ReceiverData(broadcastReceiver, intentFilter, fakeExecutor, USER_HANDLE))
+        userBroadcastDispatcher.registerReceiver(receiverData)
        testableLooper.processAllMessages()

-        assertTrue(userBroadcastDispatcher.isRegistered())
-        verify(mockContext).registerReceiverAsUser(
-                any(),
-                eq(USER_HANDLE),
-                capture(argumentCaptor),
-                any(),
-                any())
-        assertEquals(1, argumentCaptor.value.countActions())
-        assertTrue(argumentCaptor.value.hasAction(ACTION_1))
-        assertEquals(0, argumentCaptor.value.countCategories())
+        val actionReceiver = userBroadcastDispatcher.getActionReceiver(ACTION_1)
+        assertNotNull(actionReceiver)
+        verify(actionReceiver)?.addReceiverData(receiverData)
    }

    @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())
-        verify(mockContext, never()).registerReceiverAsUser(any(), any(), any(), any(), any())
-        assertFalse(userBroadcastDispatcher.isRegistered())
+        val actionReceiver = userBroadcastDispatcher.getActionReceiver(ACTION_1)
+        assertNotNull(actionReceiver)
+        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)
+    }
 }