Change BroadcastDispatcher to 1receiver/action

With the aggregated filter, as it had to be registered when any single
action changed (added for first time or removed the last receiver for
that action), if there was any action that got periodically
registered/unregistered, it would cause frequent cycles of
unregister/register with Context. This leads to interruption of service.

With this change, changes in one action would not affect receivers for
the others.

Test: manual, SystemUI works as expected
Test: manual, dumps and logs
Test: atest UserBroadcastDispatcherTest ActionReceiverTest
Bug: 148390204
Bug: 157165818

Change-Id: I7961f6c60600dfa0407615319c43bb72e43b4795

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)}")
-            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) {

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)
             }
         }
     }

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 =
-                new BroadcastDispatcher(context, backgroundLooper, dumpManager, logger);
+        BroadcastDispatcher bD = new BroadcastDispatcher(context, backgroundLooper,
+                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
@@ -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()
 }

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.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)
+    }
 }