Merge "Controls API - Publisher change - Phase 1"

11 files changed, 294 insertions, 246 deletions

diff --git a/api/current.txt b/api/current.txt
index 47d203ea9863..20933e4ad762 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -43404,11 +43404,12 @@ package android.service.controls {
 
   public abstract class ControlsProviderService extends android.app.Service {
     ctor public ControlsProviderService();
-    method public abstract void loadAvailableControls(@NonNull java.util.function.Consumer<java.util.List<android.service.controls.Control>>);
-    method public void loadSuggestedControls(int, @NonNull java.util.function.Consumer<java.util.List<android.service.controls.Control>>);
+    method @Deprecated public void loadAvailableControls(@NonNull java.util.function.Consumer<java.util.List<android.service.controls.Control>>);
     method @NonNull public final android.os.IBinder onBind(@NonNull android.content.Intent);
     method public abstract void performControlAction(@NonNull String, @NonNull android.service.controls.actions.ControlAction, @NonNull java.util.function.Consumer<java.lang.Integer>);
     method @NonNull public abstract java.util.concurrent.Flow.Publisher<android.service.controls.Control> publisherFor(@NonNull java.util.List<java.lang.String>);
+    method @Nullable public java.util.concurrent.Flow.Publisher<android.service.controls.Control> publisherForAllAvailable();
+    method @Nullable public java.util.concurrent.Flow.Publisher<android.service.controls.Control> publisherForSuggested();
     field public static final String SERVICE_CONTROLS = "android.service.controls.ControlsProviderService";
     field @NonNull public static final String TAG = "ControlsProviderService";
   }
diff --git a/core/java/android/service/controls/ControlsProviderService.java b/core/java/android/service/controls/ControlsProviderService.java
index de4c056e5501..cb20db90f549 100644
--- a/core/java/android/service/controls/ControlsProviderService.java
+++ b/core/java/android/service/controls/ControlsProviderService.java
@@ -16,6 +16,7 @@
 package android.service.controls;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.SdkConstant;
 import android.annotation.SdkConstant.SdkConstantType;
 import android.app.Service;
@@ -34,7 +35,6 @@ import android.util.Log;
 
 import com.android.internal.util.Preconditions;
 
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.Flow.Publisher;
@@ -70,25 +70,48 @@ public abstract class ControlsProviderService extends Service {
      * Retrieve all available controls, using the stateless builder
      * {@link Control.StatelessBuilder} to build each Control, then use the
      * provided consumer to callback to the call originator.
+     *
+     * @deprecated Removing consumer-based load apis. Use publisherForAllAvailable() instead
      */
-    public abstract void loadAvailableControls(@NonNull Consumer<List<Control>> consumer);
+    @Deprecated
+    public void loadAvailableControls(@NonNull Consumer<List<Control>> consumer) {
+        // pending removal
+        consumer.accept(Collections.emptyList());
+    }
+
+    /**
+     * Publisher for all available controls
+     *
+     * Retrieve all available controls. Use the stateless builder {@link Control.StatelessBuilder}
+     * to build each Control. Call {@link Subscriber#onComplete} when done loading all unique
+     * controls, or {@link Subscriber#onError} for error scenarios. Duplicate Controls will
+     * replace the original.
+     */
+    @Nullable
+    public Publisher<Control> publisherForAllAvailable() {
+        // will be abstract and @nonnull when consumers are removed
+        return null;
+    }
 
     /**
-     * (Optional) The service may be asked to provide a small number of recommended controls, in
+     * (Optional) Publisher for suggested controls
+     *
+     * The service may be asked to provide a small number of recommended controls, in
      * order to suggest some controls to the user for favoriting. The controls shall be built using
-     * the stateless builder {@link Control.StatelessBuilder}, followed by an invocation to the
-     * provided consumer to callback to the call originator. If the number of controls
-     * is greater than maxNumber, the list will be truncated.
+     * the stateless builder {@link Control.StatelessBuilder}. The number of controls requested
+     * through {@link Subscription#request} will be limited. Call {@link Subscriber#onComplete}
+     * when done, or {@link Subscriber#onError} for error scenarios.
      */
-    public void loadSuggestedControls(int maxNumber, @NonNull Consumer<List<Control>> consumer) {
-        // Override to change the default behavior
-        consumer.accept(Collections.emptyList());
+    @Nullable
+    public Publisher<Control> publisherForSuggested() {
+        return null;
     }
 
     /**
-     * Return a valid Publisher for the given controlIds. This publisher will be asked
-     * to provide updates for the given list of controlIds as long as the Subscription
-     * is valid.
+     * Return a valid Publisher for the given controlIds. This publisher will be asked to provide
+     * updates for the given list of controlIds as long as the {@link Subscription} is valid.
+     * Calls to {@link Subscriber#onComplete} will not be expected. Instead, wait for the call from
+     * {@link Subscription#cancel} to indicate that updates are no longer required.
      */
     @NonNull
     public abstract Publisher<Control> publisherFor(@NonNull List<String> controlIds);
@@ -113,13 +136,13 @@ public abstract class ControlsProviderService extends Service {
         mToken = bundle.getBinder(CALLBACK_TOKEN);
 
         return new IControlsProvider.Stub() {
-            public void load(IControlsLoadCallback cb) {
-                mHandler.obtainMessage(RequestHandler.MSG_LOAD, cb).sendToTarget();
+            public void load(IControlsSubscriber subscriber) {
+                mHandler.obtainMessage(RequestHandler.MSG_LOAD, subscriber).sendToTarget();
             }
 
-            public void loadSuggested(int maxNumber, IControlsLoadCallback cb) {
-                LoadMessage msg = new LoadMessage(maxNumber, cb);
-                mHandler.obtainMessage(RequestHandler.MSG_LOAD_SUGGESTED, msg).sendToTarget();
+            public void loadSuggested(IControlsSubscriber subscriber) {
+                mHandler.obtainMessage(RequestHandler.MSG_LOAD_SUGGESTED, subscriber)
+                        .sendToTarget();
             }
 
             public void subscribe(List<String> controlIds,
@@ -148,73 +171,56 @@ public abstract class ControlsProviderService extends Service {
         private static final int MSG_ACTION = 3;
         private static final int MSG_LOAD_SUGGESTED = 4;
 
-        /**
-         * This the maximum number of controls that can be loaded via
-         * {@link ControlsProviderService#loadAvailablecontrols}. Anything over this number
-         * will be truncated.
-         */
-        private static final int MAX_NUMBER_OF_CONTROLS_ALLOWED = 1000;
-
         RequestHandler(Looper looper) {
             super(looper);
         }
 
         public void handleMessage(Message msg) {
             switch(msg.what) {
-                case MSG_LOAD:
-                    final IControlsLoadCallback cb = (IControlsLoadCallback) msg.obj;
-                    ControlsProviderService.this.loadAvailableControls(consumerFor(
-                            MAX_NUMBER_OF_CONTROLS_ALLOWED, cb));
+                case MSG_LOAD: {
+                    final IControlsSubscriber cs = (IControlsSubscriber) msg.obj;
+                    final SubscriberProxy proxy = new SubscriberProxy(true, mToken, cs);
+
+                    Publisher<Control> publisher =
+                            ControlsProviderService.this.publisherForAllAvailable();
+                    if (publisher == null) {
+                        ControlsProviderService.this.loadAvailableControls(consumerFor(proxy));
+                    } else {
+                        publisher.subscribe(proxy);
+                    }
                     break;
+                }
+
+                case MSG_LOAD_SUGGESTED: {
+                    final IControlsSubscriber cs = (IControlsSubscriber) msg.obj;
+                    final SubscriberProxy proxy = new SubscriberProxy(true, mToken, cs);
 
-                case MSG_LOAD_SUGGESTED:
-                    final LoadMessage lMsg = (LoadMessage) msg.obj;
-                    ControlsProviderService.this.loadSuggestedControls(lMsg.mMaxNumber,
-                            consumerFor(lMsg.mMaxNumber, lMsg.mCb));
+                    Publisher<Control> publisher =
+                            ControlsProviderService.this.publisherForSuggested();
+                    if (publisher == null) {
+                        Log.i(TAG, "No publisher provided for suggested controls");
+                        proxy.onComplete();
+                    } else {
+                        publisher.subscribe(proxy);
+                    }
                     break;
+                }
 
-                case MSG_SUBSCRIBE:
+                case MSG_SUBSCRIBE: {
                     final SubscribeMessage sMsg = (SubscribeMessage) msg.obj;
-                    final IControlsSubscriber cs = sMsg.mSubscriber;
-                    Subscriber<Control> s = new Subscriber<Control>() {
-                            public void onSubscribe(Subscription subscription) {
-                                try {
-                                    cs.onSubscribe(mToken, new SubscriptionAdapter(subscription));
-                                } catch (RemoteException ex) {
-                                    ex.rethrowAsRuntimeException();
-                                }
-                            }
-                            public void onNext(@NonNull Control statefulControl) {
-                                Preconditions.checkNotNull(statefulControl);
-                                try {
-                                    cs.onNext(mToken, statefulControl);
-                                } catch (RemoteException ex) {
-                                    ex.rethrowAsRuntimeException();
-                                }
-                            }
-                            public void onError(Throwable t) {
-                                try {
-                                    cs.onError(mToken, t.toString());
-                                } catch (RemoteException ex) {
-                                    ex.rethrowAsRuntimeException();
-                                }
-                            }
-                            public void onComplete() {
-                                try {
-                                    cs.onComplete(mToken);
-                                } catch (RemoteException ex) {
-                                    ex.rethrowAsRuntimeException();
-                                }
-                            }
-                        };
-                    ControlsProviderService.this.publisherFor(sMsg.mControlIds).subscribe(s);
+                    final SubscriberProxy proxy = new SubscriberProxy(false, mToken,
+                            sMsg.mSubscriber);
+
+                    ControlsProviderService.this.publisherFor(sMsg.mControlIds).subscribe(proxy);
                     break;
+                }
 
-                case MSG_ACTION:
+                case MSG_ACTION: {
                     final ActionMessage aMsg = (ActionMessage) msg.obj;
                     ControlsProviderService.this.performControlAction(aMsg.mControlId,
                             aMsg.mAction, consumerFor(aMsg.mControlId, aMsg.mCb));
                     break;
+                }
             }
         }
 
@@ -234,39 +240,88 @@ public abstract class ControlsProviderService extends Service {
             };
         }
 
-        private Consumer<List<Control>> consumerFor(int maxNumber, IControlsLoadCallback cb) {
-            return (@NonNull List<Control> controls) -> {
+        /**
+         * Method will be removed during migration to publisher
+         */
+        private Consumer<List<Control>> consumerFor(final Subscriber<Control> subscriber) {
+            return (@NonNull final List<Control> controls) -> {
                 Preconditions.checkNotNull(controls);
-                if (controls.size() > maxNumber) {
-                    Log.w(TAG, "Too many controls. Provided: " + controls.size() + ", Max allowed: "
-                            + maxNumber + ". Truncating the list.");
-                    controls = controls.subList(0, maxNumber);
-                }
 
-                List<Control> list = new ArrayList<>();
-                for (Control control: controls) {
-                    if (control == null) {
-                        Log.e(TAG, "onLoad: null control.");
-                    }
-                    if (isStatelessControl(control)) {
-                        list.add(control);
-                    } else {
-                        Log.w(TAG, "onLoad: control is not stateless.");
-                        list.add(new Control.StatelessBuilder(control).build());
-                    }
-                }
-                try {
-                    cb.accept(mToken, list);
-                } catch (RemoteException ex) {
-                    ex.rethrowAsRuntimeException();
-                }
+                subscriber.onSubscribe(new Subscription() {
+                        public void request(long n) {
+                            for (Control control: controls) {
+                                Control c;
+                                if (control == null) {
+                                    Log.e(TAG, "onLoad: null control.");
+                                }
+                                if (isStatelessControl(control)) {
+                                    c = control;
+                                } else {
+                                    Log.w(TAG, "onLoad: control is not stateless.");
+                                    c = new Control.StatelessBuilder(control).build();
+                                }
+
+                                subscriber.onNext(c);
+                            }
+                            subscriber.onComplete();
+                        }
+
+                        public void cancel() {}
+                    });
             };
         }
+    }
+
+    private static boolean isStatelessControl(Control control) {
+        return (control.getStatus() == Control.STATUS_UNKNOWN
+                && control.getControlTemplate().getTemplateType() == ControlTemplate.TYPE_NONE
+                && TextUtils.isEmpty(control.getStatusText()));
+    }
+
+    private static class SubscriberProxy implements Subscriber<Control> {
+        private IBinder mToken;
+        private IControlsSubscriber mCs;
+        private boolean mEnforceStateless;
 
-        private boolean isStatelessControl(Control control) {
-            return (control.getStatus() == Control.STATUS_UNKNOWN
-                    && control.getControlTemplate().getTemplateType() == ControlTemplate.TYPE_NONE
-                    && TextUtils.isEmpty(control.getStatusText()));
+        SubscriberProxy(boolean enforceStateless, IBinder token, IControlsSubscriber cs) {
+            mEnforceStateless = enforceStateless;
+            mToken = token;
+            mCs = cs;
+        }
+
+        public void onSubscribe(Subscription subscription) {
+            try {
+                mCs.onSubscribe(mToken, new SubscriptionAdapter(subscription));
+            } catch (RemoteException ex) {
+                ex.rethrowAsRuntimeException();
+            }
+        }
+        public void onNext(@NonNull Control control) {
+            Preconditions.checkNotNull(control);
+            try {
+                if (mEnforceStateless && !isStatelessControl(control)) {
+                    Log.w(TAG, "onNext(): control is not stateless. Use the "
+                            + "Control.StatelessBuilder() to build the control.");
+                    control = new Control.StatelessBuilder(control).build();
+                }
+                mCs.onNext(mToken, control);
+            } catch (RemoteException ex) {
+                ex.rethrowAsRuntimeException();
+            }
+        }
+        public void onError(Throwable t) {
+            try {
+                mCs.onError(mToken, t.toString());
+            } catch (RemoteException ex) {
+                ex.rethrowAsRuntimeException();
+            }
+        }
+        public void onComplete() {
+            try {
+                mCs.onComplete(mToken);
+            } catch (RemoteException ex) {
+                ex.rethrowAsRuntimeException();
+            }
         }
     }
 
@@ -307,14 +362,4 @@ public abstract class ControlsProviderService extends Service {
             this.mSubscriber = subscriber;
         }
     }
-
-    private static class LoadMessage {
-        final int mMaxNumber;
-        final IControlsLoadCallback mCb;
-
-        LoadMessage(int maxNumber, IControlsLoadCallback cb) {
-            this.mMaxNumber = maxNumber;
-            this.mCb = cb;
-        }
-    }
 }
diff --git a/core/java/android/service/controls/IControlsLoadCallback.aidl b/core/java/android/service/controls/IControlsLoadCallback.aidl
deleted file mode 100644
index bfc61cdb54db..000000000000
--- a/core/java/android/service/controls/IControlsLoadCallback.aidl
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * 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 android.service.controls;
-
-import android.service.controls.Control;
-
-/**
- * @hide
- */
-oneway interface IControlsLoadCallback {
-    void accept(in IBinder token, in List<Control> controls);
-}
\ No newline at end of file
diff --git a/core/java/android/service/controls/IControlsProvider.aidl b/core/java/android/service/controls/IControlsProvider.aidl
index 4375fbb289db..0cb06b4a5cb6 100644
--- a/core/java/android/service/controls/IControlsProvider.aidl
+++ b/core/java/android/service/controls/IControlsProvider.aidl
@@ -17,7 +17,6 @@
 package android.service.controls;
 
 import android.service.controls.IControlsActionCallback;
-import android.service.controls.IControlsLoadCallback;
 import android.service.controls.IControlsSubscriber;
 import android.service.controls.actions.ControlActionWrapper;
 
@@ -25,9 +24,9 @@ import android.service.controls.actions.ControlActionWrapper;
  * @hide
  */
 oneway interface IControlsProvider {
-    void load(IControlsLoadCallback cb);
+    void load(IControlsSubscriber subscriber);
 
-    void loadSuggested(int maxNumber, IControlsLoadCallback cb);
+    void loadSuggested(IControlsSubscriber subscriber);
 
     void subscribe(in List<String> controlIds,
              IControlsSubscriber subscriber);
diff --git a/core/tests/coretests/src/android/service/controls/ControlProviderServiceTest.java b/core/tests/coretests/src/android/service/controls/ControlProviderServiceTest.java
index a24b4e06225a..78c88d71d953 100644
--- a/core/tests/coretests/src/android/service/controls/ControlProviderServiceTest.java
+++ b/core/tests/coretests/src/android/service/controls/ControlProviderServiceTest.java
@@ -21,6 +21,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -62,8 +63,6 @@ public class ControlProviderServiceTest {
     @Mock
     private IControlsActionCallback.Stub mActionCallback;
     @Mock
-    private IControlsLoadCallback.Stub mLoadCallback;
-    @Mock
     private IControlsSubscriber.Stub mSubscriber;
     @Mock
     private IIntentSender mIIntentSender;
@@ -79,8 +78,6 @@ public class ControlProviderServiceTest {
 
         when(mActionCallback.asBinder()).thenCallRealMethod();
         when(mActionCallback.queryLocalInterface(any())).thenReturn(mActionCallback);
-        when(mLoadCallback.asBinder()).thenCallRealMethod();
-        when(mLoadCallback.queryLocalInterface(any())).thenReturn(mLoadCallback);
         when(mSubscriber.asBinder()).thenCallRealMethod();
         when(mSubscriber.queryLocalInterface(any())).thenReturn(mSubscriber);
 
@@ -102,22 +99,28 @@ public class ControlProviderServiceTest {
         Control control2 = new Control.StatelessBuilder("TEST_ID_2", mPendingIntent)
                 .setDeviceType(DeviceTypes.TYPE_AIR_FRESHENER).build();
 
-        @SuppressWarnings("unchecked")
-        ArgumentCaptor<List<Control>> captor = ArgumentCaptor.forClass(List.class);
+        ArgumentCaptor<IControlsSubscription.Stub> subscriptionCaptor =
+                ArgumentCaptor.forClass(IControlsSubscription.Stub.class);
+        ArgumentCaptor<Control> controlCaptor =
+                ArgumentCaptor.forClass(Control.class);
 
         ArrayList<Control> list = new ArrayList<>();
         list.add(control1);
         list.add(control2);
 
         mControlsProviderService.setControls(list);
-        mControlsProvider.load(mLoadCallback);
+        mControlsProvider.load(mSubscriber);
         InstrumentationRegistry.getInstrumentation().waitForIdleSync();
 
-        verify(mLoadCallback).accept(eq(mToken), captor.capture());
-        List<Control> l = captor.getValue();
-        assertEquals(2, l.size());
-        assertTrue(equals(control1, l.get(0)));
-        assertTrue(equals(control2, l.get(1)));
+        verify(mSubscriber).onSubscribe(eq(mToken), subscriptionCaptor.capture());
+        subscriptionCaptor.getValue().request(1000);
+
+        verify(mSubscriber, times(2)).onNext(eq(mToken), controlCaptor.capture());
+        List<Control> values = controlCaptor.getAllValues();
+        assertTrue(equals(values.get(0), list.get(0)));
+        assertTrue(equals(values.get(1), list.get(1)));
+
+        verify(mSubscriber).onComplete(eq(mToken));
     }
 
     @Test
@@ -128,50 +131,57 @@ public class ControlProviderServiceTest {
                 .build();
         Control statelessControl = new Control.StatelessBuilder(control).build();
 
-        @SuppressWarnings("unchecked")
-        ArgumentCaptor<List<Control>> captor = ArgumentCaptor.forClass(List.class);
+        ArgumentCaptor<IControlsSubscription.Stub> subscriptionCaptor =
+                ArgumentCaptor.forClass(IControlsSubscription.Stub.class);
+        ArgumentCaptor<Control> controlCaptor =
+                ArgumentCaptor.forClass(Control.class);
 
         ArrayList<Control> list = new ArrayList<>();
         list.add(control);
 
         mControlsProviderService.setControls(list);
-        mControlsProvider.load(mLoadCallback);
+        mControlsProvider.load(mSubscriber);
         InstrumentationRegistry.getInstrumentation().waitForIdleSync();
 
-        verify(mLoadCallback).accept(eq(mToken), captor.capture());
-        List<Control> l = captor.getValue();
-        assertEquals(1, l.size());
-        assertFalse(equals(control, l.get(0)));
-        assertTrue(equals(statelessControl, l.get(0)));
-        assertEquals(Control.STATUS_UNKNOWN, l.get(0).getStatus());
+        verify(mSubscriber).onSubscribe(eq(mToken), subscriptionCaptor.capture());
+        subscriptionCaptor.getValue().request(1000);
+
+        verify(mSubscriber).onNext(eq(mToken), controlCaptor.capture());
+        Control c = controlCaptor.getValue();
+        assertFalse(equals(control, c));
+        assertTrue(equals(statelessControl, c));
+        assertEquals(Control.STATUS_UNKNOWN, c.getStatus());
+
+        verify(mSubscriber).onComplete(eq(mToken));
     }
 
     @Test
-    public void testLoadSuggested_withMaxNumber() throws RemoteException {
+    public void testOnLoadSuggested_allStateless() throws RemoteException {
         Control control1 = new Control.StatelessBuilder("TEST_ID", mPendingIntent).build();
         Control control2 = new Control.StatelessBuilder("TEST_ID_2", mPendingIntent)
                 .setDeviceType(DeviceTypes.TYPE_AIR_FRESHENER).build();
 
-        @SuppressWarnings("unchecked")
-        ArgumentCaptor<List<Control>> captor = ArgumentCaptor.forClass(List.class);
+        ArgumentCaptor<IControlsSubscription.Stub> subscriptionCaptor =
+                ArgumentCaptor.forClass(IControlsSubscription.Stub.class);
+        ArgumentCaptor<Control> controlCaptor =
+                ArgumentCaptor.forClass(Control.class);
 
         ArrayList<Control> list = new ArrayList<>();
         list.add(control1);
         list.add(control2);
 
-        final int maxSuggested = 1;
-
         mControlsProviderService.setControls(list);
-        mControlsProvider.loadSuggested(maxSuggested, mLoadCallback);
+        mControlsProvider.loadSuggested(mSubscriber);
         InstrumentationRegistry.getInstrumentation().waitForIdleSync();
 
-        verify(mLoadCallback).accept(eq(mToken), captor.capture());
-        List<Control> l = captor.getValue();
-        assertEquals(maxSuggested, l.size());
+        verify(mSubscriber).onSubscribe(eq(mToken), subscriptionCaptor.capture());
+        subscriptionCaptor.getValue().request(1);
 
-        for (int i = 0; i < maxSuggested; ++i) {
-            assertTrue(equals(list.get(i), l.get(i)));
-        }
+        verify(mSubscriber).onNext(eq(mToken), controlCaptor.capture());
+        Control c = controlCaptor.getValue();
+        assertTrue(equals(c, list.get(0)));
+
+        verify(mSubscriber).onComplete(eq(mToken));
     }
 
     @Test
@@ -244,22 +254,19 @@ public class ControlProviderServiceTest {
         }
 
         @Override
-        public void loadSuggestedControls(int maxNumber, Consumer<List<Control>> cb) {
-            cb.accept(mControls);
+        public Publisher<Control> publisherFor(List<String> ids) {
+            return new Publisher<Control>() {
+                public void subscribe(final Subscriber s) {
+                    s.onSubscribe(createSubscription(s, mControls));
+                }
+            };
         }
 
         @Override
-        public Publisher<Control> publisherFor(List<String> ids) {
+        public Publisher<Control> publisherForSuggested() {
             return new Publisher<Control>() {
                 public void subscribe(final Subscriber s) {
-                    s.onSubscribe(new Subscription() {
-                            public void request(long n) {
-                                for (Control c : mControls) {
-                                    s.onNext(c);
-                                }
-                            }
-                            public void cancel() {}
-                        });
+                    s.onSubscribe(createSubscription(s, mControls));
                 }
             };
         }
@@ -269,7 +276,19 @@ public class ControlProviderServiceTest {
                 Consumer<Integer> cb) {
             cb.accept(ControlAction.RESPONSE_OK);
         }
+
+        private Subscription createSubscription(Subscriber s, List<Control> controls) {
+            return new Subscription() {
+                public void request(long n) {
+                    int i = 0;
+                    for (Control c : mControls) {
+                        if (i++ < n) s.onNext(c);
+                        else break;
+                    }
+                    s.onComplete();
+                }
+                public void cancel() {}
+            };
+        }
     }
 }
-
-
diff --git a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt
index 87bdfa8e04ec..6ff1bbce672f 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt
@@ -22,7 +22,6 @@ import android.os.IBinder
 import android.os.UserHandle
 import android.service.controls.Control
 import android.service.controls.IControlsActionCallback
-import android.service.controls.IControlsLoadCallback
 import android.service.controls.IControlsSubscriber
 import android.service.controls.IControlsSubscription
 import android.service.controls.actions.ControlAction
@@ -63,12 +62,6 @@ open class ControlsBindingControllerImpl @Inject constructor(
     private val componentMap: MutableMap<Key, ControlsProviderLifecycleManager> =
             ArrayMap<Key, ControlsProviderLifecycleManager>()
 
-    private val loadCallbackService = object : IControlsLoadCallback.Stub() {
-        override fun accept(token: IBinder, controls: MutableList<Control>) {
-            backgroundExecutor.execute(OnLoadRunnable(token, controls))
-        }
-    }
-
     private val actionCallbackService = object : IControlsActionCallback.Stub() {
         override fun accept(
             token: IBinder,
@@ -106,7 +99,6 @@ open class ControlsBindingControllerImpl @Inject constructor(
         return ControlsProviderLifecycleManager(
                 context,
                 backgroundExecutor,
-                loadCallbackService,
                 actionCallbackService,
                 subscriberService,
                 currentUser,
@@ -130,7 +122,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
         callback: ControlsBindingController.LoadCallback
     ) {
         val provider = retrieveLifecycleManager(component)
-        provider.maybeBindAndLoad(callback)
+        provider.maybeBindAndLoad(LoadSubscriber(callback))
     }
 
     override fun subscribe(controls: List<ControlInfo>) {
@@ -216,7 +208,8 @@ open class ControlsBindingControllerImpl @Inject constructor(
 
     private inner class OnLoadRunnable(
         token: IBinder,
-        val list: List<Control>
+        val list: List<Control>,
+        val callback: ControlsBindingController.LoadCallback
     ) : CallbackRunnable(token) {
         override fun run() {
             if (provider == null) {
@@ -233,9 +226,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
                     return
                 }
             }
-            provider.lastLoadCallback?.accept(list) ?: run {
-                Log.w(TAG, "Null callback")
-            }
+            callback.accept(list)
             provider.unbindService()
         }
     }
@@ -277,7 +268,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
     ) : CallbackRunnable(token) {
         override fun run() {
             provider?.let {
-                Log.i(TAG, "onComplete receive from '${provider.componentName}'")
+                Log.i(TAG, "onComplete receive from '${it.componentName}'")
             }
         }
     }
@@ -288,7 +279,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
     ) : CallbackRunnable(token) {
         override fun run() {
             provider?.let {
-                Log.e(TAG, "onError receive from '${provider.componentName}': $error")
+                Log.e(TAG, "onError receive from '${it.componentName}': $error")
             }
         }
     }
@@ -308,6 +299,44 @@ open class ControlsBindingControllerImpl @Inject constructor(
             }
         }
     }
+
+    private inner class OnLoadErrorRunnable(
+        token: IBinder,
+        val error: String,
+        val callback: ControlsBindingController.LoadCallback
+    ) : CallbackRunnable(token) {
+        override fun run() {
+            callback.error(error)
+            provider?.let {
+                Log.e(TAG, "onError receive from '${it.componentName}': $error")
+            }
+        }
+    }
+
+    private inner class LoadSubscriber(
+        val callback: ControlsBindingController.LoadCallback
+    ) : IControlsSubscriber.Stub() {
+        val loadedControls = ArrayList<Control>()
+        var hasError = false
+
+        override fun onSubscribe(token: IBinder, subs: IControlsSubscription) {
+            backgroundExecutor.execute(OnSubscribeRunnable(token, subs))
+        }
+
+        override fun onNext(token: IBinder, c: Control) {
+            backgroundExecutor.execute { loadedControls.add(c) }
+        }
+        override fun onError(token: IBinder, s: String) {
+            hasError = true
+            backgroundExecutor.execute(OnLoadErrorRunnable(token, s, callback))
+        }
+
+        override fun onComplete(token: IBinder) {
+            if (!hasError) {
+                backgroundExecutor.execute(OnLoadRunnable(token, loadedControls, callback))
+            }
+        }
+    }
 }
 
-private data class Key(val component: ComponentName, val user: UserHandle)
\ No newline at end of file
+private data class Key(val component: ComponentName, val user: UserHandle)
diff --git a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManager.kt b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManager.kt
index 9ec71c754bb5..a53fcd498236 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManager.kt
@@ -29,7 +29,6 @@ import android.service.controls.ControlsProviderService
 import android.service.controls.ControlsProviderService.CALLBACK_BUNDLE
 import android.service.controls.ControlsProviderService.CALLBACK_TOKEN
 import android.service.controls.IControlsActionCallback
-import android.service.controls.IControlsLoadCallback
 import android.service.controls.IControlsProvider
 import android.service.controls.IControlsSubscriber
 import android.service.controls.IControlsSubscription
@@ -48,8 +47,6 @@ import java.util.concurrent.TimeUnit
  *
  * @property context A SystemUI context for binding to the services
  * @property executor A delayable executor for posting timeouts
- * @property loadCallbackService a callback interface to hand the remote service for loading
- *                               controls
  * @property actionCallbackService a callback interface to hand the remote service for sending
  *                                 action responses
  * @property subscriberService an "subscriber" interface for requesting and accepting updates for
@@ -60,15 +57,12 @@ import java.util.concurrent.TimeUnit
 class ControlsProviderLifecycleManager(
     private val context: Context,
     private val executor: DelayableExecutor,
-    private val loadCallbackService: IControlsLoadCallback.Stub,
     private val actionCallbackService: IControlsActionCallback.Stub,
     private val subscriberService: IControlsSubscriber.Stub,
     val user: UserHandle,
     val componentName: ComponentName
 ) : IBinder.DeathRecipient {
 
-    var lastLoadCallback: ControlsBindingController.LoadCallback? = null
-        private set
     val token: IBinder = Binder()
     @GuardedBy("subscriptions")
     private val subscriptions = mutableListOf<IControlsSubscription>()
@@ -156,9 +150,12 @@ class ControlsProviderLifecycleManager(
             bindService(false)
             return
         }
-        if (Message.Load in queue) {
-            load()
+
+        queue.filter { it is Message.Load }.forEach {
+            val msg = it as Message.Load
+            load(msg.subscriber)
         }
+
         queue.filter { it is Message.Subscribe }.flatMap { (it as Message.Subscribe).list }.run {
             if (this.isNotEmpty()) {
                 subscribe(this)
@@ -193,12 +190,12 @@ class ControlsProviderLifecycleManager(
         }
     }
 
-    private fun load() {
+    private fun load(subscriber: IControlsSubscriber.Stub) {
         if (DEBUG) {
             Log.d(TAG, "load $componentName")
         }
-        if (!(wrapper?.load(loadCallbackService) ?: false)) {
-            queueMessage(Message.Load)
+        if (!(wrapper?.load(subscriber) ?: false)) {
+            queueMessage(Message.Load(subscriber))
             binderDied()
         }
     }
@@ -213,27 +210,23 @@ class ControlsProviderLifecycleManager(
     }
 
     /**
-     * Request a call to [ControlsProviderService.loadAvailableControls].
+     * Request a call to [IControlsProvider.load].
      *
      * If the service is not bound, the call will be queued and the service will be bound first.
-     * The service will be bound after the controls are returned or the call times out.
+     * The service will be unbound after the controls are returned or the call times out.
      *
-     * @param callback a callback in which to return the result back. If the call times out
-     *                 [ControlsBindingController.LoadCallback.error] will be called instead.
+     * @param subscriber the subscriber that manages coordination for loading controls
      */
-    fun maybeBindAndLoad(callback: ControlsBindingController.LoadCallback) {
+    fun maybeBindAndLoad(subscriber: IControlsSubscriber.Stub) {
         unqueueMessage(Message.Unbind)
-        lastLoadCallback = callback
         onLoadCanceller = executor.executeDelayed({
             // Didn't receive a response in time, log and send back error
             Log.d(TAG, "Timeout waiting onLoad for $componentName")
-            callback.error("Timeout waiting onLoad")
-            // Don't accept load callbacks after this
-            lastLoadCallback = null
+            subscriber.onError(token, "Timeout waiting onLoad")
             unbindService()
         }, LOAD_TIMEOUT_SECONDS, TimeUnit.SECONDS)
 
-        invokeOrQueue(::load, Message.Load)
+        invokeOrQueue({ load(subscriber) }, Message.Load(subscriber))
     }
 
     /**
@@ -324,7 +317,6 @@ class ControlsProviderLifecycleManager(
      * Request unbind from the service.
      */
     fun unbindService() {
-        lastLoadCallback = null
         onLoadCanceller?.run()
         onLoadCanceller = null
 
@@ -344,7 +336,7 @@ class ControlsProviderLifecycleManager(
      */
     sealed class Message {
         abstract val type: Int
-        object Load : Message() {
+        class Load(val subscriber: IControlsSubscriber.Stub) : Message() {
             override val type = MSG_LOAD
         }
         object Unbind : Message() {
diff --git a/packages/SystemUI/src/com/android/systemui/controls/controller/ServiceWrapper.kt b/packages/SystemUI/src/com/android/systemui/controls/controller/ServiceWrapper.kt
index b90f892d5d26..b2afd3c144d1 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/controller/ServiceWrapper.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/controller/ServiceWrapper.kt
@@ -18,7 +18,6 @@ package com.android.systemui.controls.controller
 
 import android.service.controls.actions.ControlAction
 import android.service.controls.IControlsActionCallback
-import android.service.controls.IControlsLoadCallback
 import android.service.controls.IControlsProvider
 import android.service.controls.IControlsSubscriber
 import android.service.controls.IControlsSubscription
@@ -45,9 +44,9 @@ class ServiceWrapper(val service: IControlsProvider) {
         }
     }
 
-    fun load(cb: IControlsLoadCallback): Boolean {
+    fun load(subscriber: IControlsSubscriber): Boolean {
         return callThroughService {
-            service.load(cb)
+            service.load(subscriber)
         }
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt
index 40075c8413d2..02bfc19e07d4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt
@@ -45,7 +45,7 @@ import org.mockito.MockitoAnnotations
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
-class ControlsBindingControllerTest : SysuiTestCase() {
+class ControlsBindingControllerImplTest : SysuiTestCase() {
 
     companion object {
         fun <T> any(): T = Mockito.any<T>()
@@ -95,7 +95,7 @@ class ControlsBindingControllerTest : SysuiTestCase() {
 
         assertEquals(1, providers.size)
         val provider = providers.first()
-        verify(provider).maybeBindAndLoad(callback)
+        verify(provider).maybeBindAndLoad(any())
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManagerTest.kt
index ddd6b125dcfa..a3e59e52abe5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsProviderLifecycleManagerTest.kt
@@ -19,7 +19,6 @@ package com.android.systemui.controls.controller
 import android.content.ComponentName
 import android.os.UserHandle
 import android.service.controls.IControlsActionCallback
-import android.service.controls.IControlsLoadCallback
 import android.service.controls.IControlsProvider
 import android.service.controls.IControlsSubscriber
 import android.service.controls.actions.ControlAction
@@ -32,13 +31,13 @@ import com.android.systemui.util.time.FakeSystemClock
 import org.junit.After
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
-import org.junit.Assert.assertNull
 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
+import org.mockito.ArgumentMatchers.any
 import org.mockito.ArgumentMatchers.anyString
 import org.mockito.ArgumentMatchers.eq
 import org.mockito.Captor
@@ -54,8 +53,6 @@ class ControlsProviderLifecycleManagerTest : SysuiTestCase() {
     @Mock
     private lateinit var actionCallbackService: IControlsActionCallback.Stub
     @Mock
-    private lateinit var loadCallbackService: IControlsLoadCallback.Stub
-    @Mock
     private lateinit var subscriberService: IControlsSubscriber.Stub
     @Mock
     private lateinit var service: IControlsProvider.Stub
@@ -85,7 +82,6 @@ class ControlsProviderLifecycleManagerTest : SysuiTestCase() {
         manager = ControlsProviderLifecycleManager(
                 context,
                 executor,
-                loadCallbackService,
                 actionCallbackService,
                 subscriberService,
                 UserHandle.of(0),
@@ -113,31 +109,29 @@ class ControlsProviderLifecycleManagerTest : SysuiTestCase() {
 
     @Test
     fun testMaybeBindAndLoad() {
-        manager.maybeBindAndLoad(loadCallback)
+        manager.maybeBindAndLoad(subscriberService)
 
-        verify(service).load(loadCallbackService)
+        verify(service).load(subscriberService)
 
         assertTrue(mContext.isBound(componentName))
-        assertEquals(loadCallback, manager.lastLoadCallback)
     }
 
     @Test
     fun testMaybeUnbind_bindingAndCallback() {
-        manager.maybeBindAndLoad(loadCallback)
+        manager.maybeBindAndLoad(subscriberService)
 
         manager.unbindService()
         assertFalse(mContext.isBound(componentName))
-        assertNull(manager.lastLoadCallback)
     }
 
     @Test
     fun testMaybeBindAndLoad_timeout() {
-        manager.maybeBindAndLoad(loadCallback)
+        manager.maybeBindAndLoad(subscriberService)
 
         executor.advanceClockToLast()
         executor.runAllReady()
 
-        verify(loadCallback).error(anyString())
+        verify(subscriberService).onError(any(), anyString())
     }
 
     @Test
@@ -160,4 +154,4 @@ class ControlsProviderLifecycleManagerTest : SysuiTestCase() {
                 eq(actionCallbackService))
         assertEquals(action, wrapperCaptor.getValue().getWrappedAction())
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ServiceWrapperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ServiceWrapperTest.kt
index 9e7ce06bb74f..cd82844fcc50 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ServiceWrapperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ServiceWrapperTest.kt
@@ -18,7 +18,6 @@ package com.android.systemui.controls.controller
 
 import android.os.RemoteException
 import android.service.controls.IControlsActionCallback
-import android.service.controls.IControlsLoadCallback
 import android.service.controls.IControlsProvider
 import android.service.controls.IControlsSubscriber
 import android.service.controls.IControlsSubscription
@@ -56,9 +55,6 @@ class ServiceWrapperTest : SysuiTestCase() {
     private lateinit var subscriber: IControlsSubscriber
 
     @Mock
-    private lateinit var loadCallback: IControlsLoadCallback
-
-    @Mock
     private lateinit var actionCallback: IControlsActionCallback
 
     @Captor
@@ -81,16 +77,16 @@ class ServiceWrapperTest : SysuiTestCase() {
 
     @Test
     fun testLoad_happyPath() {
-        val result = wrapper.load(loadCallback)
+        val result = wrapper.load(subscriber)
 
         assertTrue(result)
-        verify(service).load(loadCallback)
+        verify(service).load(subscriber)
     }
 
     @Test
     fun testLoad_error() {
         `when`(service.load(any())).thenThrow(exception)
-        val result = wrapper.load(loadCallback)
+        val result = wrapper.load(subscriber)
 
         assertFalse(result)
     }